ShapeNetViz / app.py
Edward Beeching
added get code
52bd690
raw
history blame contribute delete
No virus
5.36 kB
import time
import gradio as gr
import os
import numpy as np
from PIL import Image
import math
from collections import defaultdict
import os
from huggingface_hub import HfApi
def greet(name):
return "Hello " + name + "!"
def get_dataset_classes():
# return {
# "airplane": [
# "airplane/0.glb",
# "airplane/1.glb",
# "airplane/2.glb",
# "airplane/3.glb",
# "airplane/4.glb",
# "airplane/5.glb",
# "airplane/6.glb",
# "airplane/7.glb",
# ],
# "bench": [
# "bench/0.glb",
# "bench/1.glb",
# "bench/10.glb",
# "bench/11.glb",
# "bench/12.glb",
# "bench/13.glb",
# "bench/14.glb",
# "bench/2.glb",
# "bench/3.glb",
# "bench/4.glb",
# "bench/5.glb",
# "bench/6.glb",
# "bench/7.glb",
# "bench/8.glb",
# "bench/9.glb",
# ],
# "camera": [
# "camera/0.glb",
# "camera/1.glb",
# "camera/2.glb",
# "camera/3.glb",
# "camera/4.glb",
# "camera/5.glb",
# "camera/6.glb",
# "camera/7.glb",
# ],
# "dishwasher": [
# "dishwasher/0.glb",
# "dishwasher/1.glb",
# "dishwasher/10.glb",
# "dishwasher/11.glb",
# "dishwasher/2.glb",
# "dishwasher/3.glb",
# "dishwasher/4.glb",
# "dishwasher/5.glb",
# "dishwasher/6.glb",
# "dishwasher/7.glb",
# "dishwasher/8.glb",
# "dishwasher/9.glb",
# ],
# "jar": [
# "jar/0.glb",
# "jar/1.glb",
# "jar/2.glb",
# "jar/3.glb",
# "jar/4.glb",
# "jar/5.glb",
# "jar/6.glb",
# "jar/7.glb",
# "jar/8.glb",
# ],
# "motorcycle": [
# "motorcycle/0.glb",
# "motorcycle/1.glb",
# "motorcycle/10.glb",
# "motorcycle/2.glb",
# "motorcycle/3.glb",
# "motorcycle/4.glb",
# "motorcycle/5.glb",
# "motorcycle/6.glb",
# "motorcycle/7.glb",
# "motorcycle/8.glb",
# "motorcycle/9.glb",
# ],
# "printer": [
# "printer/0.glb",
# "printer/1.glb",
# "printer/10.glb",
# "printer/11.glb",
# "printer/2.glb",
# "printer/3.glb",
# "printer/4.glb",
# "printer/5.glb",
# "printer/6.glb",
# "printer/7.glb",
# "printer/8.glb",
# "printer/9.glb",
# ],
# "sofa": [
# "sofa/0.glb",
# "sofa/1.glb",
# "sofa/10.glb",
# "sofa/11.glb",
# "sofa/12.glb",
# "sofa/2.glb",
# "sofa/3.glb",
# "sofa/4.glb",
# "sofa/5.glb",
# "sofa/6.glb",
# "sofa/7.glb",
# "sofa/8.glb",
# "sofa/9.glb",
# ],
# "washer": [
# "washer/0.glb",
# "washer/1.glb",
# "washer/10.glb",
# "washer/11.glb",
# "washer/2.glb",
# "washer/3.glb",
# "washer/4.glb",
# "washer/5.glb",
# "washer/6.glb",
# "washer/7.glb",
# "washer/8.glb",
# "washer/9.glb",
# ],
# }
hf_api = HfApi()
info = hf_api.dataset_info(repo_id="ShapeNet/Example")
dataset_classes = defaultdict(list)
for file in info.siblings:
if ".glb" 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(model_name):
# wget the glb file from the datasets repo
print(model_name)
return f"{model_name}/0.glb"
def update_model_list(choice_class):
print(f"inp1 changed {choice_class}")
return {"choices":dataset_dict[choice_class]}
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])
out2 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="3D Model")
inp.change(fn=update, inputs=inp, outputs=out2)
# with gr.Row():
# btn = gr.Button("Load model")
# btn.click(fn=update, inputs=inp, outputs=out2)
# demo = gr.Interface(
# fn=load_mesh,
# inputs=gr.Model3D(),
# outputs=[
# gr.Model3D(
# clear_color=[0.0, 0.0, 0.0, 0.0], label="3D Model"),
# gr.File(label="Download 3D Model")
# ],
# examples=[
# ],
demo.launch()