annotation.py

import streamlit as st
import os
from PIL import Image
import csv

# ---------------- CONFIG ----------------
IMAGE_FOLDER = '/path/to/patches'
# ----------------------------------------

st.set_page_config(page_title="Cell Segmentation Agreement Survey", layout="wide")

# Title always shown
st.markdown("<h1>🔬 Cell Segmentation Agreement Survey</h1>", unsafe_allow_html=True)

# --- Name input with session state locking ---
if 'annotator' not in st.session_state:
    st.session_state.annotator = ""

if not st.session_state.annotator:
    st.markdown(
        "<span style='font-size:22px; font-weight:500;'>Enter your name:</span>",
        unsafe_allow_html=True
    )
    annotator_input = st.text_input(
        "Annotator Name",  # Non-empty label for accessibility
        key="annotator_name_box",
        label_visibility="collapsed"  # Hide label visually, show for screen readers
    )
    if annotator_input:
        st.session_state.annotator = annotator_input.strip()
        st.rerun()
    else:
        st.markdown(
            "<div style='color:#FF4B4B; font-size:20px; margin-bottom: 18px;'>"
            "Please enter your name to begin the survey. (Press Enter on your keyboard after typing your name.)"
            "</div>",
            unsafe_allow_html=True
        )
        st.stop()

# Now annotator name is stored in st.session_state.annotator
annotator = st.session_state.annotator

# Display annotator name in top-right corner (only, not as editable box)
st.markdown(
    f"<div style='position: absolute; top: 10px; right: 20px; color: gray; font-size: 18px;'>Annotator: <strong>{annotator}</strong></div>",
    unsafe_allow_html=True
)

# --- Get the output folder (parent of IMAGE_FOLDER) ---
OUTPUT_FOLDER = os.path.dirname(IMAGE_FOLDER)
csv_file = os.path.join(OUTPUT_FOLDER, f"annotations_{annotator.lower().replace(' ', '_')}.csv")

if not os.path.exists(csv_file):
    with open(csv_file, mode='w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(['annotator', 'patch_id', 'agreement'])

# Load images
all_images = sorted([img for img in os.listdir(IMAGE_FOLDER) if img.lower().endswith('.png')])
if not all_images:
    st.error("No .png images found in the folder.")
    st.stop()

# Check for already annotated images
annotated = set()
with open(csv_file, newline='') as f:
    reader = csv.DictReader(f)
    for row in reader:
        annotated.add(row['patch_id'])

unannotated = [img for img in all_images if img not in annotated]
num_annotated = len(annotated)
num_total = len(all_images)
current_index = num_annotated + 1  # 1-based index for display

if not unannotated:
    st.success("✅ All images have been annotated!")
    st.stop()

# Progress
st.markdown(
    f"<b><span style='font-size:22px;'>Progress: {current_index}/{num_total}</span></b>",
    unsafe_allow_html=True
)

# --- Display image (smaller and centered) ---
current_image = unannotated[0]
image_path = os.path.join(IMAGE_FOLDER, current_image)
image = Image.open(image_path)

# Center the image using markdown, set width to 800
st.markdown("<div style='text-align:center;'>", unsafe_allow_html=True)
st.image(image, caption=current_image, width=1200)
st.markdown("</div>", unsafe_allow_html=True)

# Display question
st.markdown("### How much do you agree with the cell segmentation result (left: segmentation result, right: original tissue)?")

# Create 5-level buttons
cols = st.columns(5)
choices = [
    ("1", "Strong Disagree"),
    ("2", "Disagree"),
    ("3", "Neutral"),
    ("4", "Agree"),
    ("5", "Strong Agree")
]

response = None
for i, (val, label) in enumerate(choices):
    if cols[i].button(f"{val} {label}"):
        response = val
        break

# Save annotation
if response:
    with open(csv_file, mode='a', newline='') as f:
        writer = csv.writer(f)
        writer.writerow([annotator, current_image, response])
    st.success(f"Saved: {current_image} → {response} ({choices[int(response)-1][1]})")
    st.rerun()