CarlCochet's picture
Update app.py
9fe0c24
import gradio as gr
from collections import defaultdict
import os
from huggingface_hub import HfApi, hf_hub_download
HUB_TOKEN = os.getenv("HUB_TOKEN")
REPO_ID = "ShapeNet/shapenetcore-gltf"
def get_dataset_classes():
hf_api = HfApi()
info = hf_api.dataset_info(repo_id=REPO_ID, token=HUB_TOKEN)
dataset_classes = defaultdict(list)
for file in info.siblings:
if ".gltf" in file.rfilename:
class_name = file.rfilename.split("/")[0]
dataset_classes[class_name].append(file.rfilename)
print(dataset_classes)
return dataset_classes
dataset_dict = get_dataset_classes()
dataset_classes = list(dataset_dict.keys())
default_models = dataset_dict[dataset_classes[0]]
def load_mesh(mesh_file_name):
return mesh_file_name, mesh_file_name
def update(asset_name):
split_model_path = asset_name.split("/")
asset_path = hf_hub_download(
repo_id=REPO_ID,
filename=split_model_path[1],
subfolder=split_model_path[0],
repo_type="dataset",
use_auth_token=HUB_TOKEN,
)
print(asset_name)
return asset_path
def update_models(class_name):
model_choices = dataset_dict[class_name]
return gr.Dropdown.update(choices=model_choices)
def update_model_list(class_name):
model_choices = dataset_dict[class_name]
return gr.Dropdown.update(choices=model_choices, value=model_choices[0])
def update_asset_path(model_name):
return REPO_ID + "/" + model_name
with gr.Blocks() as demo:
with gr.Row():
with gr.Column():
inp = gr.Dropdown(choices=dataset_classes, interactive=True, label="3D Model Class", value=dataset_classes[0])
out1 = gr.Dropdown(choices=default_models, interactive=True, label="3D Model", value=default_models[0])
out3 = gr.Textbox(value=REPO_ID + "/" + out1.value, label="Asset Path")
out2 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="3D Model")
# Update second dropdown
inp.change(fn=update_model_list, inputs=inp, outputs=out1)
# Update 3D model view
inp.change(fn=update, inputs=out1, outputs=out2)
out1.change(fn=update, inputs=out1, outputs=out2)
# Update path to asset
inp.change(fn=update_asset_path, inputs=out1, outputs=out3)
out1.change(fn=update_asset_path, inputs=out1, outputs=out3)
demo.launch()