app.py

#!/usr/bin/env python3
"""HF Space entry point for 3DReflecNet dataset preview.

Loads only data/preview/preview.parquet so the Space exposes the configured
preview instance subset instead of the full dataset metadata.
"""
from __future__ import annotations

import atexit
import io
import os
import shutil
import tempfile
from pathlib import Path
from typing import Any

import gradio as gr
import pandas as pd
from datasets import load_dataset
from huggingface_hub import hf_hub_download
from PIL import Image

from utils import (
    BOOL_FILTER_CHOICES,
    FILTER_ALL,
    filter_dataframe_advanced,
    get_distinct_text_choices,
    logger,
    require_bool_columns,
    require_columns,
    require_text_columns,
    setup_logging,
)

DATASET_REPO = os.environ.get("DATASET_REPO", "3DReflecNet/3DReflecNet")
HF_TOKEN = os.environ.get("HF_TOKEN", None)
MAX_RESULTS = 300
BOOL_COLUMNS = ["hasGlass", "isGenerated", "transparent", "near_light"]
TAG_TEXT_COLUMNS = [
    "indoor_outdoor", "glossiness", "reflectivity", "texture_scale",
    "anisotropy", "metallic_hint", "light_type", "light_intensity",
    "key_light_direction", "shadow_hardness", "description",
]
TAG_NUMERIC_COLUMNS = ["roughness"]
TAG_COLUMNS = TAG_TEXT_COLUMNS + TAG_NUMERIC_COLUMNS
TEXT_FILTER_COLUMNS = ["main_category", "material_name", "env_name"]

_GLB_CACHE_DIR = Path(tempfile.mkdtemp(prefix="glb_cache_"))
atexit.register(shutil.rmtree, str(_GLB_CACHE_DIR), True)


# ---------------------------------------------------------------------------
# Data loading
# ---------------------------------------------------------------------------

def load_preview_dataframe() -> pd.DataFrame:
    """Load the small preview Parquet into memory."""
    PREVIEW_COLS = [
        "instance_id", "split", "frame_id", "rgb", "mask",
        "depth_preview", "normal_preview",
        "main_category", "sub_category", "model_name",
        "material_name", "env_name", "glb_path",
        "hasGlass", "isGenerated", "transparent", "near_light",
        "indoor_outdoor", "glossiness", "roughness", "reflectivity",
        "texture_scale", "anisotropy", "metallic_hint", "light_type",
        "light_intensity", "key_light_direction", "shadow_hardness",
        "description",
    ]
    ds = load_dataset(
        DATASET_REPO,
        data_files="data/preview/preview.parquet",
        split="train",
        streaming=False,
        token=HF_TOKEN,
    ).select_columns(PREVIEW_COLS)
    df = pd.DataFrame(list(ds))
    require_columns(df, PREVIEW_COLS, "preview parquet")
    require_text_columns(
        df,
        [
            "instance_id", "split", "main_category", "sub_category",
            "model_name", "material_name", "env_name", "glb_path",
            *TAG_TEXT_COLUMNS,
        ],
        "preview parquet",
    )
    require_bool_columns(df, BOOL_COLUMNS, "preview parquet")
    if df["roughness"].isna().any() or not pd.api.types.is_numeric_dtype(df["roughness"]):
        raise TypeError(f"Expected numeric dtype for column 'roughness' in preview parquet, got {df['roughness'].dtype}.")
    if df["frame_id"].isna().any() or not pd.api.types.is_integer_dtype(df["frame_id"]):
        raise TypeError(f"Expected integer dtype for column 'frame_id' in preview parquet, got {df['frame_id'].dtype}.")
    for col in ["rgb", "mask", "depth_preview", "normal_preview"]:
        invalid = df[col].map(lambda value: not isinstance(value, (bytes, bytearray)) or len(value) == 0)
        if invalid.any():
            raise TypeError(f"Expected non-empty binary values for column {col!r} in preview parquet.")
    return df


def decode_image_bytes(img_bytes: bytes | bytearray, context: str) -> Image.Image:
    if not isinstance(img_bytes, (bytes, bytearray)) or not img_bytes:
        raise TypeError(f"Expected non-empty image bytes for {context}.")
    with Image.open(io.BytesIO(img_bytes)) as img:
        return img.copy()


def build_preview_instance_dataframe(preview_df: pd.DataFrame) -> pd.DataFrame:
    """Derive one row per preview instance from preview frame rows."""
    instance_cols = [
        "instance_id", "main_category", "sub_category", "model_name",
        "material_name", "env_name", "hasGlass", "isGenerated",
        "transparent", "near_light", "glb_path", *TAG_COLUMNS,
    ]
    require_columns(preview_df, instance_cols, "preview parquet")

    rows: list[dict[str, Any]] = []
    for instance_id, group in preview_df.groupby("instance_id", sort=True):
        row: dict[str, Any] = {}
        for col in instance_cols:
            values = group[col].drop_duplicates().tolist()
            if len(values) != 1:
                raise ValueError(f"Inconsistent {col!r} values for preview instance {instance_id!r}.")
            row[col] = values[0]
        rows.append(row)

    df = pd.DataFrame(rows, columns=instance_cols)
    require_text_columns(
        df,
        [
            "instance_id", "main_category", "sub_category",
            "model_name", "material_name", "env_name", "glb_path",
            *TAG_TEXT_COLUMNS,
        ],
        "preview instance dataframe",
    )
    require_bool_columns(df, BOOL_COLUMNS, "preview instance dataframe")
    if df["roughness"].isna().any() or not pd.api.types.is_numeric_dtype(df["roughness"]):
        raise TypeError(
            f"Expected numeric dtype for column 'roughness' in preview instance dataframe, got {df['roughness'].dtype}."
        )
    if df["glb_path"].map(lambda value: not value.strip()).any():
        raise ValueError("Preview instance dataframe contains empty GLB paths.")
    return df


def train_frame_rows(preview_df: pd.DataFrame, instance_id: str, max_frames: int | None = None) -> list[dict[str, Any]]:
    selected = preview_df[
        (preview_df["instance_id"].astype(str) == str(instance_id))
        & (preview_df["split"].astype(str) == "train")
    ].copy()
    if selected.empty:
        raise ValueError(f"Preview instance {instance_id!r} has no train split rows.")
    selected = selected.sort_values("frame_id")
    if max_frames is not None:
        selected = selected.head(max_frames)
    return selected.to_dict(orient="records")


def get_instance_thumbnail(preview_df: pd.DataFrame, instance_id: str) -> Image.Image:
    row = train_frame_rows(preview_df, instance_id, max_frames=1)[0]
    return decode_image_bytes(row["rgb"], f"{instance_id} thumbnail RGB")


def instance_caption(row: dict[str, Any]) -> str:
    return f"{row['model_name']} | {row['material_name']} | {row['env_name']}"


def build_llm_tags(row: dict[str, Any]) -> dict[str, Any]:
    """Reconstruct the tags.json shape from preview parquet columns."""
    return {
        "instance_id": row["instance_id"],
        "category": {
            "main": row["main_category"],
            "sub": row["sub_category"],
        },
        "material_properties": {
            "glossiness": row["glossiness"],
            "roughness": float(row["roughness"]),
            "reflectivity": row["reflectivity"],
            "texture_scale": row["texture_scale"],
            "anisotropy": row["anisotropy"],
            "metallic_hint": row["metallic_hint"],
        },
        "environment": {
            "indoor_outdoor": row["indoor_outdoor"],
            "light_type": row["light_type"],
            "light_intensity": row["light_intensity"],
            "key_light_direction": row["key_light_direction"],
            "shadow_hardness": row["shadow_hardness"],
        },
        "description": row["description"],
    }


def build_instance_metadata(row: dict[str, Any]) -> dict[str, Any]:
    return {key: value for key, value in row.items() if key not in TAG_COLUMNS}


def build_instance_gallery_items(
    rows: list[dict[str, Any]],
    preview_df: pd.DataFrame,
) -> list[tuple[Image.Image, str]]:
    return [
        (get_instance_thumbnail(preview_df, row["instance_id"]), instance_caption(row))
        for row in rows
    ]


def load_instance_frames(
    preview_df: pd.DataFrame, instance_id: str, max_frames: int = 50,
) -> list[dict[str, Any]]:
    """Load train preview image payloads for one instance from preview Parquet."""
    rows = train_frame_rows(preview_df, instance_id, max_frames=max_frames)
    frames: list[dict[str, Any]] = []
    for example in rows:
        frame_id = int(example["frame_id"])
        frame_item: dict[str, Any] = {"frame_id": frame_id}
        for key in ("rgb", "mask", "depth_preview", "normal_preview"):
            frame_item[key] = decode_image_bytes(example[key], f"{key} frame {frame_id}")
        frames.append(frame_item)
    return frames


def render_frame_images(frame_items: list[dict[str, Any]], frame_index: float) -> list[Any | None]:
    """Render RGB/Mask/Depth/Normal images for one selected frame index (1-based)."""
    if not frame_items:
        return [
            gr.update(value=None, label="RGB"),
            gr.update(value=None, label="Mask"),
            gr.update(value=None, label="Depth"),
            gr.update(value=None, label="Normal"),
        ]

    idx = int(round(frame_index)) - 1
    idx = max(0, min(idx, len(frame_items) - 1))
    selected = frame_items[idx]
    frame_id = int(selected["frame_id"])
    return [
        gr.update(value=selected["rgb"], label=f"RGB frame_{frame_id:05d}"),
        gr.update(value=selected["mask"], label=f"Mask frame_{frame_id:05d}"),
        gr.update(value=selected["depth_preview"], label=f"Depth frame_{frame_id:05d}"),
        gr.update(value=selected["normal_preview"], label=f"Normal frame_{frame_id:05d}"),
    ]


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def download_glb(glb_path: str) -> str:
    """Download pre-converted GLB file from HF dataset repo."""
    if not glb_path:
        raise ValueError("GLB path is required.")
    local = _GLB_CACHE_DIR / Path(glb_path).name
    if local.exists():
        return str(local)
    downloaded = hf_hub_download(
        repo_id=DATASET_REPO,
        filename=glb_path,
        repo_type="dataset",
        token=HF_TOKEN,
    )
    shutil.copy2(downloaded, str(local))
    logger.info("GLB ready: %s", local)
    return str(local)


# ---------------------------------------------------------------------------
# App builder
# ---------------------------------------------------------------------------

def build_app(instance_df: pd.DataFrame, preview_df: pd.DataFrame) -> gr.Blocks:
    main_category_choices = get_distinct_text_choices(instance_df, "main_category")
    material_name_choices = get_distinct_text_choices(instance_df, "material_name")
    env_name_choices = get_distinct_text_choices(instance_df, "env_name")

    def text_choices_for(df: pd.DataFrame, column: str) -> list[str]:
        if df.empty:
            return [FILTER_ALL]
        return get_distinct_text_choices(df, column)

    def filter_with_values(values: dict[str, str], skip_text_column: str | None = None) -> pd.DataFrame:
        return filter_dataframe_advanced(
            instance_df,
            main_category=FILTER_ALL if skip_text_column == "main_category" else values["main_category"],
            material_name=FILTER_ALL if skip_text_column == "material_name" else values["material_name"],
            env_name=FILTER_ALL if skip_text_column == "env_name" else values["env_name"],
            has_glass=values["has_glass"],
            is_generated=values["is_generated"],
            transparent=values["transparent"],
            near_light=values["near_light"],
        )

    def normalize_filter_values(values: dict[str, str], changed_column: str | None) -> dict[str, str]:
        normalized = values.copy()
        for column in TEXT_FILTER_COLUMNS:
            value = (normalized[column] or "").strip()
            normalized[column] = value if value else FILTER_ALL
        for column in ["has_glass", "is_generated", "transparent", "near_light"]:
            value = (normalized[column] or "").strip()
            normalized[column] = value if value else FILTER_ALL

        protected = {changed_column} if changed_column in TEXT_FILTER_COLUMNS else set()
        for _ in TEXT_FILTER_COLUMNS:
            changed = False
            for column in TEXT_FILTER_COLUMNS:
                if column in protected:
                    continue
                choices = text_choices_for(filter_with_values(normalized, skip_text_column=column), column)
                if normalized[column] not in choices:
                    normalized[column] = FILTER_ALL
                    changed = True
            if not changed:
                break

        for column in protected:
            choices = text_choices_for(filter_with_values(normalized, skip_text_column=column), column)
            if normalized[column] not in choices:
                normalized[column] = FILTER_ALL
        return normalized

    def text_filter_updates(values: dict[str, str]) -> list[Any]:
        return [
            gr.update(
                choices=text_choices_for(filter_with_values(values, skip_text_column=column), column),
                value=values[column],
            )
            for column in TEXT_FILTER_COLUMNS
        ]

    def filtered_instance_rows(values: dict[str, str]) -> tuple[pd.DataFrame, list[dict[str, Any]]]:
        filtered = filter_with_values(values)
        shown = filtered.head(MAX_RESULTS).copy()
        rows = shown.to_dict(orient="records")
        return filtered, rows

    def filter_gallery(
        main_category: str,
        material_name: str,
        env_name: str,
        has_glass: str,
        is_generated: str,
        transparent: str,
        near_light: str,
        changed_column: str | None = None,
    ):
        values = normalize_filter_values(
            {
                "main_category": main_category,
                "material_name": material_name,
                "env_name": env_name,
                "has_glass": has_glass,
                "is_generated": is_generated,
                "transparent": transparent,
                "near_light": near_light,
            },
            changed_column=changed_column,
        )
        filtered, rows = filtered_instance_rows(values)
        summary = f"Matched **{len(filtered)}** preview instances, showing **{len(rows)}**."
        gallery_items = build_instance_gallery_items(rows, preview_df)
        slider_empty = gr.update(minimum=1, maximum=1, step=1, value=1, interactive=False)
        return [
            *text_filter_updates(values),
            summary,
            gallery_items,
            rows,
            {},
            {},
            None,
            None,
            None,
            None,
            None,
            slider_empty,
            [],
        ]

    def make_filter_handler(changed_column: str):
        def handler(
            main_category: str,
            material_name: str,
            env_name: str,
            has_glass: str,
            is_generated: str,
            transparent: str,
            near_light: str,
        ):
            return filter_gallery(
                main_category,
                material_name,
                env_name,
                has_glass,
                is_generated,
                transparent,
                near_light,
                changed_column=changed_column,
            )

        return handler

    def on_instance_select(rows: list[dict[str, Any]], evt: gr.SelectData):
        if not rows:
            slider_empty = gr.update(minimum=1, maximum=1, step=1, value=1, interactive=False)
            return {}, {}, None, None, None, None, None, slider_empty, []
        idx = evt.index[0] if isinstance(evt.index, tuple) else evt.index
        if not isinstance(idx, int) or idx < 0 or idx >= len(rows):
            raise IndexError(f"Selected gallery index is out of range: {evt.index!r}")

        row = rows[idx]
        instance_id = row["instance_id"]
        if not isinstance(instance_id, str) or not instance_id.strip():
            raise ValueError(f"Selected instance row has invalid instance_id: {rows[idx]!r}")
        logger.info("Loading images for instance: %s", instance_id)
        frame_items = load_instance_frames(preview_df, instance_id, max_frames=50)
        slider_ready = gr.update(
            minimum=1,
            maximum=len(frame_items),
            step=1,
            value=1,
            interactive=True,
        )
        return (
            build_instance_metadata(row),
            build_llm_tags(row),
            download_glb(row["glb_path"]),
            *render_frame_images(frame_items, 1),
            slider_ready,
            frame_items,
        )

    def on_frame_change(frame_idx: float, frame_items: list[dict[str, Any]]):
        return render_frame_images(frame_items, frame_idx)

    initial_rows = instance_df.head(MAX_RESULTS).to_dict(orient="records")
    initial_gallery = build_instance_gallery_items(initial_rows, preview_df)
    initial_summary = f"Matched **{len(instance_df)}** preview instances, showing **{len(initial_rows)}**."

    with gr.Blocks(title="3DReflecNet Dataset Explorer") as demo:
        gr.Markdown("# 3DReflecNet Dataset Explorer")
        gr.Markdown(
            "Browse the configured preview subset. Select an RGB thumbnail to inspect the instance."
        )

        with gr.Row():
            main_category = gr.Dropdown(label="main_category", choices=main_category_choices, value=FILTER_ALL)
            material_name = gr.Dropdown(label="material_name", choices=material_name_choices, value=FILTER_ALL)
            env_name = gr.Dropdown(label="env_name", choices=env_name_choices, value=FILTER_ALL)
        with gr.Row():
            has_glass = gr.Dropdown(label="hasGlass", choices=BOOL_FILTER_CHOICES, value=FILTER_ALL)
            is_generated = gr.Dropdown(label="isGenerated", choices=BOOL_FILTER_CHOICES, value=FILTER_ALL)
            transparent = gr.Dropdown(label="transparent", choices=BOOL_FILTER_CHOICES, value=FILTER_ALL)
            near_light = gr.Dropdown(label="near_light", choices=BOOL_FILTER_CHOICES, value=FILTER_ALL)

        summary = gr.Markdown(initial_summary)
        instance_gallery = gr.Gallery(
            label="Preview Instances",
            value=initial_gallery,
            columns=5,
            object_fit="contain",
            height="auto",
        )

        with gr.Row():
            instance_meta = gr.JSON(label="Instance Metadata", height=280, scale=1)
            llm_tags = gr.JSON(label="LLM Tags", height=280, scale=1, open=True)

        model_viewer = gr.Model3D(
            label="3D Preview (GLB)",
            clear_color=(0.35, 0.35, 0.38, 1.0),
            camera_position=(35, 70, 3.5),
            height=460,
        )

        with gr.Row():
            rgb_image = gr.Image(label="RGB", height=320, interactive=False, scale=1, min_width=160)
            mask_image = gr.Image(label="Mask", height=320, interactive=False, scale=1, min_width=160)
            depth_image = gr.Image(label="Depth", height=320, interactive=False, scale=1, min_width=160)
            normal_image = gr.Image(label="Normal", height=320, interactive=False, scale=1, min_width=160)
        frame_slider = gr.Slider(
            label="Frame",
            minimum=1,
            maximum=1,
            step=1,
            value=1,
            interactive=False,
        )

        instance_state = gr.State(initial_rows)
        frame_state = gr.State([])

        filter_inputs = [
            main_category,
            material_name,
            env_name,
            has_glass,
            is_generated,
            transparent,
            near_light,
        ]
        filter_outputs = [
            main_category,
            material_name,
            env_name,
            summary,
            instance_gallery,
            instance_state,
            instance_meta,
            llm_tags,
            model_viewer,
            rgb_image,
            mask_image,
            depth_image,
            normal_image,
            frame_slider,
            frame_state,
        ]
        for filter_component, changed_column in [
            (main_category, "main_category"),
            (material_name, "material_name"),
            (env_name, "env_name"),
            (has_glass, "has_glass"),
            (is_generated, "is_generated"),
            (transparent, "transparent"),
            (near_light, "near_light"),
        ]:
            filter_component.change(
                fn=make_filter_handler(changed_column),
                inputs=filter_inputs,
                outputs=filter_outputs,
            )
        instance_gallery.select(
            fn=on_instance_select,
            inputs=[instance_state],
            outputs=[
                instance_meta,
                llm_tags,
                model_viewer,
                rgb_image,
                mask_image,
                depth_image,
                normal_image,
                frame_slider,
                frame_state,
            ],
        )
        frame_slider.change(
            fn=on_frame_change,
            inputs=[frame_slider, frame_state],
            outputs=[rgb_image, mask_image, depth_image, normal_image],
        )

    return demo


def create_demo() -> gr.Blocks:
    setup_logging()
    logger.info("DATASET_REPO = %r", DATASET_REPO)
    logger.info("HF_TOKEN set = %s, length = %d", HF_TOKEN is not None, len(HF_TOKEN) if HF_TOKEN else 0)
    logger.info("Loading preview subset from Hugging Face Hub...")
    preview_df = load_preview_dataframe()
    instance_df = build_preview_instance_dataframe(preview_df)
    logger.info("Loaded %d preview rows for %d preview instance(s).", len(preview_df), len(instance_df))
    return build_app(instance_df, preview_df)


def gradio_hot_reload_enabled() -> bool:
    return os.environ.get("GRADIO_HOT_RELOAD", "false").lower() in {"1", "true", "yes"}


def main() -> None:
    launch_kwargs: dict[str, Any] = {}
    server_name = os.environ.get("GRADIO_SERVER_NAME")
    server_port = os.environ.get("GRADIO_SERVER_PORT") or os.environ.get("PORT")
    if server_name:
        launch_kwargs["server_name"] = server_name
    if server_port:
        launch_kwargs["server_port"] = int(server_port)
    demo.launch(**launch_kwargs)


demo = create_demo()


if __name__ == "__main__" and not gradio_hot_reload_enabled():
    main()