src/streamlit_app.py

# app.py
import json
import re
from pathlib import Path
from urllib.parse import urlparse
import streamlit as st

# ──────────────────────────────────────────────────────────────────────────────
# Page config
# ──────────────────────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="ImagenWorld Data Visualization",
    page_icon="🖼️",
    layout="wide",
    initial_sidebar_state="expanded"
)

# ──────────────────────────────────────────────────────────────────────────────
# Basics
# ──────────────────────────────────────────────────────────────────────────────
MANIFEST_PATH = Path(__file__).parent / "manifest_v1.json"

DATA_DIRS = ["TIG", "TIE", "SRIG", "SRIE", "MRIG", "MRIE"]
TASK_DESCRIPTIONS = {
    "TIG":  "Text to Image Generation - Generate images from text prompts only",
    "TIE":  "Text and Image to Image Editing - Edit images based on text prompts",
    "SRIG": "Single-Reference Image Generation - Generate images using text + single reference image",
    "SRIE": "Single-Reference Image Editing - Edit images using text + single reference image",
    "MRIG": "Multi-Reference Image Generation - Generate images using text + multiple reference images",
    "MRIE": "Multi-Reference Image Editing - Edit images using text + multiple reference images"
}

DEFAULT_CHUNK = 10
FIXED_HEIGHT_PX = 320  # always stabilized

if "show_counts" not in st.session_state:
    st.session_state.show_counts = {task: DEFAULT_CHUNK for task in DATA_DIRS}

# ──────────────────────────────────────────────────────────────────────────────
# Helpers: natural sorting for ids and filenames
# ──────────────────────────────────────────────────────────────────────────────
_num_re = re.compile(r"(\d+)")

def _natural_key(s: str):
    parts = _num_re.split(s)
    return [int(p) if p.isdigit() else p.lower() for p in parts]

def _basename_from_url(url: str) -> str:
    try:
        return Path(urlparse(url).path).name
    except Exception:
        return url

def _sorted_urls(urls):
    return sorted(urls, key=lambda u: _natural_key(_basename_from_url(u)))

# ──────────────────────────────────────────────────────────────────────────────
# Load manifest
# ──────────────────────────────────────────────────────────────────────────────
@st.cache_data(show_spinner=False)
def load_manifest():
    with open(MANIFEST_PATH, "r", encoding="utf-8") as f:
        man = json.load(f)
    items = man.get("items", [])
    per_task = {t: [] for t in DATA_DIRS}
    topics = set()

    for it in items:
        if it.get("cond_image_urls"):
            it["cond_image_urls"] = _sorted_urls(it["cond_image_urls"])
        if it.get("model_output_urls"):
            it["model_output_urls"] = _sorted_urls(it["model_output_urls"])

        per_task.setdefault(it.get("task", "Unknown"), []).append(it)
        if it.get("topic"):
            topics.add(it["topic"])

    for t, lst in per_task.items():
        lst.sort(key=lambda it: _natural_key(str(it.get("item_id", ""))))
    return per_task, sorted(list(topics))

# ──────────────────────────────────────────────────────────────────────────────
# Stable image grid
# ──────────────────────────────────────────────────────────────────────────────
def _display_images(urls, caption_prefix="", max_per_row=3):
    if not urls:
        st.write("No images found.")
        return

    cols = st.columns(max_per_row, vertical_alignment="top")
    for i, url in enumerate(urls):
        col = cols[i % max_per_row]
        with col:
            st.markdown(
                f"""
                <div class="img-frame" style="height:{FIXED_HEIGHT_PX}px; display:flex; align-items:center; justify-content:center; overflow:hidden; border-radius:12px;">
                  <img src="{url}" alt="{_basename_from_url(url)}" style="max-height:100%; width:100%; object-fit:contain;" />
                </div>
                <div class="img-cap" style="text-align:center; font-size:0.9rem; opacity:0.8; margin-top:6px;">
                  {caption_prefix} {_basename_from_url(url)}
                </div>
                """,
                unsafe_allow_html=True,
            )

# ──────────────────────────────────────────────────────────────────────────────
# Global CSS (adds top padding for HF Space header)
# ──────────────────────────────────────────────────────────────────────────────
def _inject_css():
    css = """
    <style>
      /* Offset for HF Spaces top bar (~60–70px) so the main title is visible */
      .block-container { padding-top: 5rem; }
      h1:first-of-type { margin-top: 0; }

      [data-testid="stImage"] img {
        width: 100%;
        height: auto;
        object-fit: contain;
        display: block;
      }
      html * {
        -webkit-font-smoothing: antialiased;
        -moz-osx-font-smoothing: grayscale;
        text-rendering: optimizeLegibility;
      }
      .img-frame { background: rgba(0,0,0,0.03); }
    </style>
    """
    st.markdown(css, unsafe_allow_html=True)

# ──────────────────────────────────────────────────────────────────────────────
# App
# ──────────────────────────────────────────────────────────────────────────────
def main():
    st.title("🖼️ ImagenWorld Data Visualization")
    st.markdown("Each task starts with **10** items — click **Show more** to load **+10**.")
    _inject_css()

    with st.spinner("Loading manifest…"):
        per_task, _all_topics = load_manifest()

    # Sidebar filters
    st.sidebar.header("Filters")
    selected_tasks = st.sidebar.multiselect("Select Tasks", DATA_DIRS, default=DATA_DIRS)
    search_query = st.sidebar.text_input("🔍 Search in prompts", "")
    topic_filter = st.sidebar.multiselect("Select Topics", _all_topics, default=[])
    subtopic_filter = st.sidebar.text_input("Filter by subtopic (optional)", "")

    st.sidebar.header("Task Descriptions")
    for t in selected_tasks:
        st.sidebar.write(f"**{t}**: {TASK_DESCRIPTIONS.get(t, '')}")

    # Tabs per task
    tabs = st.tabs(selected_tasks) if selected_tasks else []
    for task, tab in zip(selected_tasks, tabs):
        with tab:
            st.subheader(task)
            limit = st.session_state.show_counts.get(task, DEFAULT_CHUNK)
            all_items = per_task.get(task, [])

            def _match(it):
                sq = search_query.strip().lower()
                if sq and (sq not in it.get("prompt", "").lower()
                           and sq not in it.get("prompt_refined", "").lower()):
                    return False
                if topic_filter and it.get("topic", "") not in topic_filter:
                    return False
                if subtopic_filter and it.get("subtopic", "") != subtopic_filter:
                    return False
                return True

            filtered = [it for it in all_items if _match(it)]
            batch = filtered[:limit]

            st.caption(f"Showing {len(batch)} / {len(filtered)} (from {len(all_items)} items in {task})")

            if not batch:
                st.warning("No items match current filters.")
            else:
                for it in batch:
                    header = f"**{it.get('item_id','?')}** — {it.get('topic','Unknown')} / {it.get('subtopic','Unknown')}"
                    with st.expander(header, expanded=False):
                        c1, c2, c3 = st.columns([1,1,1], vertical_alignment="top")
                        with c1: st.write(f"**Task:** {it.get('task','Unknown')}")
                        with c2: st.write(f"**Topic:** {it.get('topic','Unknown')}")
                        with c3: st.write(f"**Subtopic:** {it.get('subtopic','Unknown')}")
                        st.write("**Original Prompt:**")
                        st.write(it.get("prompt", "—"))
                        if it.get("prompt_refined"):
                            st.write("**Refined Prompt:**")
                            st.write(it.get("prompt_refined", ""))
                        if it.get("remarks"):
                            st.write("**Remarks:**")
                            st.write(it.get("remarks", ""))

                        cond_urls  = it.get("cond_image_urls", [])
                        model_urls = it.get("model_output_urls", [])

                        if cond_urls:
                            st.write("**Condition Images:**")
                            _display_images(cond_urls, "Condition")
                        if model_urls:
                            st.write("**Model Output:**")
                            _display_images(model_urls, "Model")
                st.divider()

            # Pagination controls
            c_more, c_reset, c_info = st.columns(3)
            with c_more:
                if st.button("Show more", key=f"more_{task}"):
                    st.session_state.show_counts[task] = limit + DEFAULT_CHUNK
                    st.rerun()
            with c_reset:
                if st.button("Reset", key=f"reset_{task}"):
                    st.session_state.show_counts[task] = DEFAULT_CHUNK
                    st.rerun()
            with c_info:
                st.caption(f"Current limit: {limit}")

if __name__ == "__main__":
    main()