| |
| """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) |
|
|
|
|
| |
| |
| |
|
|
| 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}"), |
| ] |
|
|
|
|
| |
| |
| |
|
|
| 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) |
|
|
|
|
| |
| |
| |
|
|
| 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() |
|
|
