File size: 1,234 Bytes
f74bb58 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
import gradio as gr
import os
import pytorch_lightning as pl
import torch as th
import open3d as o3d
import numpy as np
import trimesh as tm
from models.model import Model
model = Model()
ckpg = th.load("./checkpoints/epoch=99-step=6000.ckpt")
model.load_state_dict(ckpg["state_dict"])
def process_mesh(mesh_file_name):
mesh = tm.load_mesh(mesh_file_name)
v = th.tensor(mesh.vertices, dtype=th.float)
n = th.tensor(mesh.vertex_normals, dtype=th.float)
with th.no_grad():
v, f, n, _ = model(v.unsqueeze(0), n.unsqueeze(0))
mesh = tm.Trimesh(vertices=v.squeeze(0),
faces=f.squeeze(0),
vertex_normals=n.squeeze(0))
obj_path = "./sample.obj"
mesh.export(obj_path)
return obj_path
demo = gr.Interface(
fn=process_mesh,
inputs=gr.Model3D(),
outputs=gr.Model3D(
clear_color=[0.0, 0.0, 0.0, 0.0], label="3D Model"),
examples=[
[os.path.join(os.path.dirname(__file__), "files\\bunny_n1_hi_50.obj")],
[os.path.join(os.path.dirname(__file__), "files\\child_n2_80.obj")],
[os.path.join(os.path.dirname(__file__), "files\\eight_n3_70.obj")],
],
)
if __name__ == "__main__":
demo.launch() |