Create app.py

app.py

@@ -0,0 +1,971 @@
+import os
+import json
+import time
+import random
+import hashlib
+import threading
+from dataclasses import dataclass
+from typing import List, Dict, Any
+
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+import gradio as gr
+from filelock import FileLock
+# ----------------------
+# Configuration
+# ----------------------
+
+SR_DIR = os.environ.get("SR_DIR", "./data/sr")
+GT_DIR = os.environ.get("GT_DIR", "./data/gt")
+RESULTS_DIR = os.environ.get("RESULTS_DIR", "./results")
+PROGRESS_PATH = os.path.join(RESULTS_DIR, "progress.json")
+ALL_RESULTS_JSONL = os.path.join(RESULTS_DIR, "all_results.jsonl")
+SAVE_PII = True  # set False to only store hashed user_id
+
+
+# NEW — global coordination files
+ASSIGN_PATH = os.path.join(RESULTS_DIR, "assignments.json")  # in-flight reservations
+ROUND_PATH  = os.path.join(RESULTS_DIR, "round.json")        # current round state
+
+# NEW — cross-process locks for those files
+ASSIGN_LOCK = FileLock(ASSIGN_PATH + ".lock")
+ROUND_LOCK  = FileLock(ROUND_PATH + ".lock")
+
+# NEW — reservation expires if user disappears before submitting (seconds)
+RESERVATION_TTL = 600
+
+
+
+# Tile grid config: 20 tiles -> 4 rows x 5 columns
+TILE_ROWS = 4
+TILE_COLS = 5
+N_TILES = TILE_ROWS * TILE_COLS
+
+# Which image is used for the RTS detection / tile selection part:
+#   "sr" -> use SR image; "gt" -> use GT image
+TEST_IMAGE_SOURCE = "sr"
+
+# Thread lock for writing files safely in multi-user scenarios
+WRITE_LOCK = threading.Lock()
+
+# ----------------------
+# Data model
+# ----------------------
+@dataclass
+class Sample:
+    sample_id: str
+    sr_path: str
+    gt_path: str
+
+
+# === GLOBAL ROUND / ASSIGNMENT HELPERS ===
+
+def _load_json(path: str) -> Dict[str, Any]:
+    if not os.path.exists(path):
+        return {}
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            return json.load(f)
+    except Exception:
+        return {}
+
+def _save_json(path: str, data: Dict[str, Any]):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(data, f, ensure_ascii=False, indent=2)
+
+def get_round_and_left(total: int) -> tuple[int, int]:
+    """Return (current_round_number, images_left_in_this_round)."""
+    with ROUND_LOCK:
+        st = _load_json(ROUND_PATH) or {}
+        r = int(st.get("round", 1))
+        done = set(st.get("done", []))
+        left = max(0, total - len(done))
+        return r, left
+
+def mark_done_and_maybe_rollover(sample_id: str, total: int) -> tuple[int, int, bool]:
+    """
+    Mark sample as done for the current global round.
+    Returns (round_number_after_op, images_left_after_op, rolled_over_bool).
+    If we finished the round, it automatically starts the next round.
+    """
+    with ROUND_LOCK:
+        st = _load_json(ROUND_PATH) or {}
+        st.setdefault("round", 1)
+        done = set(st.get("done", []))
+        rolled = False
+
+        if sample_id not in done:
+            done.add(sample_id)
+
+        if len(done) >= total:
+            # complete round, roll to next
+            st["round"] = int(st["round"]) + 1
+            st["done"] = []
+            _save_json(ROUND_PATH, st)
+            return int(st["round"]), total, True
+        else:
+            st["done"] = sorted(done)
+            _save_json(ROUND_PATH, st)
+            left = total - len(done)
+            return int(st["round"]), left, False
+
+def _cleanup_and_load_assignments() -> Dict[str, Any]:
+    """Drop expired reservations and return the map."""
+    now = time.time()
+    with ASSIGN_LOCK:
+        a = _load_json(ASSIGN_PATH) or {}
+        fresh = {k: v for k, v in a.items() if now - v.get("ts", 0) < RESERVATION_TTL}
+        if len(fresh) != len(a):
+            _save_json(ASSIGN_PATH, fresh)
+        return fresh
+
+def release_reservation(sample_id: str):
+    with ASSIGN_LOCK:
+        a = _load_json(ASSIGN_PATH) or {}
+        if a.pop(sample_id, None) is not None:
+            _save_json(ASSIGN_PATH, a)
+
+def reserve_next_sample(uid: str, samples: List[Sample]) -> int:
+    """
+    Reserve and return an index into `samples` that:
+      - is NOT already done in the current round
+      - is NOT reserved by someone else (unless reservation expired)
+    If the round is already complete, it rolls over and picks from the new round.
+    """
+    total = len(samples)
+    with ASSIGN_LOCK, ROUND_LOCK:
+        # Clean old reservations first
+        assignments = _cleanup_and_load_assignments()
+
+        # Load round state
+        st = _load_json(ROUND_PATH) or {}
+        st.setdefault("round", 1)
+        st.setdefault("done", [])
+        done = set(st["done"])
+
+        # If current round is exhausted, roll now so users can continue seamlessly
+        if len(done) >= total:
+            st["round"] = int(st["round"]) + 1
+            st["done"] = []
+            done = set()
+            _save_json(ROUND_PATH, st)
+
+        assigned = set(assignments.keys())
+
+        # Prefer: not-done + not-assigned
+        candidates = [i for i, s in enumerate(samples)
+                      if s.sample_id not in done and s.sample_id not in assigned]
+
+        # If everything not-done is currently reserved, fall back to any not-assigned (rare)
+        if not candidates:
+            candidates = [i for i, s in enumerate(samples) if s.sample_id not in assigned]
+            if not candidates:
+                candidates = list(range(total))  # last resort
+
+        idx = random.choice(candidates)
+        assignments[samples[idx].sample_id] = {"uid": uid, "ts": time.time()}
+        _save_json(ASSIGN_PATH, assignments)
+        return idx
+# === END GLOBAL ROUND / ASSIGNMENT HELPERS ===
+
+
+def ensure_paths():
+    os.makedirs(RESULTS_DIR, exist_ok=True)
+    if not os.path.isdir(SR_DIR) or not os.path.isdir(GT_DIR):
+        raise FileNotFoundError(
+            f"Expected directories '{SR_DIR}' and '{GT_DIR}'.\n"
+            "Create them and add matching SR/GT image pairs."
+        )
+
+
+def list_image_files(folder: str) -> List[str]:
+    supported = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"}
+    return sorted(
+        [os.path.join(folder, f) for f in os.listdir(folder)
+         if os.path.splitext(f.lower())[1] in supported]
+    )
+
+
+def load_dataset(sr_dir: str, gt_dir: str) -> List[Sample]:
+    sr_files = list_image_files(sr_dir)
+    gt_files = list_image_files(gt_dir)
+
+    def key(path: str) -> str:
+        return os.path.splitext(os.path.basename(path))[0]
+
+    gt_map = {key(p): p for p in gt_files}
+    samples: List[Sample] = []
+
+    for sr_path in sr_files:
+        k = key(sr_path)
+        if k in gt_map:
+            samples.append(Sample(sample_id=k, sr_path=sr_path, gt_path=gt_map[k]))
+
+    if not samples:
+        raise RuntimeError(
+            f"No matched SR/GT pairs found. Ensure filenames match in {sr_dir} and {gt_dir}."
+        )
+    return samples
+
+
+# ----------------------
+# Progress & results I/O
+# ----------------------
+
+def hash_user_id(name: str, address: str) -> str:
+    norm = (name or "").strip().lower() + "|" + (address or "").strip().lower()
+    return hashlib.sha256(norm.encode("utf-8")).hexdigest()[:16]
+
+
+def load_progress() -> Dict[str, Dict[str, Any]]:
+    if not os.path.exists(PROGRESS_PATH):
+        return {}
+    try:
+        with open(PROGRESS_PATH, "r", encoding="utf-8") as f:
+            return json.load(f)
+    except Exception:
+        return {}
+
+
+def save_progress(progress: Dict[str, Dict[str, Any]]):
+    with WRITE_LOCK:
+        with open(PROGRESS_PATH, "w", encoding="utf-8") as f:
+            json.dump(progress, f, ensure_ascii=False, indent=2)
+
+
+def append_jsonl(path: str, record: Dict[str, Any]):
+    line = json.dumps(record, ensure_ascii=False)
+    with WRITE_LOCK:
+        with open(path, "a", encoding="utf-8") as f:
+            f.write(line + "\n")
+
+
+# ----------------------
+# Image utilities
+# ----------------------
+
+def load_image(path: str) -> Image.Image:
+    img = Image.open(path).convert("RGB")
+    return img
+
+
+def draw_grid_overlay(
+    base_img: Image.Image,
+    tiles_selected: List[int] | None,
+    rows: int = TILE_ROWS,
+    cols: int = TILE_COLS,
+    show_numbers: bool = True,
+) -> Image.Image:
+    """Draw a grid and (optionally) a semi-transparent overlay over selected tiles.
+    tiles_selected: list of 1-based tile indices
+    """
+    # Ensure PIL Image (Gradio may pass numpy arrays)
+    if isinstance(base_img, np.ndarray):
+        arr = base_img
+        if arr.dtype != np.uint8:
+            # Normalize float [0,1] to [0,255] or clip other ranges
+            if arr.max() <= 1.0:
+                arr = (arr * 255).clip(0, 255).astype(np.uint8)
+            else:
+                arr = np.clip(arr, 0, 255).astype(np.uint8)
+        if arr.ndim == 2:
+            img = Image.fromarray(arr, mode="L").convert("RGB")
+        else:
+            img = Image.fromarray(arr).convert("RGB")
+    else:
+        img = base_img.copy()
+    W, H = img.size
+    draw = ImageDraw.Draw(img, "RGBA")
+
+    # grid lines
+    for r in range(1, rows):
+        y = int(H * r / rows)
+        draw.line([(0, y), (W, y)], fill=(255, 255, 255, 180), width=2)
+    for c in range(1, cols):
+        x = int(W * c / cols)
+        draw.line([(x, 0), (x, H)], fill=(255, 255, 255, 180), width=2)
+
+    # overlay selected
+    tiles_selected = tiles_selected or []
+    for idx in tiles_selected:
+        idx0 = idx - 1
+        r = idx0 // cols
+        c = idx0 % cols
+        x0 = int(W * c / cols)
+        y0 = int(H * r / rows)
+        x1 = int(W * (c + 1) / cols)
+        y1 = int(H * (r + 1) / rows)
+        draw.rectangle([(x0, y0), (x1, y1)], fill=(255, 0, 0, 70), outline=(255, 0, 0, 160), width=3)
+
+    if show_numbers:
+        try:
+            font = ImageFont.load_default()
+        except Exception:
+            font = None
+        for i in range(rows * cols):
+            idx = i + 1
+            r = i // cols
+            c = i % cols
+            x_center = int(W * (c + 0.5) / cols)
+            y_center = int(H * (r + 0.5) / rows)
+            text = str(idx)
+            bbox = draw.textbbox((x_center, y_center), text, font=font, anchor="mm")
+            draw.rectangle(bbox, fill=(0, 0, 0, 110))
+            draw.text((x_center, y_center), text, fill=(255, 255, 255, 230), anchor="mm", font=font)
+
+    return img
+
+
+# ----------------------
+# App state helpers
+# ----------------------
+
+STEP_DETECTION = 1
+STEP_CONFIDENCE = 2
+STEP_TILES = 3
+STEP_CLARITY = 4
+STEP_REASONING = 5
+
+
+def pick_next_index(user_seen: List[str], samples: List[Sample]) -> int:
+    seen_set = set(user_seen)
+    remaining = [i for i, s in enumerate(samples) if s.sample_id not in seen_set]
+    if remaining:
+        return random.choice(remaining)
+    else:
+        return random.randrange(len(samples))
+
+
+def get_test_image_path(sample: Sample) -> str:
+    return sample.sr_path if TEST_IMAGE_SOURCE == "sr" else sample.gt_path
+
+
+def allowed_steps(detect_choice: str | None) -> List[int]:
+    steps = [STEP_DETECTION, STEP_CONFIDENCE]
+    if detect_choice == "Yes":
+        steps.append(STEP_TILES)
+    steps.extend([STEP_CLARITY, STEP_REASONING])
+    return steps
+
+
+def step_display_index(step: int, detect_choice: str | None) -> tuple[int, int]:
+    seq = allowed_steps(detect_choice)
+    total = len(seq)
+    if step not in seq:
+        # if on TILES but detect != Yes, normalize to CLARITY
+        step = STEP_CLARITY
+    idx = seq.index(step) + 1
+    return idx, total
+
+
+# ----------------------
+# Gradio callbacks (core)
+# ----------------------
+
+def start_or_resume(name: str, address: str):
+    if not name or not address:
+        raise gr.Error("Please enter your name and address to begin.")
+
+    ensure_paths()
+    samples = load_dataset(SR_DIR, GT_DIR)
+
+    uid = hash_user_id(name, address)
+    progress = load_progress()
+    if uid not in progress:
+        progress[uid] = {"seen": []}
+        save_progress(progress)
+
+    user_seen: List[str] = progress[uid].get("seen", [])
+
+    # idx = pick_next_index(user_seen, samples)
+    # sample = samples[idx]
+    idx = reserve_next_sample(uid, samples)
+    sample = samples[idx]
+
+    # Randomize A/B mapping
+    a_is_sr = bool(random.getrandbits(1))
+
+    # Prepare images
+    test_img = load_image(get_test_image_path(sample))
+    overlay_img = draw_grid_overlay(test_img, [])
+
+    img_a = load_image(sample.sr_path if a_is_sr else sample.gt_path)
+    img_b = load_image(sample.gt_path if a_is_sr else sample.sr_path)
+
+    # left_in_round = max(0, len(samples) - len(set(user_seen))) if len(set(user_seen)) < len(samples) else 0
+
+    # status = (
+    #     f"Welcome, {name}. Images left in your first round: {left_in_round}.\n"
+    #     f"Current sample: {sample.sample_id}"
+    # )
+    rnum, left = get_round_and_left(len(samples))
+    status = (
+        f"Welcome, {name}. Global round {rnum} — images left in this round: {left}.\n"
+        f"Current sample: {sample.sample_id}"
+    )
+
+    os.makedirs(RESULTS_DIR, exist_ok=True)
+    user_file_path = os.path.join(RESULTS_DIR, f"{uid}.jsonl")
+
+    return (
+        uid, samples, user_seen, idx, a_is_sr,
+        test_img, overlay_img, img_a, img_b, status,
+        user_file_path,
+        gr.update(visible=True),  # reveal the evaluation panel
+        [],  # reset tile selections
+        1,   # current step -> detection
+        gr.update(visible=False),  # hide intro
+        gr.update(visible=False),  # hide start form
+    )
+
+
+def update_overlay(test_img: Image.Image, selected_tiles: List[str] | None):
+    tiles = sorted([int(x) for x in (selected_tiles or [])])
+    if test_img is None:
+        return None
+    return draw_grid_overlay(test_img, tiles)
+
+
+def _save_record_and_progress(
+    name: str,
+    address: str,
+    uid: str,
+    samples: List[Sample],
+    user_seen: List[str],
+    idx: int,
+    a_is_sr: bool,
+    detection_choice: str,
+    confidence: int,
+    selected_tiles: List[str] | None,
+    clarity_choice: str,
+    reasoning: str,
+):
+    """Common save logic used by both 'finish' and 'next image' paths."""
+    if not name or not address:
+        raise gr.Error("Please enter your name and address.")
+    if detection_choice not in {"Yes", "No", "Unsure"}:
+        detection_choice = "Unsure"
+    valid_clarity = {
+        "Image A is much clearer",
+        "Image A is slightly clearer",
+        "Both are about the same",
+        "Image B is slightly clearer",
+        "Image B is much clearer",
+    }
+    if clarity_choice not in valid_clarity:
+        clarity_choice = None
+
+    confidence = int(confidence)
+    selected_tiles_int = sorted([int(x) for x in (selected_tiles or [])])
+    sample = samples[idx]
+
+    record: Dict[str, Any] = {
+        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+        "user_id": uid,
+        "name": name if SAVE_PII else None,
+        "address": address if SAVE_PII else None,
+        "sample_id": sample.sample_id,
+        "sr_path": sample.sr_path,
+        "gt_path": sample.gt_path,
+        "test_image_source": TEST_IMAGE_SOURCE,
+        "ab_mapping": {
+            "image_a": "SR" if a_is_sr else "GT",
+            "image_b": "GT" if a_is_sr else "SR",
+        },
+        "responses": {
+            "contains_rts": detection_choice,
+            "confidence": confidence,
+            "selected_tiles": selected_tiles_int if detection_choice == "Yes" else [],
+            "clarity": clarity_choice,
+            "reasoning": reasoning or "",
+        },
+    }
+
+    os.makedirs(RESULTS_DIR, exist_ok=True)
+    append_jsonl(os.path.join(RESULTS_DIR, f"{uid}.jsonl"), record)
+    append_jsonl(ALL_RESULTS_JSONL, record)
+
+    progress = load_progress()
+    progress.setdefault(uid, {"seen": []})
+    seen = set(progress[uid].get("seen", []))
+    if sample.sample_id not in seen:
+        seen.add(sample.sample_id)
+        progress[uid]["seen"] = sorted(list(seen))
+        save_progress(progress)
+
+    return progress
+
+
+def submit_finish(
+    name: str,
+    address: str,
+    uid: str,
+    samples: List[Sample],
+    user_seen: List[str],
+    idx: int,
+    a_is_sr: bool,
+    detection_choice: str,
+    confidence: int,
+    selected_tiles: List[str] | None,
+    clarity_choice: str,
+    reasoning: str,
+):
+    """Save and show Thank You (no auto-next)."""
+    progress = _save_record_and_progress(
+        name, address, uid, samples, user_seen, idx, a_is_sr,
+        detection_choice, confidence, selected_tiles, clarity_choice, reasoning
+    )
+
+    sample = samples[idx]
+    # status = f"Saved! Current sample {sample.sample_id} recorded."
+    # NEW — mark global done & release reservation
+    round_num, left, rolled = mark_done_and_maybe_rollover(sample.sample_id, len(samples))
+    release_reservation(sample.sample_id)
+
+    if rolled:
+        status = f"Saved! Global round {round_num-1} complete — starting Round {round_num}."
+    else:
+        status = f"Saved! Global round {round_num} — images left: {left}."
+
+    # Return current state (eval panel will be hidden by to_thanks)
+    return (
+        progress[uid]["seen"],
+        idx,
+        a_is_sr,
+        load_image(get_test_image_path(sample)),
+        draw_grid_overlay(load_image(get_test_image_path(sample)), []),
+        load_image(sample.sr_path if a_is_sr else sample.gt_path),
+        load_image(sample.gt_path if a_is_sr else sample.sr_path),
+        status,
+        [],
+        gr.update(value=""),
+        gr.update(value=None),
+        gr.update(value=3),
+        gr.update(value=None),
+        # 1,
+    )
+
+
+def submit_next_image(
+    name: str,
+    address: str,
+    uid: str,
+    samples: List[Sample],
+    user_seen: List[str],
+    idx: int,
+    a_is_sr: bool,
+    detection_choice: str,
+    confidence: int,
+    selected_tiles: List[str] | None,
+    clarity_choice: str,
+    reasoning: str,
+):
+    """Save and immediately load the next sample."""
+    progress = _save_record_and_progress(
+        name, address, uid, samples, user_seen, idx, a_is_sr,
+        detection_choice, confidence, selected_tiles, clarity_choice, reasoning
+    )
+    # NEW — mark global done & release current reservation
+    current_sample = samples[idx]
+    round_num, left, rolled = mark_done_and_maybe_rollover(current_sample.sample_id, len(samples))
+    release_reservation(current_sample.sample_id)
+
+
+    # Pick next sample
+    # idx_next = pick_next_index(progress[uid]["seen"], samples)
+    # sample_next = samples[idx_next]
+
+    # NEW: reserve next globally available sample
+    idx_next = reserve_next_sample(uid, samples)
+    sample_next = samples[idx_next]
+
+    remaining_after = [s for s in samples if s.sample_id not in set(progress[uid]["seen"])]
+    round_reset_note = ""
+    if not remaining_after:
+        round_reset_note = " First round complete — starting a new round (images may repeat)."
+
+    a_is_sr_next = bool(random.getrandbits(1))
+
+    test_img_next = load_image(get_test_image_path(sample_next))
+    overlay_img_next = draw_grid_overlay(test_img_next, [])
+    img_a_next = load_image(sample_next.sr_path if a_is_sr_next else sample_next.gt_path)
+    img_b_next = load_image(sample_next.gt_path if a_is_sr_next else sample_next.sr_path)
+
+    # left_in_round = max(0, len(samples) - len(set(progress[uid]["seen"])))
+    # status = (
+    #     f"Saved! Images left in your first round: {left_in_round}.{round_reset_note}\n"
+    #     f"Next sample: {sample_next.sample_id}"
+    # )
+
+    rnum, left_after = get_round_and_left(len(samples))
+
+    round_reset_note = ""
+    if rolled:
+        round_reset_note = f" Round {round_num-1} complete — starting Round {round_num}."
+
+    status = (
+        f"Saved! Global round {rnum} — images left: {left_after}.{round_reset_note}\n"
+        f"Next sample: {sample_next.sample_id}"
+    )
+
+    return (
+        progress[uid]["seen"],
+        idx_next,
+        a_is_sr_next,
+        test_img_next,
+        overlay_img_next,
+        img_a_next,
+        img_b_next,
+        status,
+        [],
+        gr.update(value=""),
+        gr.update(value=None),
+        gr.update(value=3),
+        gr.update(value=None),
+        1,
+    )
+
+
+# ----------------------
+# Wizard helpers (visibility + navigation)
+# ----------------------
+
+def step_label(step: int, detect_choice: str | None) -> str:
+    names = {
+        STEP_DETECTION: "RTS detection",
+        STEP_CONFIDENCE: "Confidence",
+        STEP_TILES: "Select tiles",
+        STEP_CLARITY: "A/B clarity",
+        STEP_REASONING: "Reasoning",
+    }
+    i, total = step_display_index(step, detect_choice)
+    return f"**Step {i} of {total} — {names.get(step, '')}**"
+
+
+def apply_step_visibility(current_step: int, detect_choice: str | None):
+    # Normalize step if needed
+    seq = allowed_steps(detect_choice)
+    if current_step not in seq:
+        current_step = STEP_CLARITY
+
+    # Determine which groups are visible
+    show_detection = current_step == STEP_DETECTION
+    show_confidence = current_step == STEP_CONFIDENCE
+    show_tiles = current_step == STEP_TILES and detect_choice == "Yes"
+    show_clarity = current_step == STEP_CLARITY
+    show_reasoning = current_step == STEP_REASONING
+
+    # Nav buttons
+    idx, total = step_display_index(current_step, detect_choice)
+    show_back = idx > 1
+    show_next = idx < total  # Submit buttons only on last step
+
+    return (
+        gr.update(value=step_label(current_step, detect_choice)),
+        gr.update(visible=show_detection),
+        gr.update(visible=show_confidence),
+        gr.update(visible=show_tiles),
+        gr.update(visible=show_clarity),
+        gr.update(visible=show_reasoning),
+        gr.update(visible=show_back),
+        gr.update(visible=show_next),
+    )
+
+
+def go_next_step(current_step: int, detect_choice: str | None, clarity_choice: str | None):
+    # Basic validation at key steps
+    if current_step == STEP_DETECTION and detect_choice not in {"Yes", "No", "Unsure"}:
+        raise gr.Error("Please answer the RTS detection question before continuing.")
+    if current_step == STEP_CLARITY and clarity_choice not in {
+        "Image A is much clearer",
+        "Image A is slightly clearer",
+        "Both are about the same",
+        "Image B is slightly clearer",
+        "Image B is much clearer",
+    }:
+        raise gr.Error("Please select your A/B clarity rating before continuing.")
+
+    seq = allowed_steps(detect_choice)
+    if current_step not in seq:
+        current_step = seq[0]
+    pos = seq.index(current_step)
+    next_pos = min(pos + 1, len(seq) - 1)
+    return seq[next_pos]
+
+
+def go_prev_step(current_step: int, detect_choice: str | None):
+    seq = allowed_steps(detect_choice)
+    if current_step not in seq:
+        current_step = seq[0]
+    pos = seq.index(current_step)
+    prev_pos = max(pos - 1, 0)
+    return seq[prev_pos]
+
+
+def reset_to_step1():
+    return 1
+
+
+# --- Thank-you / restart helpers ---
+def to_thanks(name: str):
+    """Hide eval panel and show the thank-you screen with a personalized message."""
+    msg = f"### ✅ Thanks, {name}! Your responses were saved.\n\nClick **Start Again** to evaluate another image."
+    return gr.update(visible=False), gr.update(visible=True), gr.update(value=msg)
+
+
+def hide_thanks():
+    """Hide the thank-you screen (used when starting/resuming)."""
+    return gr.update(visible=False)
+
+
+# ----------------------
+# UI
+# ----------------------
+with gr.Blocks(title="RTS Human Evaluation", theme=gr.themes.Soft()) as demo:
+    intro_md = gr.Markdown(
+        """
+        # Retrogressive Thaw Slump (RTS) Human Evaluation
+
+        **Instructions:**
+        1. Enter your name and address, then click **Start / Resume**.
+        2. Use **Next** and **Back** to go through the questions.
+        3. On the **last step**, choose **Submit & Next Image** to continue, or **Submit & Finish** to end and see a Thank You screen (with **Start Again**).
+        """
+    )
+
+    # Hidden states
+    state_uid = gr.State("")
+    state_samples = gr.State([])
+    state_seen = gr.State([])
+    state_idx = gr.State(-1)
+    state_a_is_sr = gr.State(True)
+    state_step = gr.State(1)  # wizard step
+
+    with gr.Group() as start_group:
+        with gr.Row():
+            name = gr.Textbox(label="Full name", placeholder="Jane Doe", autofocus=True)
+            address = gr.Textbox(label="Address", placeholder="123 Pine St, City / or Email")
+            start_btn = gr.Button("Start / Resume", variant="primary")
+        status = gr.Markdown("\n")
+
+    eval_panel = gr.Group(visible=False)
+    with eval_panel:
+        step_status = gr.Markdown(visible=True)
+
+        # Step 1 — Detection
+        group_detection = gr.Group(visible=False)
+        with group_detection:
+            with gr.Row():
+                test_img = gr.Image(label="Test Image", interactive=False, type="pil")
+            detect = gr.Radio(
+                label="Does this image contain a Retrogressive Thaw Slump?",
+                choices=["Yes", "No", "Unsure"],
+            )
+
+        # Step 2 — Confidence
+        group_confidence = gr.Group(visible=False)
+        with group_confidence:
+            conf = gr.Slider(1, 5, value=3, step=1, label="Confidence (1=low, 5=high)")
+
+        # Step 3 — Tiles (only if detect == "Yes")
+        group_tiles = gr.Group(visible=False)
+        with group_tiles:
+            gr.Markdown(
+                f"If you answered **Yes**, select tiles around the main RTS features (grid {TILE_ROWS}×{TILE_COLS} = {N_TILES})."
+            )
+            with gr.Row():
+                overlay_img = gr.Image(label="Tile Overlay (selected in red)", interactive=False)
+            tile_choices = [str(i) for i in range(1, N_TILES + 1)]
+            tiles = gr.CheckboxGroup(
+                choices=tile_choices,
+                label="Select tiles",
+            )
+            tiles.change(update_overlay, inputs=[test_img, tiles], outputs=[overlay_img])
+
+        # Step 4 — A/B clarity
+        group_clarity = gr.Group(visible=False)
+        with group_clarity:
+            gr.Markdown("Which image provides a clearer overall view of the geological features?")
+            with gr.Row():
+                image_a = gr.Image(label="Image A", interactive=False)
+                image_b = gr.Image(label="Image B", interactive=False)
+            clarity = gr.Radio(
+                label="Clarity (A vs B)",
+                choices=[
+                    "Image A is much clearer",
+                    "Image A is slightly clearer",
+                    "Both are about the same",
+                    "Image B is slightly clearer",
+                    "Image B is much clearer",
+                ],
+            )
+
+        # Step 5 — Reasoning + Submit (only here)
+        group_reasoning = gr.Group(visible=False)
+        with group_reasoning:
+            reasoning = gr.Textbox(label="Briefly explain your reasoning (artifacts, enhancements, etc.)", lines=3)
+            with gr.Row():
+                submit_next_btn = gr.Button("Submit & Next Image", variant="secondary")
+                submit_finish_btn = gr.Button("Submit & Finish", variant="primary")
+            your_jsonl_path = gr.Textbox(label="Your results file path (for reference)", interactive=False)
+
+        # Nav row (no global Submit here)
+        with gr.Row():
+            back_btn = gr.Button("Back")
+            next_btn = gr.Button("Next", variant="secondary")
+
+    # Thank You screen
+    with gr.Group(visible=False) as thanks_group:
+        thanks_md = gr.Markdown("### ✅ Thanks! Your responses were saved.\n\nClick **Start Again** to evaluate another image.")
+        restart_btn = gr.Button("Start Again", variant="primary")
+
+    # --- Wiring ---
+    start_event = start_btn.click(
+        start_or_resume,
+        inputs=[name, address],
+        outputs=[
+            state_uid, state_samples, state_seen, state_idx, state_a_is_sr,
+            test_img, overlay_img, image_a, image_b, status, your_jsonl_path, eval_panel, tiles, state_step, intro_md, start_group
+        ],
+    )
+
+    # Hide Thank You (if it was open) and show step 1 visibility
+    start_event.then(hide_thanks, inputs=None, outputs=[thanks_group])
+    start_event.then(
+        apply_step_visibility,
+        inputs=[state_step, detect],
+        outputs=[step_status, group_detection, group_confidence, group_tiles, group_clarity, group_reasoning, back_btn, next_btn],
+    )
+
+    # Navigation: Next
+    next_click = next_btn.click(
+        go_next_step,
+        inputs=[state_step, detect, clarity],
+        outputs=[state_step],
+    )
+    next_click.then(
+        apply_step_visibility,
+        inputs=[state_step, detect],
+        outputs=[step_status, group_detection, group_confidence, group_tiles, group_clarity, group_reasoning, back_btn, next_btn],
+    )
+
+    # Navigation: Back
+    back_click = back_btn.click(
+        go_prev_step,
+        inputs=[state_step, detect],
+        outputs=[state_step],
+    )
+    back_click.then(
+        apply_step_visibility,
+        inputs=[state_step, detect],
+        outputs=[step_status, group_detection, group_confidence, group_tiles, group_clarity, group_reasoning, back_btn, next_btn],
+    )
+
+    # Submit & Finish -> save, then show Thank You
+    finish_event = submit_finish_btn.click(
+        submit_finish,
+        inputs=[
+            name,
+            address,
+            state_uid,
+            state_samples,
+            state_seen,
+            state_idx,
+            state_a_is_sr,
+            detect,
+            conf,
+            tiles,
+            clarity,
+            reasoning,
+        ],
+        outputs=[
+            state_seen,
+            state_idx,
+            state_a_is_sr,
+            test_img,
+            overlay_img,
+            image_a,
+            image_b,
+            status,
+            tiles,
+            reasoning,
+            detect,
+            conf,
+            clarity,
+        ],
+    )
+    finish_event.then(
+        to_thanks,
+        inputs=[name],
+        outputs=[eval_panel, thanks_group, thanks_md],
+    )
+
+    # Submit & Next Image -> save, load next image, reset to step 1
+    nextimg_event = submit_next_btn.click(
+        submit_next_image,
+        inputs=[
+            name,
+            address,
+            state_uid,
+            state_samples,
+            state_seen,
+            state_idx,
+            state_a_is_sr,
+            detect,
+            conf,
+            tiles,
+            clarity,
+            reasoning,
+        ],
+        outputs=[
+            state_seen,
+            state_idx,
+            state_a_is_sr,
+            test_img,
+            overlay_img,
+            image_a,
+            image_b,
+            status,
+            tiles,
+            reasoning,
+            detect,
+            conf,
+            clarity,
+            state_step, 
+        ],
+    )
+    # nextimg_event.then(reset_to_step1, inputs=None, outputs=[state_step])
+    nextimg_event.then(
+    apply_step_visibility,
+    inputs=[state_step, detect],
+    outputs=[step_status, group_detection, group_confidence, group_tiles, group_clarity, group_reasoning, back_btn, next_btn],
+)
+    nextimg_event.then(
+        apply_step_visibility,
+        inputs=[state_step, detect],
+        outputs=[step_status, group_detection, group_confidence, group_tiles, group_clarity, group_reasoning, back_btn, next_btn],
+    )
+
+    # Restart: start again using current name/address
+    restart_event = restart_btn.click(
+        start_or_resume,
+        inputs=[name, address],
+        outputs=[
+            state_uid, state_samples, state_seen, state_idx, state_a_is_sr,
+            test_img, overlay_img, image_a, image_b, status, your_jsonl_path, eval_panel, tiles, state_step, intro_md, start_group
+        ],
+    )
+    restart_event.then(hide_thanks, inputs=None, outputs=[thanks_group])
+    restart_event.then(
+        apply_step_visibility,
+        inputs=[state_step, detect],
+        outputs=[step_status, group_detection, group_confidence, group_tiles, group_clarity, group_reasoning, back_btn, next_btn],
+    )
+
+
+if __name__ == "__main__":
+    ensure_paths()
+    _ = load_dataset(SR_DIR, GT_DIR)
+    demo.queue()
+    demo.launch()