deploy app, storage, readme

app.py

@@ -1,192 +1,111 @@
 import gradio as gr
 import random
-import json
-import os
-import time
 import threading
-import requests
-
-# -- Dataset config ---------------------------------------------------------
-# Each entry:
-#   fetch_pair  : callable(group_name, group_cfg) -> (row_a, row_b, common_tags)
-#   get_id      : callable(row) -> unique image identifier (ELO key)
-#   get_image   : callable(row) -> image URL
-#   groups      : dict of group_name -> list of API tag strings
-#                 Both images in a pair come from the same group — prevents
-#                 cross-group bias (e.g. rating male vs female images).
-
-E621_API = "https://e621.net/posts.json"
-E621_HEADERS = {"User-Agent": "ImageRater/1.0 (HuggingFace Space)"}
-E621_MEDIA_PREFIXES = ("https://static1.e621.net/", "https://static1.e926.net/")
-E621_IMAGE_EXTS = {"jpg", "jpeg", "png", "gif", "webp"}
-_e621_last_req = 0.0
-_e621_lock = threading.Lock()
-
-def _e621_request(tags: list[str], limit: int) -> list[dict]:
-    global _e621_last_req
-    with _e621_lock:
-        # respect sustained ≤1 req/sec
-        gap = time.monotonic() - _e621_last_req
-        if gap < 1.0:
-            time.sleep(1.0 - gap)
-        resp = requests.get(
-            E621_API,
-            headers=E621_HEADERS,
-            params={"tags": " ".join(tags + ["order:random"]), "limit": limit},
-            timeout=10,
-        )
-        _e621_last_req = time.monotonic()
-    resp.raise_for_status()
-    return resp.json()["posts"]
-
-def _is_e621_media_url(url: str | None) -> bool:
-    return isinstance(url, str) and url.startswith(E621_MEDIA_PREFIXES)
-
-def _e621_media_url(row: dict) -> str | None:
-    # Prefer sample URL for faster client-side loading, fallback to original file URL.
-    sample_url = row.get("sample", {}).get("url")
-    file_url = row.get("file", {}).get("url")
-    if _is_e621_media_url(sample_url):
+import uuid
+import os
+import html
+from pathlib import Path
+
+import pandas as pd
+from huggingface_hub import hf_hub_download
+
+from storage import VoteStorage
+
+LOCAL_DATA_DIR = 'data'
+DEBUG_MODE = os.getenv("DEBUG", "0").lower() in ("1", "true", "yes", "on")
+VOTE_STORAGE = VoteStorage(mode="local" if DEBUG_MODE else "hf", local_dir=LOCAL_DATA_DIR)
+
+# -- Pool dataset -----------------------------------------------------------
+if DEBUG_MODE:
+    _pool_path = str(Path(__file__).resolve().parent / LOCAL_DATA_DIR / "pool.parquet")
+    assert Path(_pool_path).exists(), f"Missing local debug pool file: {_pool_path}"
+else:
+    _pool_path = hf_hub_download(
+        repo_id="taigasan/e6-visual-ratings",
+        filename="pool.parquet",
+        repo_type="dataset",
+    )
+_pool_df = pd.read_parquet(_pool_path)
+_pool_group_dfs = {g: gdf for g, gdf in _pool_df.groupby("group")}
+
+def _pool_fetch_pair(group_name: str) -> tuple:
+    gdf = _pool_group_dfs[group_name]
+    assert len(gdf) >= 2, f"Not enough rows for group: {group_name}"
+    sample = gdf.sample(2)
+    return sample.iloc[0], sample.iloc[1]
+
+
+def _row_image_url(row) -> str:
+    sample_url = row.get("sample_url")
+    if isinstance(sample_url, str) and sample_url:
         return sample_url
-    if _is_e621_media_url(file_url):
-        return file_url
-    return None
-
-def _valid_image_post(row: dict) -> bool:
-    ext = row.get("file", {}).get("ext")
-    return ext in E621_IMAGE_EXTS and _e621_media_url(row) is not None
-
-def _e621_fetch_pair(group_tags: list[str]) -> tuple:
-    posts = _e621_request(group_tags, limit=20)
-    valid = [p for p in posts if _valid_image_post(p)]
-    assert len(valid) >= 2, f"Not enough image posts for tags: {group_tags}"
-    row_a, row_b = valid[0], valid[1]
-    tags_a = set(row_a["tags"]["general"] + row_a["tags"].get("species", []) + row_a["tags"].get("character", []))
-    tags_b = set(row_b["tags"]["general"] + row_b["tags"].get("species", []) + row_b["tags"].get("character", []))
-    common = sorted(tags_a & tags_b)
-    return row_a, row_b, common
+    image_url = row.get("image_url")
+    if isinstance(image_url, str) and image_url:
+        return image_url
+    return ''
 
 DATASETS: dict[str, dict] = {
-    "e621": {
-        "fetch_pair": _e621_fetch_pair,
-        "get_id": lambda row: row["file"]["md5"],
-        "get_image": _e621_media_url,
-        "groups": {
-            "e_male":        ["male", "solo", "rating:e"],
-            "e_female":      ["female", "solo", "rating:e"],
-            "e_male_female": ["male", "female", "rating:e"],
-            "q_male":        ["male", "solo", "rating:q"],
-            "q_female":      ["female", "solo", "rating:q"],
-            "s_male":        ["male", "solo", "rating:s"],
-            "s_female":      ["female", "solo", "rating:s"],
-        },
+    "pool": {
+        "fetch_pair": _pool_fetch_pair,
+        "get_id": lambda row: row["md5"],
+        "get_image": _row_image_url,
+        "groups": {g: g for g in sorted(_pool_df["group"].unique())},
     },
 }
-
-RATINGS_FILE = "elo_ratings.json"
-DEFAULT_ELO = 1500
-K = 32
-RATINGS_MEM: dict[str, int] = {}
-RATING_PREFIX = {
-    "safe": "s_",
-    "questionable": "q_",
-    "explicit": "e_",
-    "all": None,
-}
-
-# -- Prefetch ---------------------------------------------------------------
-
-_prefetch: dict[tuple[str, str], tuple | None] = {}
-_prefetch_threads: dict[tuple[str, str], threading.Thread] = {}
+DEFAULT_DATASET = list(DATASETS.keys())[0]
 
 def _select_groups(cfg: dict, rating_pref: str) -> list[str]:
-    prefix = RATING_PREFIX[rating_pref]
     groups = list(cfg["groups"].keys())
-    if prefix is None:
+    if rating_pref == "all":
         return groups
-    return [g for g in groups if g.startswith(prefix)]
-
-def _do_prefetch(dataset_name: str, rating_pref: str):
-    try:
-        cfg = DATASETS[dataset_name]
-        groups = _select_groups(cfg, rating_pref)
-        assert groups, f"No groups for rating preference: {rating_pref}"
-        group = random.choice(groups)
-        row_a, row_b, common = cfg["fetch_pair"](cfg["groups"][group])
-        _prefetch[(dataset_name, rating_pref)] = (row_a, row_b, common, group)
-    except Exception:
-        _prefetch[(dataset_name, rating_pref)] = None
-
-def prefetch(dataset_name: str, rating_pref: str):
-    key = (dataset_name, rating_pref)
-    _prefetch[key] = None
-    t = threading.Thread(target=_do_prefetch, args=(dataset_name, rating_pref), daemon=True)
-    _prefetch_threads[key] = t
-    t.start()
-
-def consume_prefetch(dataset_name: str, rating_pref: str) -> tuple:
-    key = (dataset_name, rating_pref)
-    # Wait for prefetch to finish (should be near-instant since we started it earlier)
-    t = _prefetch_threads.get(key)
-    if t:
-        t.join(timeout=15)
-    result = _prefetch.pop(key, None)
-    # Kick off the next prefetch immediately
-    prefetch(dataset_name, rating_pref)
-    if result is not None:
-        return result
-    # Fallback: fetch synchronously if prefetch failed
-    cfg = DATASETS[dataset_name]
-    groups = _select_groups(cfg, rating_pref)
-    assert groups, f"No groups for rating preference: {rating_pref}"
-    group = random.choice(groups)
-    row_a, row_b, common = cfg["fetch_pair"](cfg["groups"][group])
-    return row_a, row_b, common, group
-
-# -- ELO helpers ------------------------------------------------------------
-
-def load_ratings() -> dict:
-    return RATINGS_MEM.copy()
-
-def save_ratings(ratings: dict):
-    # Stubbed persistence for now: keep ratings only in memory.
-    RATINGS_MEM.clear()
-    RATINGS_MEM.update(ratings)
-
-def elo_update(ratings: dict, winner_key: str, loser_key: str) -> dict:
-    rw = ratings.get(winner_key, DEFAULT_ELO)
-    rl = ratings.get(loser_key, DEFAULT_ELO)
-    ea = 1 / (1 + 10 ** ((rl - rw) / 400))
-    ratings[winner_key] = round(rw + K * (1 - ea))
-    ratings[loser_key] = round(rl + K * (0 - (1 - ea)))
-    return ratings
+    return [g for g in groups if g.endswith(f"_{rating_pref}")]
 
 def _commit_oldest_pending(state: dict):
     pending = state.setdefault("pending", [])
-    if len(pending) <= 2:
+    if len(pending) <= 1:
         return
     oldest = pending.pop(0)
-    winner = oldest.get("winner")
-    if winner is None:
-        return
-    ratings = load_ratings()
-    if winner == "A":
-        ratings = elo_update(ratings, oldest["key_a"], oldest["key_b"])
-    else:
-        ratings = elo_update(ratings, oldest["key_b"], oldest["key_a"])
-    save_ratings(ratings)
+    threading.Thread(target=VOTE_STORAGE.append_vote_row, args=(oldest.copy(), oldest.get("winner")), daemon=True).start()
 
 def _render_current(state: dict) -> tuple:
-    return state["url_a"], state["url_b"], state
+    img_a_html = f"<div class=\"rating-card\"><div class=\"rating-card-title\"><strong>Image A</strong></div><div class=\"rating-image-frame\"><img src=\"{html.escape(state['url_a'])}\" class=\"rating-image\" loading=\"eager\" referrerpolicy=\"no-referrer\"></div></div>"
+    img_b_html = f"<div class=\"rating-card\"><div class=\"rating-card-title\"><strong>Image B</strong></div><div class=\"rating-image-frame\"><img src=\"{html.escape(state['url_b'])}\" class=\"rating-image\" loading=\"eager\" referrerpolicy=\"no-referrer\"></div></div>"
+    link_a = f"Image A: https://e621.net/posts/{state['id_a']}"
+    link_b = f"Image B: https://e621.net/posts/{state['id_b']}"
+    can_go_back = bool(state.get("can_go_back"))
+    back_md = "[back](#back)" if can_go_back else "<span class='subtle-back-link-disabled'>back</span>"
+    details = f"<span class='subtle-note'>Group: {state['group']}</span>"
+    return img_a_html, img_b_html, link_a, link_b, back_md, details, state
+
+
+
+def _normalize_rating_pref(pref: str | None) -> str:
+    return pref if pref in ("safe", "all") else "safe"
+
+
+def _initial_load(state: dict, pref: str | None):
+    rating_pref = _normalize_rating_pref(pref)
+    return rating_pref, *new_round(DEFAULT_DATASET, rating_pref, state)
+
+
+def _on_rating_change(rating_pref: str, state: dict):
+    rating_pref = _normalize_rating_pref(rating_pref)
+    return *new_round(DEFAULT_DATASET, rating_pref, state), rating_pref
 
 # -- Gradio callbacks -------------------------------------------------------
 
 def new_round(dataset_name: str, rating_pref: str, state: dict) -> tuple:
     cfg = DATASETS[dataset_name]
-    row_a, row_b, common, group = consume_prefetch(dataset_name, rating_pref)
+    groups = _select_groups(cfg, rating_pref)
+    assert groups, f"No groups for rating preference: {rating_pref}"
+    group = random.choice(groups)
+    row_a, row_b = cfg["fetch_pair"](cfg["groups"][group])
+    state.setdefault("session_id", uuid.uuid4().hex)
     key_a = cfg["get_id"](row_a)
     key_b = cfg["get_id"](row_b)
-    state.update(dataset=dataset_name, rating_pref=rating_pref, key_a=key_a, key_b=key_b)
+    id_a = int(row_a["id"])
+    id_b = int(row_b["id"])
+    state.update(dataset=dataset_name, rating_pref=rating_pref, key_a=key_a, key_b=key_b, id_a=id_a, id_b=id_b, group=group)
     url_a = cfg["get_image"](row_a)
     url_b = cfg["get_image"](row_b)
     state["url_a"] = url_a
@@ -194,19 +113,24 @@ def new_round(dataset_name: str, rating_pref: str, state: dict) -> tuple:
     return _render_current(state)
 
 def _queue_decision(winner: str | None, state: dict):
+    assert state.get("session_id"), "Missing session_id: refusing to record vote"
     state.setdefault("pending", [])
-    state.setdefault("decision_history", [])
     decision = {
         "winner": winner,
         "key_a": state["key_a"],
         "key_b": state["key_b"],
+        "id_a": state["id_a"],
+        "id_b": state["id_b"],
         "url_a": state["url_a"],
         "url_b": state["url_b"],
         "dataset": state["dataset"],
         "rating_pref": state["rating_pref"],
+        "group": state["group"],
+        "session_id": state["session_id"],
     }
     state["pending"].append(decision)
-    state["decision_history"].append(decision)
+    state["last_decision"] = decision
+    state["can_go_back"] = True
     _commit_oldest_pending(state)
 
 def vote(winner: str | None, state: dict) -> tuple:
@@ -215,31 +139,58 @@ def vote(winner: str | None, state: dict) -> tuple:
     return new_round(state["dataset"], state["rating_pref"], state)
 
 def go_back(state: dict) -> tuple:
-    history = state.setdefault("decision_history", [])
     pending = state.setdefault("pending", [])
-    if not history:
+    if not state.get("can_go_back"):
         return _render_current(state)
-    last = history.pop()
-    if pending and pending[-1] is last:
+    last = state.get("last_decision")
+    if not last:
+        state["can_go_back"] = False
+        return _render_current(state)
+    if pending and pending[-1] == last:
         pending.pop()
+    state["can_go_back"] = False
+    state["last_decision"] = None
     state.update(
         dataset=last["dataset"],
         rating_pref=last["rating_pref"],
         key_a=last["key_a"],
         key_b=last["key_b"],
+        id_a=last["id_a"],
+        id_b=last["id_b"],
         url_a=last["url_a"],
         url_b=last["url_b"],
+        group=last["group"],
     )
     return _render_current(state)
 
-# Warm up prefetch for all datasets at startup (safe by default)
-for _ds in DATASETS:
-    prefetch(_ds, "safe")
-
 # -- UI ---------------------------------------------------------------------
 
 with gr.Blocks(
     title="Image Rater",
+    head="""
+    <script>
+    window.addEventListener('keydown', function (e) {
+        const t = e.target;
+        if (t && (t.tagName === 'INPUT' || t.tagName === 'TEXTAREA' || t.isContentEditable)) return;
+        if (e.key === 'ArrowLeft') {
+            e.preventDefault();
+            document.querySelector('#btn-vote-a button, button#btn-vote-a')?.click();
+        } else if (e.key === 'ArrowRight') {
+            e.preventDefault();
+            document.querySelector('#btn-vote-b button, button#btn-vote-b')?.click();
+        } else if (e.key === 'Backspace') {
+            e.preventDefault();
+            document.querySelector('#btn-back-action button, button#btn-back-action')?.click();
+        }
+    });
+    document.addEventListener('click', function (e) {
+        const a = e.target.closest('a[href="#back"]');
+        if (!a) return;
+        e.preventDefault();
+        document.querySelector('#btn-back-action button, button#btn-back-action')?.click();
+    });
+    </script>
+    """,
     css="""
     .subtle-link button {
         background: none !important;
@@ -255,42 +206,95 @@ with gr.Blocks(
     .subtle-link button:hover {
         color: #5a5a5a !important;
     }
+    .subtle-link {
+        width: fit-content !important;
+    }
+    .subtle-link button {
+        width: fit-content !important;
+    }
+    .subtle-note {
+        color: #888;
+        font-size: 0.9em;
+    }
+    .rating-card {
+        width: 100%;
+    }
+    .rating-card-title {
+        min-height: 24px;
+        margin-bottom: 8px;
+    }
+    .rating-image-frame {
+        width: 100%;
+        height: 512px;
+        border: 1px solid #e6e6e6;
+        border-radius: 8px;
+        background: #fafafa;
+        display: flex;
+        align-items: center;
+        justify-content: center;
+        overflow: hidden;
+    }
+    .rating-image {
+        width: 100%;
+        height: 100%;
+        object-fit: contain;
+    }
+    .subtle-back-link-wrap a {
+        color: #7a7a7a !important;
+        font-size: 0.9em;
+        text-decoration: underline;
+    }
+    .subtle-back-link-wrap a:hover {
+        color: #5a5a5a !important;
+    }
+    .subtle-back-link-disabled {
+        color: #b8b8b8 !important;
+        pointer-events: none;
+        text-decoration: none;
+    }
+    .hidden-action-btn {
+        display: none !important;
+    }
     """,
 ) as demo:
-    gr.Markdown("## Image Rater\nPick the image you prefer. Both images are drawn from the same group to avoid cross-group bias.")
+    gr.Markdown("## Image Quality Rater\nRate relative image quality. Choose the image with better quality, or select same quality if they are comparable. Both images are drawn from the same group to avoid cross-group bias.")
 
     state = gr.State({})
+    rating_pref_store = gr.BrowserState(default_value="safe", storage_key="rating_pref")
 
     with gr.Row():
-        img_a = gr.Image(label="Image A", interactive=False, height=512)
-        img_b = gr.Image(label="Image B", interactive=False, height=512)
+        img_a = gr.HTML()
+        img_b = gr.HTML()
 
     with gr.Row():
-        btn_a = gr.Button("👍 Prefer A", variant="primary")
-        btn_skip = gr.Button("About the same")
-        btn_b = gr.Button("👍 Prefer B", variant="primary")
+        btn_a = gr.Button("👍 Prefer A", variant="primary", elem_id="btn-vote-a")
+        btn_skip = gr.Button("Same quality", elem_id="btn-skip")
+        btn_b = gr.Button("👍 Prefer B", variant="primary", elem_id="btn-vote-b")
 
     with gr.Accordion("Settings", open=False):
         gr.Markdown("<span style='color:#888;font-size:0.9em;'>Advanced options</span>")
         rating_dd = gr.Dropdown(
-            choices=["safe", "questionable", "explicit", "all"],
+            choices=["safe", "all"],
             value="safe",
             label="Rating",
+            elem_id="rating-pref",
         )
-        dataset_dd = gr.Dropdown(choices=list(DATASETS.keys()), value=list(DATASETS.keys())[0], label="Dataset")
-        btn_back = gr.Button("Back", elem_classes=["subtle-link"])
-
-    outputs = [img_a, img_b, state]
-
-    btn_a.click(fn=lambda s: vote("A", s), inputs=[state], outputs=outputs)
-    btn_b.click(fn=lambda s: vote("B", s), inputs=[state], outputs=outputs)
-    btn_skip.click(fn=lambda s: vote(None, s), inputs=[state], outputs=outputs)
-    btn_back.click(fn=go_back, inputs=[state], outputs=outputs)
-    img_a.select(fn=lambda s, evt: vote("A", s), inputs=[state], outputs=outputs)
-    img_b.select(fn=lambda s, evt: vote("B", s), inputs=[state], outputs=outputs)
-    dataset_dd.change(fn=new_round, inputs=[dataset_dd, rating_dd, state], outputs=outputs)
-    rating_dd.change(fn=new_round, inputs=[dataset_dd, rating_dd, state], outputs=outputs)
-    demo.load(fn=lambda s: new_round(list(DATASETS.keys())[0], "safe", s), inputs=[state], outputs=outputs)
+
+    link_a = gr.Markdown(label="Image A link")
+    link_b = gr.Markdown(label="Image B link")
+    back_link = gr.Markdown(elem_classes=["subtle-back-link-wrap"])
+    btn_back_action = gr.Button("back", elem_id="btn-back-action", elem_classes=["hidden-action-btn"])
+    details_md = gr.Markdown()
+    gr.Markdown("<span class='subtle-note'>Dataset: <a href='https://huggingface.co/datasets/taigasan/e6-visual-ratings' target='_blank' rel='noopener noreferrer'>taigasan/e6-visual-ratings</a></span>")
+    gr.Markdown("<span class='subtle-note'>Shortcuts: Left = vote A, Right = vote B, Backspace = back</span>")
+    outputs = [img_a, img_b, link_a, link_b, back_link, details_md, state]
+
+    btn_a.click(fn=lambda s: vote("A", s), inputs=[state], outputs=outputs, queue=False, show_progress="hidden")
+    btn_b.click(fn=lambda s: vote("B", s), inputs=[state], outputs=outputs, queue=False, show_progress="hidden")
+    btn_skip.click(fn=lambda s: vote(None, s), inputs=[state], outputs=outputs, queue=False, show_progress="hidden")
+    btn_back_action.click(fn=go_back, inputs=[state], outputs=outputs, queue=False, show_progress="hidden")
+    rating_dd.change(fn=_on_rating_change, inputs=[rating_dd, state], outputs=[*outputs, rating_pref_store], queue=False, show_progress="hidden")
+    demo.load(fn=_initial_load, inputs=[state, rating_pref_store], outputs=[rating_dd, *outputs], queue=False, show_progress="hidden")
 
 if __name__ == "__main__":
     demo.launch()