Allow picking from individual image groups.

app.py

@@ -250,16 +250,6 @@ DATASETS: dict[str, dict] = {
 }
 DEFAULT_DATASET = list(DATASETS.keys())[0]
 
-def _select_groups(cfg: dict, rating_pref: str) -> list[str]:
-    if rating_pref == "all":
-        return cfg["groups"]
-
-    return [
-        g
-        for g in cfg["groups"]
-        if g.endswith(f"_{rating_pref}")
-    ]
-
 def _format_rating_post_title(post_id: int, votes: int, label: str) -> str:
     return f"<strong>{label}</strong>: <a href=\"https://e621.net/posts/{post_id}\" target=\"_blank\" rel=\"noreferrer\">Post #{post_id}</a> | {votes} {'Vote' if votes == 1 else 'Votes'}"
 
@@ -279,15 +269,13 @@ def _render_current(state: dict, submit_status: str = "") -> tuple:
 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, submit_key: str | None, image_height: str):
-    rating_pref = _normalize_rating_pref(pref)
+def _initial_load(state: dict, rating_pref: str | None, submit_key: str | None, image_height: str, groups: list[str]):
+    rating_pref = _normalize_rating_pref(rating_pref)
     submit_key = _normalize_submit_key(submit_key)
-    return rating_pref, submit_key, image_height, image_height, *new_round(DEFAULT_DATASET, rating_pref, state)
+    return rating_pref, submit_key, image_height, image_height, groups, *new_round(DEFAULT_DATASET, groups, 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
+def _on_groups_change(groups: list[str], state: dict):
+    return *new_round(DEFAULT_DATASET, groups, state), groups
 
 def _on_image_height_change(image_height: str) -> tuple[str, str]:
     return image_height, image_height
@@ -298,7 +286,6 @@ def _normalize_submit_key(submit_key: str | None) -> str:
 def _filtered_explorer_df(rating_pref: str) -> pd.DataFrame:
     return _filtered_explorer_df_by_classifier(rating_pref, ALLOWED_CLASSIFIER_FILTERS[0])
 
-
 def _filtered_explorer_df_by_classifier(rating_pref: str, classifier_name: str) -> pd.DataFrame:
     if rating_pref == "all":
         rating_filtered = _explorer_df
@@ -308,11 +295,10 @@ def _filtered_explorer_df_by_classifier(rating_pref: str, classifier_name: str)
     assert classifier_name in ALLOWED_CLASSIFIER_FILTERS, f"Unsupported classifier filter: {classifier_name}"
     return rating_filtered[rating_filtered["classifier"] == classifier_name]
 
-
-def _load_results(rating_pref_value: str, sort_mode_value: str, classifier_filter_value: str):
-    rating_pref = _normalize_rating_pref(rating_pref_value)
-    sort_mode = _normalize_sort_mode(sort_mode_value)
-    classifier_name = _normalize_classifier_filter(classifier_filter_value)
+def _load_results(rating_pref: str, sort_mode: str, classifier_filter: str):
+    rating_pref = _normalize_rating_pref(rating_pref)
+    sort_mode = _normalize_sort_mode(sort_mode)
+    classifier_name = _normalize_classifier_filter(classifier_filter)
     filtered_explorer_df = _filtered_explorer_df_by_classifier(rating_pref, classifier_name)
     summary, score_distribution_plot, distribution_data, gallery_items, page_meta, next_offset, btn_update = build_results_data(
         filtered_explorer_df,
@@ -323,13 +309,11 @@ def _load_results(rating_pref_value: str, sort_mode_value: str, classifier_filte
     )
     return summary, score_distribution_plot, distribution_data, gallery_items, btn_update, "Click an image to reveal its ID and link.", page_meta, next_offset
 
-
 def _normalize_sort_mode(sort_mode: str | None) -> str:
     if sort_mode in ("Default", "Rating: Low to High", "Rating: High to Low"):
         return sort_mode
     return "Default"
 
-
 def _normalize_classifier_filter(classifier_name: str | None) -> str:
     if classifier_name in ALLOWED_CLASSIFIER_FILTERS:
         return str(classifier_name)
@@ -337,11 +321,11 @@ def _normalize_classifier_filter(classifier_name: str | None) -> str:
 
 # -- Gradio callbacks -------------------------------------------------------
 
-def new_round(dataset_name: str, rating_pref: str, state: dict) -> tuple:
-    cfg = DATASETS[dataset_name]
-    groups = _select_groups(cfg, rating_pref)
-    assert groups, f"No groups for rating preference: {rating_pref}"
+def new_round(dataset_name: str, groups: list[str], state: dict) -> tuple:
+    if not groups:
+        return "", "", gr.skip(), "", "Please select at least one group.", state
 
+    cfg = DATASETS[dataset_name]
     group = random.choice(groups)
     row_a, row_b, reason_remaining, pair_reason = cfg["fetch_pair"](group)
 
@@ -352,7 +336,7 @@ def new_round(dataset_name: str, rating_pref: str, state: dict) -> tuple:
     key_b = cfg["get_id"](row_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, pair_reason=pair_reason)
+    state.update(dataset=dataset_name, key_a=key_a, key_b=key_b, id_a=id_a, id_b=id_b, group=group, pair_reason=pair_reason)
     url_a = cfg["get_image"](row_a)
     url_b = cfg["get_image"](row_b)
     state["url_a"] = url_a
@@ -372,7 +356,6 @@ def _queue_decision(winner: str | None, state: dict):
         "url_a": state["url_a"],
         "url_b": state["url_b"],
         "dataset": state["dataset"],
-        "rating_pref": state["rating_pref"],
         "group": state["group"],
         "pair_reason": state.get("pair_reason", ""),
         "session_id": state["session_id"],
@@ -387,10 +370,13 @@ def _add_vote(idx: int, col_loc: int, delta: int = 1) -> None:
     wins, ties, votes = _pool_df.iloc[idx, [WINS_LOC, TIES_LOC, VOTES_LOC]]
     _pool_df.iloc[idx, WINRATE_LOC] = (wins + 0.5 * ties) / max(votes, 1)
 
-def vote(winner: str | None, state: dict, submit_key: str | None) -> tuple:
+def vote(winner: str | None, state: dict, groups: list[str], submit_key: str | None) -> tuple:
     if _normalize_submit_key(submit_key) != SUBMIT_KEY:
         return _render_current(state, "Wrong submission key.")
 
+    if not groups:
+        return "", "", gr.skip(), "", "Please select at least one group.", state
+
     _queue_decision(winner, state)
 
     a_idx = _md5_to_idx[state["key_a"]]
@@ -409,7 +395,7 @@ def vote(winner: str | None, state: dict, submit_key: str | None) -> tuple:
             case _:
                 raise AssertionError
 
-    return new_round(state["dataset"], state["rating_pref"], state)
+    return new_round(state["dataset"], groups, state)
 
 def go_back(state: dict) -> tuple:
     pending = state.setdefault("pending", [])
@@ -418,7 +404,6 @@ def go_back(state: dict) -> tuple:
         last = pending.pop()
         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"],
@@ -674,6 +659,11 @@ with gr.Blocks(
     results_sort_store = gr.BrowserState(default_value="Default", storage_key="results_sort_mode")
     results_classifier_store = gr.BrowserState(default_value=ALLOWED_CLASSIFIER_FILTERS[0], storage_key="results_classifier")
     image_height_store = gr.BrowserState(default_value=768, storage_key="image_height")
+    groups_store = gr.BrowserState(default_value=[
+        group
+        for group in DATASETS[DEFAULT_DATASET]["groups"]
+        if group.endswith("_safe")
+    ], storage_key="groups")
 
     with gr.Tabs():
         with gr.Tab("Image Quality Rater"):
@@ -693,16 +683,16 @@ with gr.Blocks(
                 btn_b = gr.Button("➡️ Prefer B", variant="primary", elem_id="btn-vote-b")
 
             with gr.Accordion("Settings", open=False):
+                groups_select = gr.CheckboxGroup(
+                    choices=DATASETS[DEFAULT_DATASET]["groups"],
+                    label="Categories",
+                    show_label=True,
+                    show_select_all=True
+                )
                 image_height_slider = gr.Slider(
                     minimum=512, maximum=2048, step=16, precision=0,
                     label="Image Size",
                 )
-                rating_dd = gr.Dropdown(
-                    choices=["safe", "all"],
-                    value="safe",
-                    label="Rating",
-                    elem_id="rating-pref",
-                )
                 submit_key_tb = gr.Textbox(
                     value="",
                     type="password",
@@ -717,6 +707,7 @@ with gr.Blocks(
 
         (
             results_summary_md,
+            results_rating_dd,
             results_sort_dd,
             results_classifier_dd,
             results_score_distribution_plot,
@@ -740,23 +731,26 @@ with gr.Blocks(
         results_page_offset_state,
     ]
 
-    btn_a.click(fn=lambda s, k: vote("A", s, k), inputs=[state, submit_key_store], outputs=outputs, queue=False, show_progress="hidden")
-    btn_b.click(fn=lambda s, k: vote("B", s, k), inputs=[state, submit_key_store], outputs=outputs, queue=False, show_progress="hidden")
-    btn_skip.click(fn=lambda s, k: vote(None, s, k), inputs=[state, submit_key_store], outputs=outputs, queue=False, show_progress="hidden")
+    btn_a.click(fn=lambda s, g, k: vote("A", s, g, k), inputs=[state, groups_store, submit_key_store], outputs=outputs, queue=False, show_progress="hidden")
+    btn_b.click(fn=lambda s, g, k: vote("B", s, g, k), inputs=[state, groups_store, submit_key_store], outputs=outputs, queue=False, show_progress="hidden")
+    btn_skip.click(fn=lambda s, g, k: vote(None, s, g, k), inputs=[state, groups_store, submit_key_store], 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")
-    submit_key_tb.input(fn=_normalize_submit_key, inputs=[submit_key_tb], outputs=[submit_key_store], queue=False, show_progress="hidden")
     submit_key_tb.change(fn=_normalize_submit_key, inputs=[submit_key_tb], outputs=[submit_key_store], queue=False, show_progress="hidden")
-    rating_dd.change(fn=_load_results, inputs=[rating_dd, results_sort_store, results_classifier_store], outputs=results_outputs, queue=False, show_progress="hidden")
+    groups_select.change(fn=_on_groups_change, inputs=[groups_select, state], outputs=[*outputs, groups_store], queue=False, show_progress="hidden")
+    image_height_slider.change(fn=_on_image_height_change, inputs=[image_height_slider], outputs=[image_height_store, image_height], queue=False, show_progress="hidden")
+
+    results_rating_dd.change(fn=_normalize_rating_pref, inputs=[results_rating_dd], outputs=[rating_pref_store], queue=False, show_progress="hidden")
+    results_rating_dd.change(fn=_load_results, inputs=[results_rating_dd, results_sort_store, results_classifier_store], outputs=results_outputs, queue=False, show_progress="hidden")
     results_sort_dd.change(fn=_normalize_sort_mode, inputs=[results_sort_dd], outputs=[results_sort_store], queue=False, show_progress="hidden")
     results_sort_dd.change(fn=_load_results, inputs=[rating_pref_store, results_sort_dd, results_classifier_store], outputs=results_outputs, queue=False, show_progress="hidden")
     results_classifier_dd.change(fn=_normalize_classifier_filter, inputs=[results_classifier_dd], outputs=[results_classifier_store], queue=False, show_progress="hidden")
     results_classifier_dd.change(fn=_load_results, inputs=[rating_pref_store, results_sort_store, results_classifier_dd], outputs=results_outputs, queue=False, show_progress="hidden")
-    image_height_slider.change(fn=_on_image_height_change, inputs=[image_height_slider], outputs=[image_height_store, image_height], queue=False, show_progress="hidden")
-    demo.load(fn=_initial_load, inputs=[state, rating_pref_store, submit_key_store, image_height_store], outputs=[rating_dd, submit_key_tb, image_height_slider, image_height, *outputs], queue=False, show_progress="hidden")
+
+    demo.load(fn=_initial_load, inputs=[state, rating_pref_store, submit_key_store, image_height_store, groups_store], outputs=[results_rating_dd, submit_key_tb, image_height_slider, image_height, groups_select, *outputs], queue=False, show_progress="hidden")
     demo.load(fn=_load_results, inputs=[rating_pref_store, results_sort_store, results_classifier_store], outputs=results_outputs, queue=False, show_progress="hidden")
     demo.load(fn=_normalize_sort_mode, inputs=[results_sort_store], outputs=[results_sort_dd], queue=False, show_progress="hidden")
     demo.load(fn=_normalize_classifier_filter, inputs=[results_classifier_store], outputs=[results_classifier_dd], queue=False, show_progress="hidden")
+
     results_load_more_btn.click(
         fn=lambda r, s, c, o: load_more_results(_filtered_explorer_df_by_classifier(_normalize_rating_pref(r), _normalize_classifier_filter(c)), _explorer_df, s, o),
         inputs=[rating_pref_store, results_sort_store, results_classifier_store, results_page_offset_state],

explorer.py

@@ -208,6 +208,11 @@ def add_results_tab(pool_df: pd.DataFrame):
         results_load_more_btn = gr.Button("Load more (ArrowDown)", elem_id="btn-results-load-more")
         selected_image_md = gr.Markdown("Click an image to reveal its ID and link.")
         results_score_distribution_plot = gr.Plot(label="Classifier score distribution")
+        results_rating_dd = gr.Dropdown(
+            choices=["safe", "all"],
+            value="safe",
+            label="Rating",
+        )
         results_sort_dd = gr.Dropdown(
             choices=SORT_MODES,
             value="Default",
@@ -225,6 +230,7 @@ def add_results_tab(pool_df: pd.DataFrame):
         results_page_offset_state = gr.State(0)
     return (
         results_summary_md,
+        results_rating_dd,
         results_sort_dd,
         results_classifier_dd,
         results_score_distribution_plot,