update app.py

README.md

@@ -7,6 +7,11 @@ sdk: docker
 app_port: 7860
 app_file: app.py
 pinned: false
+models:
+- timm/eva02_enormous_patch14_plus_clip_224.laion2b_s9b_b144k
+- BAAI/Uni3D
+datasets:
+- VAST-AI/LD-T3D
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,8 +1,9 @@
 import os
+import random
 import gradio as gr
-import numpy as np
 import torch
 import functools
+from PIL import Image
 from datasets import load_dataset
 from feature_extractors.uni3d_embedding_encoder import Uni3dEmbeddingEncoder
 
@@ -19,6 +20,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 source_id_list = torch.load("data/source_id_list.pt")
 source_to_id = {source_id: i for i, source_id in enumerate(source_id_list)}
 dataset = load_dataset("VAST-AI/LD-T3D", name=f"rendered_imgs_diag_above", split="base", cache_dir=cache_dir)
+relation = load_dataset("VAST-AI/LD-T3D", split="full", cache_dir=cache_dir)
 
 @functools.lru_cache()
 def get_embedding(option, modality, angle=None):
@@ -34,8 +36,8 @@ def predict(xb, xq, top_k):
     _, indices = sim.topk(k=top_k, largest=True)
     return indices
 
-def get_image(index):
-    return dataset[index]["image"]
+def get_image_and_id(index):
+    return dataset[index]["image"], dataset[index]["source_id"]
 
 def retrieve_3D_models(textual_query, top_k, modality_list):
     if textual_query == "":
@@ -88,35 +90,106 @@ def retrieve_3D_models(textual_query, top_k, modality_list):
         return pred_ind_list[0].cpu().tolist() # we have only one query
 
     indices = _retrieve_3D_models(textual_query, top_k, modality_list)
-    return [get_image(index) for index in indices]
+    return [get_image_and_id(index) for index in indices]
+
+def get_sub_dataset(sub_dataset_id):
+    """
+    get sub-dataset by sub_dataset_id [1, 1000]
+
+    Returns:
+        caption: str
+        difficulty: str
+        images: list of tuple (PIL.Image, str)
+    """
+    rel = relation[sub_dataset_id - 1]
+    target_ids, GT_ids, caption, difficulty = set(rel["target_ids"]), set(rel["GT_ids"]), rel["caption"], rel["difficulty"]
+    negative_ids = target_ids - GT_ids
+
+    def handle_image(image, is_gt=False):
+        "image is a PIL.Image object, surround the image with green border if is_gt, else red border"
+        border_color = (0, 255, 0) if is_gt else (255, 0, 0)
+        border_width = 5
+        new_image = Image.new("RGBA", (image.width + 2 * border_width, image.height + 2 * border_width), border_color)
+        new_image.paste(image, (border_width, border_width))
+        return new_image
+
+    results = []
+    for gt_id in GT_ids:
+        image, source_id = get_image_and_id(source_to_id[gt_id])
+        results.append((handle_image(image, True), source_id))
+    for neg_id in negative_ids:
+        image, source_id = get_image_and_id(source_to_id[neg_id])
+        results.append((handle_image(image, False), source_id))
+
+    return caption, difficulty, results
+
+def feel_lucky():
+    sub_dataset_id = random.randint(1, 1000)
+    return sub_dataset_id, *get_sub_dataset(sub_dataset_id)
 
 def launch():
-    with gr.Blocks() as demo:
-        with gr.Row():
-            textual_query = gr.Textbox(label="Textual Query", autofocus=True,
-                                       placeholder="A chair with a wooden frame and a cushioned seat")
-            modality_list = gr.CheckboxGroup(label="Modality List", value=[],
-                                             choices=["text", "front", "back", "left", "right", "above", 
-                                                      "below", "diag_above", "diag_below", "3D"])
-        with gr.Row():
-            top_k = gr.Slider(minimum=1, maximum=MAX_K_RETRIEVAL, step=1, label="Top K Retrieval Result", 
-                              value=5, scale=2)
-            run = gr.Button("Search", scale=1)
-            clear_button = gr.ClearButton(scale=1)
-        with gr.Row():
-            output = gr.Gallery(format="webp", label="Retrieval Result", columns=5, type="pil")
-        run.click(retrieve_3D_models, [textual_query, top_k, modality_list], output, 
-                #   batch=True, max_batch_size=MAX_BATCH_SIZE
-                  )
-        clear_button.click(lambda: ["", 5, [], []], outputs=[textual_query, top_k, modality_list, output])
-        examples = gr.Examples(examples=[["An ice cream with a cherry on top", 10, ["text", "front", "back", "left", "right", "above", "below", "diag_above", "diag_below", "3D"]],
-                                         ["A mid-age castle", 10, ["text", "front", "back", "left", "right", "above", "below", "diag_above", "diag_below", "3D"]],
-                                         ["A coke", 10, ["text", "front", "back", "left", "right", "above", "below", "diag_above", "diag_below", "3D"]]],
-                            inputs=[textual_query, top_k, modality_list],
-                            # cache_examples=True,
-                            outputs=output,
-                            fn=retrieve_3D_models)
-        
+    with gr.Blocks() as demo: # https://sketchfab.com/3d-models/fd30f87848c9454c9225eccc39726787
+        md = gr.Markdown(r"""## LD-T3D: A Large-scale and Diverse Benchmark for Text-based 3D Model Retrieval
+**Official 🤗 Gradio demo** for [LD-T3D: A Large-scale and Diverse Benchmark for Text-based 3D Model Retrieval](arxiv.org)""")
+        with gr.Tab("Retrieval Visualization"):
+            with gr.Row():
+                md2 = gr.Markdown(r"""### Visualization for Text-Based-3D Model Retrieval
+We build a visualization demo to demonstrate the text-based-3D model retrievals. Due to the memory limitation of HF Space, we only support the [Uni3D](https://github.com/baaivision/Uni3D) which has shown an excellent performance in our benchmark. 
+
+**Note**: 
+                                  
+The *Modality List* refers to the features ensembled by the retrieval methods. According to our experiment results, basically the more modalities, the better performance the methods gets.""")
+            with gr.Row():
+                textual_query = gr.Textbox(label="Textual Query", autofocus=True,
+                                        placeholder="A chair with a wooden frame and a cushioned seat")
+                modality_list = gr.CheckboxGroup(label="Modality List", value=[],
+                                                choices=["text", "front", "back", "left", "right", "above", 
+                                                        "below", "diag_above", "diag_below", "3D"])
+            with gr.Row():
+                top_k = gr.Slider(minimum=1, maximum=MAX_K_RETRIEVAL, step=1, label="Top K Retrieval Result", 
+                                value=5, scale=2)
+                run = gr.Button("Search", scale=1, variant='primary')
+                clear_button = gr.ClearButton(scale=1)
+            with gr.Row():
+                output = gr.Gallery(format="webp", label="Retrieval Result", columns=5, type="pil", interactive=False)
+            run.click(retrieve_3D_models, [textual_query, top_k, modality_list], output, 
+                    #   batch=True, max_batch_size=MAX_BATCH_SIZE
+                    )
+            clear_button.click(lambda: ["", 5, [], []], outputs=[textual_query, top_k, modality_list, output])
+            examples = gr.Examples(examples=[["An ice cream with a cherry on top", 10, ["text", "front", "back", "left", "right", "above", "below", "diag_above", "diag_below", "3D"]],
+                                            ["A mid-age castle", 10, ["text", "front", "back", "left", "right", "above", "below", "diag_above", "diag_below", "3D"]],
+                                            ["A coke", 10, ["text", "front", "back", "left", "right", "above", "below", "diag_above", "diag_below", "3D"]]],
+                                inputs=[textual_query, top_k, modality_list],
+                                outputs=output,
+                                fn=retrieve_3D_models)
+        with gr.Tab("Federated Dataset"):
+            md3 = gr.Markdown(r"""### Visualization for Federated Dataset
+We provide a federated dataset that contains **1000** textual queries and **89K** 3D models, which corresponds to **1000** sub-datasets with around **100** 3D models. 
+In total, there is 100K pairs of text-to-3D model relationships. 
+                            
+Here is a visualization of the dataset.
+
+**Usage:** 
+                            
+1. You can click the "I'm Feeling Lucky !" button to randomly select a sub-dataset.
+2. Or you can **Enter** to submit a Sub-dataset ID in **[1, 1000]**, which you can find details in our dataset [LD-T3D](https://huggingface.co/datasets/VAST-AI/LD-T3D), to search for the corresponding sub-dataset.
+
+**Note:**
+
+The *Query* is used in this sub-dataset. The *Difficulty* is a coarse label for the textual query, which is divided into **easy**, **medium**, and **hard**, basically submit to the rule in our paper. 
+The color surrounding the 3D model indicates whether it is a good fit for the textual query. A **<span style="color:#00FF00">green</span>** color suggests a Ground Truth, while a **<span style="color:#FF0000">red</span>** color indicates a mismatch.""")
+            with gr.Row():
+                lucky = gr.Button("I'm Feeling Lucky !", scale=1, variant='primary')
+                query_id = gr.Number(label="Sub-dataset ID", scale=1, minimum=1, maximum=1000, step=1, interactive=True)
+                query = gr.Textbox(label="Textual Query", scale=3, interactive=False)
+                difficulty = gr.Textbox(label="Query Difficulty", scale=1, interactive=False)
+                # model3d = gr.Model3D(interactive=False, scale=1)
+            with gr.Row():
+                output2 = gr.Gallery(format="webp", label="3D Models in Sub-dataset", columns=5, type="pil", interactive=False)
+            
+            lucky.click(feel_lucky, outputs=[query_id, query, difficulty, output2])
+            query_id.submit(get_sub_dataset, query_id, [query, difficulty, output2])
+
     demo.queue(max_size=10)
 
     os.environ.pop('HTTP_PROXY')