update app

Dockerfile

@@ -22,7 +22,7 @@ RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
 # note that you may need to modify the TORCH_CUDA_ARCH_LIST in the setup.py file
 ENV TORCH_CUDA_ARCH_LIST="8.6"
 
-# Install Pointnet2_PyTorch
+# Install Pointnet2_PyTorch, pip install git+ won't work for unknown reason
 RUN git clone https://github.com/yuanze1024/Pointnet2_PyTorch.git && cd Pointnet2_PyTorch/pointnet2_ops_lib && pip install .
 
 COPY --chown=user:user . /code

app.py

@@ -89,13 +89,12 @@ def retrieve_3D_models(textual_query, top_k, modality_list):
     indices = _retrieve_3D_models(textual_query, top_k, modality_list)
     return [get_image_and_id(index) for index in indices]
 
-def get_sub_dataset(sub_dataset_id):
+def get_sub_dataset(sub_dataset_id, sorted=False):
     """
     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]
@@ -111,18 +110,23 @@ def get_sub_dataset(sub_dataset_id):
         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))
+    if not sorted:
+        for ind in target_ids:
+            image, source_id = get_image_and_id(source_to_id[ind])
+            results.append((handle_image(image, True if ind in GT_ids else False), source_id))
+    else:
+        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
+    return caption, results
 
-def feel_lucky():
+def feel_lucky(is_sorted):
     sub_dataset_id = random.randint(1, 1000)
-    return sub_dataset_id, *get_sub_dataset(sub_dataset_id)
+    return sub_dataset_id, *get_sub_dataset(sub_dataset_id, is_sorted)
 
 def launch():
     with gr.Blocks() as demo: # https://sketchfab.com/3d-models/fd30f87848c9454c9225eccc39726787
@@ -131,14 +135,17 @@ def launch():
         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. 
+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. 
+What's more, **we only search in a subset of Objaverse, which contains 89K 3D models**.
 
 **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.""")
+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.
+
+Also, you may want to ckeck the 3D model in a 3D model viewer, in that case, you can visit [Objaverse](https://objaverse.allenai.org/explore) for exploration.""")
             with gr.Row():
-                textual_query = gr.Textbox(label="Textual Query", autofocus=True,
-                                        placeholder="A chair with a wooden frame and a cushioned seat")
+                textual_query = gr.Textbox(label="Textual Query", autofocus=True, value="Super Mario")
                 modality_list = gr.CheckboxGroup(label="Modality List", value=[],
                                                 choices=["text", "front", "back", "left", "right", "above", 
                                                         "below", "diag_above", "diag_below", "3D"])
@@ -173,19 +180,22 @@ Here is a visualization of the 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.""")
+The *Query* is used in this sub-dataset. The *Sorted* will put the Ground Truths in the front of the results.
+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_id = gr.Number(label="Sub-dataset ID", scale=1, minimum=1, maximum=1000, step=1, interactive=True, value=986)
+                is_sorted = gr.Checkbox(value=False, label="", scale=1, info="Sorted")
                 query = gr.Textbox(label="Textual Query", scale=3, interactive=False)
-                difficulty = gr.Textbox(label="Query Difficulty", scale=1, 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])
+            lucky.click(feel_lucky, inputs=is_sorted, outputs=[query_id, query, output2])
+            query_id.submit(get_sub_dataset, [query_id, is_sorted], [query, output2])
+            is_sorted.change(get_sub_dataset, [query_id, is_sorted], [query, output2])
 
     demo.queue(max_size=10)
     demo.launch(server_name='0.0.0.0')