Create app.py

app.py

@@ -0,0 +1,209 @@
+import random
+import pandas as pd
+import gradio as gr
+from typing import Dict, Optional
+
+import unibox as ub
+
+# Store current dataset in a global dict so it persists across Gradio calls.
+CURRENT_DATASET = {
+    "id": None,
+    "df": None
+}
+
+rating_map = {
+    "g": "general",
+    "s": "sensitive",
+    "q": "questionable",
+    "e": "explicit"
+}
+
+def load_dataset_if_needed(dataset_id: str):
+    """
+    Checks if dataset_id is different from what's currently loaded.
+    If so, loads from HF again and updates CURRENT_DATASET.
+    """
+    if CURRENT_DATASET["id"] != dataset_id:
+        df = ub.loads(f"hf://{dataset_id}").to_pandas()
+        CURRENT_DATASET["id"] = dataset_id
+        CURRENT_DATASET["df"] = df
+
+
+def convert_dbr_tag_string(tag_string: str, shuffle: bool = True) -> str:
+    """
+    1girl long_hair blush -> 1girl, long_hair, blush
+    """
+    tags_list = [i.replace("_", " ") for i in tag_string.split(" ") if i]
+    if shuffle:
+        random.shuffle(tags_list)
+    return ", ".join(tags_list)
+
+
+def get_tags_dict(df_row: pd.Series) -> dict:
+    """
+    Returns a dict with rating/artist/character/copyright/general/meta
+    plus numeric score.
+    """
+    rating = df_row["rating"]
+    artist = df_row["tag_string_artist"]
+    character = df_row["tag_string_character"]
+    copyright_ = df_row["tag_string_copyright"]
+    general = df_row["tag_string_general"]
+    meta = df_row["tag_string_meta"]
+    score = df_row["score"]
+
+    rating_str = rating_map.get(rating, "")
+    artist_str = artist if artist else ""
+    character_str = convert_dbr_tag_string(character) if character else ""
+    copyright_str = f"copyright:{copyright_}" if copyright_ else ""
+    general_str = convert_dbr_tag_string(general) if general else ""
+    meta_str = convert_dbr_tag_string(meta) if meta else ""
+    _score = str(score) if score else ""
+
+    return {
+        "rating_str": rating_str,
+        "artist_str": artist_str,
+        "character_str": character_str,
+        "copyright_str": copyright_str,
+        "general_str": general_str,
+        "meta_str": meta_str,
+        "score": _score,
+    }
+
+
+def build_tags_from_tags_dict(tags_dict: dict, add_artist_tags: bool = True) -> str:
+    """
+    Build a final comma-separated string (rating, artist, character, etc.).
+    """
+    context = []
+
+    if tags_dict["rating_str"]:
+        context.append(tags_dict["rating_str"])
+
+    if tags_dict["artist_str"] and add_artist_tags:
+        context.append(f"artist:{tags_dict['artist_str']}")
+
+    if tags_dict["character_str"]:
+        context.append(tags_dict["character_str"])
+
+    if tags_dict["copyright_str"]:
+        context.append(tags_dict["copyright_str"])
+
+    if tags_dict["general_str"]:
+        context.append(tags_dict["general_str"])
+
+    return ", ".join(context)
+
+
+def get_captions_for_rows(df, start_idx: int = 0, end_idx: int = 5,
+                          tags_front: str = "", tags_back: str = "",
+                          add_artist_tags: bool = True) -> list:
+    filtered_df = df.iloc[start_idx:end_idx]
+    captions = []
+    for _, row in filtered_df.iterrows():
+        tags = get_tags_dict(row)
+        caption_base = build_tags_from_tags_dict(tags, add_artist_tags)
+        # Combine front, base, back
+        pieces = [part for part in [tags_front, caption_base, tags_back] if part]
+        final_caption = ", ".join(pieces)
+        captions.append(final_caption)
+    return captions
+
+
+def get_previews_for_rows(df: pd.DataFrame, start_idx: int = 0, end_idx: int = 5) -> list:
+    filtered_df = df.iloc[start_idx:end_idx]
+    return [row["large_file_url"] for _, row in filtered_df.iterrows()]
+
+
+def gradio_interface(
+    dataset_id: str,
+    start_idx: int = 0,
+    display_count: int = 5,
+    tags_front: str = "",
+    tags_back: str = "",
+    add_artist_tags: bool = True
+):
+    """
+    1) Loads dataset if needed
+    2) Returns (DataFrame, Gallery, InfoMessage)
+    """
+    # 1) Possibly reload
+    load_dataset_if_needed(dataset_id)
+    dset_df = CURRENT_DATASET["df"]
+    if dset_df is None:
+        return pd.DataFrame(), [], f"ERROR: Could not load dataset {dataset_id}"
+
+    # 2) Figure out total length, clamp inputs
+    total_len = len(dset_df)
+    if total_len == 0:
+        return pd.DataFrame(), [], f"Dataset {dataset_id} is empty."
+
+    start_idx = max(start_idx, 0)
+    if start_idx >= total_len:
+        start_idx = total_len - 1
+
+    end_idx = start_idx + display_count
+    if end_idx > total_len:
+        end_idx = total_len
+
+    # 3) Build results
+    idxs = range(start_idx, end_idx)
+    captions = get_captions_for_rows(dset_df, start_idx, end_idx, tags_front, tags_back, add_artist_tags)
+    previews = get_previews_for_rows(dset_df, start_idx, end_idx)
+    df_out = pd.DataFrame({"index": idxs, "Captions": captions})
+
+    # 4) Build info string
+    info_msg = (
+        f"**Current dataset:** {CURRENT_DATASET['id']}  \n"
+        f"**Dataset length:** {total_len}  \n"
+        f"**start_idx:** {start_idx}, **display_count:** {display_count}, "
+        f"**tags_front:** '{tags_front}', **tags_back:** '{tags_back}', "
+        f"**add_artist_tags:** {add_artist_tags}"
+    )
+
+    return df_out, previews, info_msg
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("## Danbooru2025 Dataset Captions and Previews")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            dataset_id_input = gr.Textbox(
+                value="dataproc5/test-danbooru2025-tag-balanced-2k", 
+                label="Dataset ID"
+            )
+            start_idx_input = gr.Number(value=0, label="Start Index")
+            display_count_input = gr.Slider(
+                value=5, minimum=1, maximum=50, step=1, 
+                label="Number of Items"
+            )
+            tags_front_input = gr.Textbox(value="", label="Tags Front")
+            tags_back_input = gr.Textbox(value="", label="Tags Back")
+            add_artist_tags_input = gr.Checkbox(label="Add artist tags", value=True)
+
+            run_button = gr.Button("Get Captions & Previews")
+
+        with gr.Column(scale=2):
+            captions_df_out = gr.DataFrame(label="Captions")
+            previews_gallery_out = gr.Gallery(label="Previews", type="filepath")
+            info_textbox_out = gr.Markdown(value="")
+
+    run_button.click(
+        fn=gradio_interface,
+        inputs=[
+            dataset_id_input,
+            start_idx_input,
+            display_count_input,
+            tags_front_input,
+            tags_back_input,
+            add_artist_tags_input
+        ],
+        outputs=[
+            captions_df_out,
+            previews_gallery_out,
+            info_textbox_out
+        ]
+    )
+
+demo.launch()