app.py

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

os.environ['HTTP_PROXY'] = 'http://192.168.48.17:18000'
os.environ['HTTPS_PROXY'] = 'http://192.168.48.17:18000'

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)

@functools.lru_cache()
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):
    """
    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: # 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')
    os.environ.pop('HTTPS_PROXY')

    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"])))