Create main.py

main.py

@@ -0,0 +1,667 @@
+import os
+import json
+import uuid
+import time
+import tempfile
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Tuple
+
+import streamlit as st
+import pandas as pd
+from PIL import Image
+from huggingface_hub import HfApi, hf_hub_download
+from huggingface_hub.utils import HfHubHTTPError
+
+# =========================
+# CONFIG
+# =========================
+
+st.set_page_config(
+    page_title="VizWiz Benchmark Arena",
+    page_icon="🏆",
+    layout="wide",
+    initial_sidebar_state="expanded",
+)
+
+DB_REPO_ID = os.getenv("DB_REPO_ID", "NidhiS09/VizWiz-submissions-db")
+DB_REPO_TYPE = "dataset"
+SUBMISSIONS_TOKEN = os.getenv("SUBMISSIONS_TOKEN", "")
+
+DAILY_SUBMISSION_CAP = 5
+
+PHASES = [
+    {"label": "Dev (query-dev2024)", "codename": "test-dev2024"},
+    {"label": "Standard (query-standard2024)", "codename": "test-standard2024"},
+    {"label": "Challenge (query-challenge2024)", "codename": "test-challenge2024"},
+]
+
+CHALLENGE_TYPES = ["Object Detection", "Instance Segmentation"]
+
+LEADERBOARD_METRICS = ["bbox_mAP", "bbox_AP50", "segm_mAP", "segm_AP50"]
+DEFAULT_SORT_METRIC = "segm_AP50"
+
+# =========================
+# AUTH HELPERS
+# =========================
+
+def get_oauth_info() -> Dict | None:
+    """
+    In Docker spaces, HF injects user info via the X-Space-User-Info request header
+    as a JSON string once the user has authenticated via /oauth/login.
+    """
+    try:
+        headers = st.context.headers
+        user_info_raw = headers.get("X-Space-User-Info")
+        if user_info_raw:
+            data = json.loads(user_info_raw)
+            return {
+                "preferred_username": data.get("preferred_username", data.get("name", "")),
+                "name": data.get("name", ""),
+                "email": data.get("email", ""),
+                "picture": data.get("picture", None),
+            }
+    except Exception:
+        pass
+    return None
+
+
+def get_login_url() -> str:
+    """
+    In HF Docker spaces, OAuth login is at /oauth/login on the space host.
+    HF injects SPACE_HOST as an env var automatically.
+    """
+    space_host = os.getenv("SPACE_HOST", "")
+    if space_host:
+        return f"https://{space_host}/oauth/login"
+    return "/oauth/login"
+
+
+def render_login_wall():
+    """Shows login prompt and stops rendering if not authenticated."""
+    user = get_oauth_info()
+    if user:
+        return user
+
+    st.markdown("---")
+    st.markdown("## 🔐 Login Required")
+    st.info(
+        "You need to log in with your Hugging Face account to submit predictions "
+        "or view your submission history."
+    )
+    login_url = get_login_url()
+    st.markdown(
+        f'''<a href="{login_url}" target="_self">
+        <button style="background:#FF6B35;color:white;border:none;padding:12px 28px;
+        font-size:16px;border-radius:8px;cursor:pointer;font-weight:bold;">
+        🤗 Login with Hugging Face</button></a>''',
+        unsafe_allow_html=True,
+    )
+    st.stop()
+
+
+# =========================
+# SUBMISSION CAP HELPERS
+# =========================
+
+def _today_utc_str() -> str:
+    return datetime.now(timezone.utc).strftime("%Y-%m-%d")
+
+
+def _count_submissions_today(hf_username: str) -> int:
+    """
+    Count how many submissions this user has made today by scanning submission metadata.
+    We do this by listing all submission files and checking meta.json for matching username + today's date.
+    Uses a lightweight approach: filter by filename listing only.
+    """
+    try:
+        files = api_client().list_repo_files(
+            repo_id=DB_REPO_ID, repo_type=DB_REPO_TYPE, token=SUBMISSIONS_TOKEN
+        )
+        today = _today_utc_str()
+        count = 0
+        for f in files:
+            if not (f.startswith("submissions/") and f.endswith("/meta.json")):
+                continue
+            try:
+                p = hf_hub_download(
+                    repo_id=DB_REPO_ID,
+                    repo_type=DB_REPO_TYPE,
+                    filename=f,
+                    token=SUBMISSIONS_TOKEN,
+                )
+                meta = json.load(open(p))
+                if meta.get("hf_username") == hf_username:
+                    ts = meta.get("timestamp", 0)
+                    sub_date = datetime.fromtimestamp(ts, tz=timezone.utc).strftime("%Y-%m-%d")
+                    if sub_date == today:
+                        count += 1
+            except Exception:
+                continue
+        return count
+    except Exception:
+        return 0
+
+
+# =========================
+# HF API HELPERS
+# =========================
+
+@st.cache_resource
+def api_client() -> HfApi:
+    return HfApi()
+
+
+def _require_token() -> None:
+    if not SUBMISSIONS_TOKEN:
+        st.error(
+            "Missing SUBMISSIONS_TOKEN. Add it in Space Settings → Secrets."
+        )
+        st.stop()
+
+
+def _validate_submission_json(obj: Any) -> Tuple[bool, str]:
+    if not isinstance(obj, list):
+        return False, "Submission must be a JSON list of annotations."
+
+    required_keys = {"image_id", "score", "category_id", "area", "bbox", "segmentation"}
+
+    for i, ann in enumerate(obj):
+        if not isinstance(ann, dict):
+            return False, f"Annotation at index {i} must be an object/dict."
+        missing = required_keys - set(ann.keys())
+        if missing:
+            return False, f"Annotation at index {i} missing keys: {sorted(list(missing))}"
+        if not isinstance(ann["image_id"], int):
+            return False, f"image_id at index {i} must be an integer."
+        if not isinstance(ann["category_id"], int):
+            return False, f"category_id at index {i} must be an integer."
+        if not isinstance(ann["score"], (int, float)):
+            return False, f"score at index {i} must be a number."
+        if not isinstance(ann["area"], (int, float)):
+            return False, f"area at index {i} must be a number."
+        bbox = ann["bbox"]
+        if not (isinstance(bbox, list) and len(bbox) == 4 and all(isinstance(x, (int, float)) for x in bbox)):
+            return False, f"bbox at index {i} must be a list of 4 numbers: [x, y, w, h]."
+        if not isinstance(ann["segmentation"], list):
+            return False, f"segmentation at index {i} must be a list."
+
+    return True, "OK"
+
+
+def _upload_json(data: Any, path_in_repo: str, commit_message: str = "") -> None:
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as tmp:
+        json.dump(data, tmp, ensure_ascii=False)
+        tmp_path = tmp.name
+    try:
+        api_client().upload_file(
+            path_or_fileobj=tmp_path,
+            path_in_repo=path_in_repo,
+            repo_id=DB_REPO_ID,
+            repo_type=DB_REPO_TYPE,
+            token=SUBMISSIONS_TOKEN,
+            commit_message=commit_message or f"Add {path_in_repo}",
+        )
+    finally:
+        try:
+            os.remove(tmp_path)
+        except OSError:
+            pass
+
+
+def _create_submission_record(
+    *,
+    pred: List[Dict[str, Any]],
+    team: str,
+    model_name: str,
+    phase_codename: str,
+    challenge_type: str,
+    original_filename: str,
+    hf_username: str,
+    hf_email: str,
+) -> str:
+    _require_token()
+
+    submission_id = str(uuid.uuid4())
+    ts = int(time.time())
+
+    meta = {
+        "submission_id": submission_id,
+        "team": team.strip(),
+        "model": model_name.strip(),
+        "phase_codename": phase_codename,
+        "challenge_type": challenge_type,
+        "timestamp": ts,
+        "original_filename": original_filename,
+        "hf_username": hf_username,
+        "hf_email": hf_email,
+    }
+
+    status = {"state": "queued", "timestamp": ts}
+    base = f"submissions/{submission_id}"
+
+    _upload_json(pred, f"{base}/pred.json", f"pred {submission_id}")
+    _upload_json(meta, f"{base}/meta.json", f"meta {submission_id}")
+    _upload_json(status, f"{base}/status.json", f"status {submission_id}")
+
+    return submission_id
+
+
+def _load_leaderboard_df() -> pd.DataFrame:
+    _require_token()
+    try:
+        path = hf_hub_download(
+            repo_id=DB_REPO_ID,
+            repo_type=DB_REPO_TYPE,
+            filename="leaderboard.jsonl",
+            token=SUBMISSIONS_TOKEN,
+        )
+    except HfHubHTTPError as e:
+        if "404" in str(e):
+            return pd.DataFrame(columns=["team", "model", "phase_codename", *LEADERBOARD_METRICS, "timestamp"])
+        raise
+
+    rows = []
+    with open(path, "r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                rows.append(json.loads(line))
+            except json.JSONDecodeError:
+                continue
+
+    if not rows:
+        return pd.DataFrame(columns=["team", "model", "phase_codename", *LEADERBOARD_METRICS, "timestamp"])
+
+    df = pd.DataFrame(rows)
+    for col in ["team", "model", "phase_codename", "timestamp", *LEADERBOARD_METRICS]:
+        if col not in df.columns:
+            df[col] = None
+
+    if DEFAULT_SORT_METRIC in df.columns:
+        df = df.sort_values(by=DEFAULT_SORT_METRIC, ascending=False, kind="mergesort")
+
+    return df
+
+
+def _load_user_submissions(hf_username: str) -> List[Dict]:
+    """Load all submissions for a specific HF user."""
+    _require_token()
+    try:
+        files = api_client().list_repo_files(
+            repo_id=DB_REPO_ID, repo_type=DB_REPO_TYPE, token=SUBMISSIONS_TOKEN
+        )
+    except Exception:
+        return []
+
+    results = []
+    for f in files:
+        if not (f.startswith("submissions/") and f.endswith("/meta.json")):
+            continue
+        try:
+            sid = f.split("/")[1]
+            meta_path = hf_hub_download(
+                repo_id=DB_REPO_ID, repo_type=DB_REPO_TYPE,
+                filename=f, token=SUBMISSIONS_TOKEN
+            )
+            meta = json.load(open(meta_path))
+
+            if meta.get("hf_username") != hf_username:
+                continue
+
+            # Get status
+            try:
+                status_path = hf_hub_download(
+                    repo_id=DB_REPO_ID, repo_type=DB_REPO_TYPE,
+                    filename=f"submissions/{sid}/status.json", token=SUBMISSIONS_TOKEN
+                )
+                status = json.load(open(status_path))
+            except Exception:
+                status = {"state": "unknown"}
+
+            # Get results if done
+            metrics = {}
+            if status.get("state") == "done":
+                metrics = status.get("metrics", {})
+
+            results.append({
+                "submission_id": sid,
+                "team": meta.get("team", ""),
+                "model": meta.get("model", ""),
+                "phase": meta.get("phase_codename", ""),
+                "challenge_type": meta.get("challenge_type", ""),
+                "timestamp": meta.get("timestamp", 0),
+                "state": status.get("state", "unknown"),
+                **metrics,
+            })
+        except Exception:
+            continue
+
+    results.sort(key=lambda x: x["timestamp"], reverse=True)
+    return results
+
+
+# =========================
+# UI PAGES
+# =========================
+
+def render_overview():
+    with st.expander("ℹ️ Overview of the VizWiz Benchmark Arena"):
+        st.markdown("""
+        **VizWiz Object Localization Challenge** — Submit your model predictions and get scored automatically.
+        
+        - 📤 **Submit**: Upload your prediction JSON for Dev, Standard, or Challenge phases
+        - 🏆 **Leaderboard**: View public rankings for the Standard phase
+        - 📋 **My Submissions**: Track all your past submissions and scores
+        """)
+        try:
+            overview_image = Image.open("src/overview_image.png").resize((600, 600))
+            st.image(overview_image, caption="Example of an object localization task")
+        except Exception:
+            pass
+
+
+def render_eval_details():
+    with st.expander("📐 How is the Score Calculated?"):
+        st.markdown("""
+        Your submission is evaluated automatically against hidden ground-truth annotations.
+        The leaderboard reports:
+        - **bbox_mAP** — bounding box mean average precision
+        - **bbox_AP50** — bounding box AP at IoU=0.50
+        - **segm_mAP** — segmentation mean average precision  
+        - **segm_AP50** — segmentation AP at IoU=0.50 *(default ranking)*
+        """)
+
+
+def page_leaderboard():
+    st.header("🏆 Leaderboard — Standard Phase")
+    st.caption(f"Ranked by **{DEFAULT_SORT_METRIC}** (descending). Showing Standard phase submissions only.")
+
+    with st.spinner("Loading leaderboard..."):
+        try:
+            df = _load_leaderboard_df()
+        except Exception as e:
+            st.error(f"Could not load leaderboard: {e}")
+            return
+
+    # Filter to standard phase only for public leaderboard
+    if not df.empty and "phase_codename" in df.columns:
+        df = df[df["phase_codename"] == "test-standard2024"]
+
+    if df.empty:
+        st.info("No scored Standard phase submissions yet. Be the first!")
+        return
+
+    df_display = df.copy()
+    df_display.insert(0, "Rank", range(1, len(df_display) + 1))
+    if "timestamp" in df_display.columns:
+        df_display["timestamp"] = pd.to_datetime(df_display["timestamp"], unit="s", errors="coerce")
+
+    # Drop hf_username / hf_email from public view
+    for col in ["hf_username", "hf_email", "phase_codename", "submission_id"]:
+        if col in df_display.columns:
+            df_display.drop(columns=[col], inplace=True)
+
+    st.dataframe(
+        df_display,
+        column_config={
+            "Rank": st.column_config.Column("Rank", width="small"),
+            "team": "Team",
+            "model": "Model",
+            "bbox_mAP": st.column_config.NumberColumn("bbox_mAP", format="%.4f"),
+            "bbox_AP50": st.column_config.NumberColumn("bbox_AP50", format="%.4f"),
+            "segm_mAP": st.column_config.NumberColumn("segm_mAP", format="%.4f"),
+            "segm_AP50": st.column_config.NumberColumn("segm_AP50", format="%.4f"),
+            "timestamp": st.column_config.DatetimeColumn("Scored at", format="D MMM YYYY, h:mm a"),
+        },
+        use_container_width=True,
+        hide_index=True,
+    )
+
+
+def page_submit():
+    user = render_login_wall()
+    hf_username = user.get("preferred_username", "unknown")
+    hf_email = user.get("email", "")
+
+    st.header("🚀 Submit Your Predictions")
+
+    # Show user info + daily cap status
+    col_info, col_cap = st.columns([3, 1])
+    with col_info:
+        avatar = user.get("picture", "")
+        name = user.get("name", hf_username)
+        if avatar:
+            st.markdown(
+                f'<img src="{avatar}" width="32" style="border-radius:50%;vertical-align:middle;margin-right:8px">'
+                f'<b>{name}</b> (@{hf_username})',
+                unsafe_allow_html=True,
+            )
+        else:
+            st.markdown(f"👤 Logged in as **{name}** (@{hf_username})")
+
+    with col_cap:
+        with st.spinner("Checking daily limit..."):
+            subs_today = _count_submissions_today(hf_username)
+        remaining = DAILY_SUBMISSION_CAP - subs_today
+        color = "green" if remaining > 1 else ("orange" if remaining == 1 else "red")
+        st.markdown(
+            f'<div style="text-align:center;padding:8px;border-radius:8px;border:1px solid {color}">'
+            f'<b style="color:{color}">{remaining}/{DAILY_SUBMISSION_CAP}</b><br>'
+            f'<small>submissions left today</small></div>',
+            unsafe_allow_html=True,
+        )
+
+    st.markdown("---")
+
+    if remaining <= 0:
+        st.error(f"⛔ You've reached your daily limit of {DAILY_SUBMISSION_CAP} submissions. Come back tomorrow!")
+        return
+
+    col1, col2 = st.columns([2, 1])
+
+    with col2:
+        st.subheader("Submission Info")
+        team = st.text_input("Team / Display Name", value=st.session_state.get("team", ""))
+        model_name = st.text_input("Model Name", value=st.session_state.get("model_name", ""))
+        phase_label = st.selectbox("Phase", [p["label"] for p in PHASES])
+        phase_codename = next(p["codename"] for p in PHASES if p["label"] == phase_label)
+        challenge_type = st.radio("Challenge type", CHALLENGE_TYPES, horizontal=False)
+        st.session_state["team"] = team
+        st.session_state["model_name"] = model_name
+        st.caption("Submissions are scored automatically. Results appear in My Submissions once processed.")
+
+    with col1:
+        st.subheader("Upload Submission File")
+        uploaded_file = st.file_uploader("Choose a JSON file", type=["json"])
+
+        if uploaded_file is None:
+            st.info("Upload a JSON file containing a list of annotations.")
+            return
+
+        try:
+            raw = uploaded_file.getvalue().decode("utf-8")
+            pred_obj = json.loads(raw)
+        except Exception:
+            st.error("Could not parse JSON. Please upload a valid JSON file.")
+            return
+
+        ok, msg = _validate_submission_json(pred_obj)
+        if not ok:
+            st.error(f"Invalid submission format: {msg}")
+            return
+
+        st.success("Submission file looks valid ✅")
+
+        submit_clicked = st.button("Submit (Queue for Evaluation)", type="primary")
+
+        if submit_clicked:
+            if not team.strip():
+                st.error("Please enter a Team / Display Name.")
+                return
+            if not model_name.strip():
+                st.error("Please enter a Model Name.")
+                return
+
+            # Re-check cap right before submitting
+            subs_today_fresh = _count_submissions_today(hf_username)
+            if subs_today_fresh >= DAILY_SUBMISSION_CAP:
+                st.error("⛔ Daily submission limit reached. Try again tomorrow.")
+                return
+
+            with st.spinner("Uploading submission..."):
+                try:
+                    submission_id = _create_submission_record(
+                        pred=pred_obj,
+                        team=team,
+                        model_name=model_name,
+                        phase_codename=phase_codename,
+                        challenge_type=challenge_type,
+                        original_filename=uploaded_file.name,
+                        hf_username=hf_username,
+                        hf_email=hf_email,
+                    )
+                except Exception as e:
+                    st.error(f"Upload failed: {e}")
+                    return
+
+            st.balloons()
+            st.success("Submission queued successfully!")
+            st.code(f"Submission ID: {submission_id}")
+            st.info("Track your results in **My Submissions** once the evaluator processes your file.")
+
+
+def page_my_submissions():
+    user = render_login_wall()
+    hf_username = user.get("preferred_username", "unknown")
+
+    st.header("📋 My Submissions")
+    st.caption(f"Showing all submissions for @{hf_username}")
+
+    with st.spinner("Loading your submissions..."):
+        submissions = _load_user_submissions(hf_username)
+
+    if not submissions:
+        st.info("You haven't made any submissions yet. Head to **Submit Model** to get started!")
+        return
+
+    # Phase filter
+    phases_present = sorted(set(s["phase"] for s in submissions))
+    selected_phase = st.selectbox("Filter by phase", ["All"] + phases_present)
+
+    filtered = submissions if selected_phase == "All" else [s for s in submissions if s["phase"] == selected_phase]
+
+    # State badges
+    state_colors = {
+        "queued": "🟡",
+        "running": "🔵",
+        "done": "🟢",
+        "failed": "🔴",
+        "unknown": "⚪",
+    }
+
+    df = pd.DataFrame(filtered)
+
+    if "timestamp" in df.columns:
+        df["submitted_at"] = pd.to_datetime(df["timestamp"], unit="s", errors="coerce")
+
+    if "state" in df.columns:
+        df["status"] = df["state"].apply(lambda s: f"{state_colors.get(s, '⚪')} {s.capitalize()}")
+
+    display_cols = ["submitted_at", "status", "team", "model", "phase", "challenge_type"]
+    metric_cols = [m for m in LEADERBOARD_METRICS if m in df.columns]
+    display_cols += metric_cols
+
+    df_display = df[[c for c in display_cols if c in df.columns]].copy()
+
+    col_config = {
+        "submitted_at": st.column_config.DatetimeColumn("Submitted At", format="D MMM YYYY, h:mm a"),
+        "status": "Status",
+        "team": "Team",
+        "model": "Model",
+        "phase": "Phase",
+        "challenge_type": "Challenge Type",
+    }
+    for m in metric_cols:
+        col_config[m] = st.column_config.NumberColumn(m, format="%.4f")
+
+    st.dataframe(df_display, column_config=col_config, use_container_width=True, hide_index=True)
+
+    # Summary stats
+    st.markdown("---")
+    total = len(submissions)
+    done = sum(1 for s in submissions if s["state"] == "done")
+    today_count = sum(
+        1 for s in submissions
+        if datetime.fromtimestamp(s["timestamp"], tz=timezone.utc).strftime("%Y-%m-%d") == _today_utc_str()
+    )
+
+    c1, c2, c3 = st.columns(3)
+    c1.metric("Total Submissions", total)
+    c2.metric("Scored", done)
+    c3.metric(f"Today ({DAILY_SUBMISSION_CAP} max)", f"{today_count}/{DAILY_SUBMISSION_CAP}")
+
+
+# =========================
+# MAIN
+# =========================
+
+def main():
+    st.sidebar.title("VizWiz Benchmark 🏆")
+
+    # Check login state for sidebar display
+    user = get_oauth_info()
+
+    if user:
+        hf_username = user.get("preferred_username", "")
+        avatar = user.get("picture", "")
+        if avatar:
+            st.sidebar.markdown(
+                f'<img src="{avatar}" width="40" style="border-radius:50%"><br>'
+                f'<b>@{hf_username}</b>',
+                unsafe_allow_html=True,
+            )
+        else:
+            st.sidebar.markdown(f"👤 @{hf_username}")
+        space_host = os.getenv("SPACE_HOST", "")
+        logout_url = f"https://{space_host}/oauth/logout" if space_host else "/oauth/logout"
+        st.sidebar.markdown(
+            f'''<a href="{logout_url}" target="_self">
+            <button style="width:100%;background:#555;color:white;border:none;
+            padding:8px;border-radius:6px;cursor:pointer;font-size:13px;">
+            Logout</button></a>''',
+            unsafe_allow_html=True,
+        )
+        st.sidebar.markdown("---")
+        menu = ["Leaderboard", "Submit Model", "My Submissions"]
+    else:
+        st.sidebar.info("Login to submit predictions and view your history.")
+        login_url = get_login_url()
+        st.sidebar.markdown(
+            f'''<a href="{login_url}" target="_self">
+            <button style="width:100%;background:#FF6B35;color:white;border:none;
+            padding:10px;border-radius:6px;cursor:pointer;font-size:14px;font-weight:bold;">
+            🤗 Login with HF</button></a>''',
+            unsafe_allow_html=True,
+        )
+        menu = ["Leaderboard", "Submit Model", "My Submissions"]
+
+    choice = st.sidebar.radio("Navigation", menu)
+    st.sidebar.markdown("---")
+    st.sidebar.caption("Submissions are queued to a private DB repo and evaluated automatically.")
+
+    render_overview()
+    render_eval_details()
+    st.markdown("---")
+
+    if choice == "Leaderboard":
+        page_leaderboard()
+    elif choice == "Submit Model":
+        page_submit()
+    elif choice == "My Submissions":
+        page_my_submissions()
+
+
+if __name__ == "__main__":
+    main()