Spaces:
Sleeping
Sleeping
| import os | |
| import random | |
| import gradio as gr | |
| import torch | |
| import functools | |
| from PIL import Image | |
| from datasets import load_dataset | |
| from feature_extractors.uni3d_embedding_encoder import Uni3dEmbeddingEncoder | |
| MAX_BATCH_SIZE = 16 | |
| MAX_QUEUE_SIZE = 10 | |
| MAX_K_RETRIEVAL = 20 | |
| cache_dir = "./.cache" | |
| encoder = Uni3dEmbeddingEncoder(cache_dir) | |
| 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) | |
| def get_embedding(option, modality, angle=None): | |
| save_path = f'data/objaverse_{option}_{modality + (("_" + str(angle)) if angle is not None else "")}_embeddings.pt' | |
| if os.path.exists(save_path): | |
| return torch.load(save_path) | |
| else: | |
| return gr.Error(f"Embedding file not found: {save_path}") | |
| def predict(xb, xq, top_k): | |
| xb = xb.to(xq.device) | |
| sim = xq @ xb.T # (nq, nb) | |
| _, indices = sim.topk(k=top_k, largest=True) | |
| return indices | |
| 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 == "": | |
| raise gr.Error("Please enter a textual query") | |
| if len(textual_query.split()) > 20: | |
| gr.Warning("Retrieval result may be inaccurate due to long textual query") | |
| if len(modality_list) == 0: | |
| raise gr.Error("Please select at least one modality") | |
| def _retrieve_3D_models(query, top_k, modals:list): | |
| option = "uni3d" | |
| op = "add" | |
| is_text = True if "text" in modals else False | |
| is_3D = True if "3D" in modals else False | |
| if is_text: | |
| modals.remove("text") | |
| if is_3D: | |
| modals.remove("3D") | |
| angles = modals | |
| # get base embeddings | |
| embeddings = [] | |
| if is_text: | |
| embeddings.append(get_embedding(option, "text")) | |
| if len(angles) > 0: | |
| for angle in angles: | |
| embeddings.append(get_embedding(option, "image", angle=angle)) | |
| if is_3D: | |
| embeddings.append(get_embedding(option, "3D")) | |
| ## fuse base embeddings | |
| if len(embeddings) > 1: | |
| if op == "concat": | |
| embeddings = torch.cat(embeddings, dim=-1) | |
| elif op == "add": | |
| embeddings = sum(embeddings) | |
| else: | |
| raise ValueError(f"Unsupported operation: {op}") | |
| embeddings /= embeddings.norm(dim=-1, keepdim=True) | |
| else: | |
| embeddings = embeddings[0] | |
| # encode query embeddings | |
| xq = encoder.encode_query(query) | |
| if op == "concat": | |
| xq = xq.repeat(1, embeddings.shape[-1] // xq.shape[-1]) # repeat to be aligned with the xb | |
| xq /= xq.norm(dim=-1, keepdim=True) | |
| pred_ind_list = predict(embeddings, xq, top_k) | |
| 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_and_id(index) for index in indices] | |
| def get_sub_dataset(sub_dataset_id, sorted=False): | |
| """ | |
| get sub-dataset by sub_dataset_id [1, 1000] | |
| Returns: | |
| caption: 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 = [] | |
| 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, results | |
| def feel_lucky(is_sorted): | |
| sub_dataset_id = random.randint(1, 1000) | |
| 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 | |
| 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 (paper not ready yet)""") | |
| 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. | |
| 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. | |
| 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, value="Super Mario") | |
| modality_list = gr.CheckboxGroup(label="Modality List", value=["text", "front", "back", "left", "right", "above", | |
| "below", "diag_above", "diag_below", "3D"], | |
| 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 *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, 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) | |
| # 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, 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') | |
| if __name__ == "__main__": | |
| launch() | |
| # print(len(retrieve_3D_models("A chair with a wooden frame and a cushioned seat", 5, ["3D", "diag_above", "diag_below"]))) |