Update app.py

app.py

@@ -1,89 +1,137 @@
-import gradio as gr
 import os
+import glob
+from pathlib import Path
+from datetime import datetime, timezone
+
+import streamlit as st
 import pandas as pd
 from PIL import Image
-import random
-
-IMAGE_FOLDER = "images"
-LABELS_FILE = "labels.csv"
-
-# Load or create labels dataframe
-if os.path.exists(LABELS_FILE):
-    labels_df = pd.read_csv(LABELS_FILE)
+from filelock import FileLock
+
+# ---------------- Config ----------------
+IMAGE_DIR = os.getenv("IMAGE_DIR", "images")
+LABELS_CSV = os.getenv("LABELS_CSV", "labels.csv")
+LABEL_LEFT = os.getenv("LABEL_A", "Class A")
+LABEL_RIGHT = os.getenv("LABEL_B", "Class B")
+SHOW_ALREADY_LABELED = os.getenv("SHOW_ALREADY_LABELED", "0") == "1"
+SUPPORTED_EXTS = (".png", ".jpg", ".jpeg", ".bmp", ".gif", ".tif", ".tiff", ".webp")
+
+st.set_page_config(page_title="Two-Button Image Labeler", layout="centered")
+
+# --------------- Helpers ----------------
+def list_images():
+    paths = []
+    for p in glob.glob(os.path.join(IMAGE_DIR, "**", "*"), recursive=True):
+        if p.lower().endswith(SUPPORTED_EXTS):
+            paths.append(p)
+    return sorted(paths)
+
+def read_labels():
+    if os.path.exists(LABELS_CSV):
+        try:
+            return pd.read_csv(LABELS_CSV)
+        except Exception:
+            return pd.DataFrame(columns=["image", "label", "annotator", "timestamp"])
+    return pd.DataFrame(columns=["image", "label", "annotator", "timestamp"])
+
+def rel_to_image_dir(p: str):
+    try:
+        return str(Path(p).resolve().relative_to(Path(IMAGE_DIR).resolve()))
+    except Exception:
+        return p
+
+def write_label(image_path: str, label: str, annotator: str):
+    os.makedirs(os.path.dirname(LABELS_CSV) or ".", exist_ok=True)
+    record = {
+        "image": rel_to_image_dir(image_path),
+        "label": label,
+        "annotator": annotator,
+        "timestamp": datetime.now(timezone.utc).isoformat()
+    }
+    with FileLock(LABELS_CSV + ".lock", timeout=10):
+        exists = os.path.exists(LABELS_CSV)
+        df = pd.DataFrame([record])
+        df.to_csv(LABELS_CSV, mode="a", header=not exists, index=False)
+
+# --------------- UI ---------------------
+st.title("🏷️ Two-Button Image Labeler")
+st.write("Enter your name, click **Start**, then label each image using the big buttons. Use **Skip** to bypass an image.")
+
+with st.sidebar:
+    annotator = st.text_input("Your name*", value="")
+    start = st.button("Start", type="primary")
+
+# session state
+if "order" not in st.session_state:
+    st.session_state.order = []
+if "idx" not in st.session_state:
+    st.session_state.idx = 0
+if "total" not in st.session_state:
+    st.session_state.total = 0
+if "started" not in st.session_state:
+    st.session_state.started = False
+
+# start
+if start:
+    if not annotator.strip():
+        st.sidebar.error("Please enter your name.")
+    else:
+        imgs = list_images()
+        labels_df = read_labels()
+        if not SHOW_ALREADY_LABELED and len(labels_df) > 0:
+            labeled = set(labels_df["image"].astype(str).tolist())
+            rel_imgs = [rel_to_image_dir(p) for p in imgs]
+            imgs = [p for p, r in zip(imgs, rel_imgs) if r not in labeled]
+
+        if not imgs:
+            st.warning("No images found (or all labeled). Upload PNG/JPG/etc. to the `images/` folder.")
+        else:
+            # Simple deterministic order; you can randomize with: random.shuffle(imgs)
+            st.session_state.order = imgs
+            st.session_state.idx = 0
+            st.session_state.total = len(imgs)
+            st.session_state.started = True
+
+# main panel
+if not st.session_state.started:
+    st.info("Fill your name on the left and press **Start**.")
 else:
-    labels_df = pd.DataFrame(columns=["user", "image", "label"])
-
-def get_images():
-    valid_exts = {".jpg", ".jpeg", ".png", ".bmp", ".gif", ".tif", ".tiff", ".webp"}
-    all_files = []
-    for root, _, files in os.walk(IMAGE_FOLDER):
-        for f in files:
-            if os.path.splitext(f.lower())[1] in valid_exts:
-                all_files.append(os.path.join(root, f))
-    return sorted(all_files)
-
-def start_labeling(user):
-    images = get_images()
-    if not images:
-        return None, f"No images found in '{IMAGE_FOLDER}'", gr.update(visible=False)
-    random.shuffle(images)
-    # Store the shuffled list in session state
-    return images[0], f"Found {len(images)} images. Starting labeling...", gr.update(visible=True), images
-
-def label_image(user, current_image, label, remaining_images):
-    global labels_df
-    labels_df = pd.concat([labels_df, pd.DataFrame([[user, current_image, label]], columns=labels_df.columns)], ignore_index=True)
-    labels_df.to_csv(LABELS_FILE, index=False)
-
-    if not remaining_images:
-        return None, "All images labeled!", gr.update(visible=False), []
-
-    next_image = remaining_images.pop(0)
-    return next_image, f"{len(remaining_images)} images remaining.", gr.update(visible=True), remaining_images
-
-with gr.Blocks() as demo:
-    gr.Markdown("# 🏷️ Two-Button Image Labeler")
-
-    with gr.Row():
-        user_input = gr.Textbox(label="Your name", placeholder="Enter your name")
-        start_btn = gr.Button("Start Labeling", variant="primary")
-
-    info = gr.Markdown()
-    image_display = gr.Image(type="filepath", label="Image to label")
-    with gr.Row(visible=False) as button_row:
-        left_btn = gr.Button("⬅️ Class A")
-        skip_btn = gr.Button("Skip")
-        right_btn = gr.Button("Class B ➡️")
-
-    state_images = gr.State([])
-
-    start_btn.click(
-        fn=lambda user: start_labeling(user),
-        inputs=user_input,
-        outputs=[image_display, info, button_row, state_images]
-    )
-
-    left_btn.click(
-        fn=lambda user, img, rem: label_image(user, img, "Class A", rem),
-        inputs=[user_input, image_display, state_images],
-        outputs=[image_display, info, button_row, state_images]
-    )
-
-    right_btn.click(
-        fn=lambda user, img, rem: label_image(user, img, "Class B", rem),
-        inputs=[user_input, image_display, state_images],
-        outputs=[image_display, info, button_row, state_images]
-    )
-
-    skip_btn.click(
-        fn=lambda user, img, rem: label_image(user, img, "Skip", rem),
-        inputs=[user_input, image_display, state_images],
-        outputs=[image_display, info, button_row, state_images]
-    )
-
-
-demo = demo.queue(max_size=64)
-
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, show_api=False)
-
+    idx = st.session_state.idx
+    total = st.session_state.total
+    if idx >= total:
+        st.success("All done 🎉 Thank you!")
+    else:
+        current_image = st.session_state.order[idx]
+        st.caption(f"{idx} / {total}")
+        try:
+            img = Image.open(current_image)
+            # If multipage TIFF, show first page
+            if getattr(img, "n_frames", 1) > 1:
+                img.seek(0)
+            # Convert to RGB for browsers
+            if img.mode not in ("RGB", "RGBA"):
+                img = img.convert("RGB")
+            st.image(img, use_column_width=True)
+        except Exception as e:
+            st.warning(f"Could not display image: {current_image}\n{e}")
+
+        col1, col2, col3 = st.columns([1,1,1])
+        with col1:
+            if st.button(f"⬅️ {LABEL_LEFT}", use_container_width=True):
+                write_label(current_image, LABEL_LEFT, annotator.strip())
+                st.session_state.idx += 1
+                st.rerun()
+        with col2:
+            if st.button("Skip", use_container_width=True):
+                write_label(current_image, "Skip", annotator.strip())
+                st.session_state.idx += 1
+                st.rerun()
+        with col3:
+            if st.button(f"{LABEL_RIGHT} ➡️", use_container_width=True):
+                write_label(current_image, LABEL_RIGHT, annotator.strip())
+                st.session_state.idx += 1
+                st.rerun()
+
+# Footer / admin
+st.divider()
+st.caption("Labels are saved to `labels.csv` in this Space. You can download it from the Files tab.")