init version of app

README.md

@@ -5,7 +5,7 @@ colorFrom: blue
 colorTo: green
 sdk: gradio
 python_version: 3.9.13
-sdk_version: 4.12.0
+sdk_version: 3.50.2
 app_file: app.py
 pinned: false
 license: cc-by-nc-sa-4.0

app.py

@@ -98,7 +98,7 @@ def expand2square(pil_img, background_color):
         return result
 
 
-def preprocess(predictor, input_image, chk_group=None, segment=True):
+def preprocess(predictor, input_image, chk_group=None, segment=False):
     RES = 1024
     input_image.thumbnail([RES, RES], Image.Resampling.LANCZOS)
     if chk_group is not None:
@@ -403,13 +403,76 @@ def create_bones_scene(bones, joint_color=[66, 91, 140], bone_color=[119, 144, 1
     return mesh
 
 
-def run_pipeline(model_items, cfgs, input_img, device):
+def save_mesh(mesh, file_path):
+    obj_file = file_path
+    idx = 0
+    print("Writing mesh: ", obj_file)
+    with open(obj_file, "w") as f:
+        # f.write(f"mtllib {fname}.mtl\n")
+        f.write("g default\n")
+
+        v_pos = mesh.v_pos[idx].detach().cpu().numpy() if mesh.v_pos is not None else None
+        v_nrm = mesh.v_nrm[idx].detach().cpu().numpy() if mesh.v_nrm is not None else None
+        v_tex = mesh.v_tex[idx].detach().cpu().numpy() if mesh.v_tex is not None else None
+
+        t_pos_idx = mesh.t_pos_idx[0].detach().cpu().numpy() if mesh.t_pos_idx is not None else None
+        t_nrm_idx = mesh.t_nrm_idx[0].detach().cpu().numpy() if mesh.t_nrm_idx is not None else None
+        t_tex_idx = mesh.t_tex_idx[0].detach().cpu().numpy() if mesh.t_tex_idx is not None else None
+
+        print("    writing %d vertices" % len(v_pos))
+        for v in v_pos:
+            f.write('v {} {} {} \n'.format(v[0], v[1], v[2]))
+
+        if v_nrm is not None:
+            print("    writing %d normals" % len(v_nrm))
+            assert(len(t_pos_idx) == len(t_nrm_idx))
+            for v in v_nrm:
+                f.write('vn {} {} {}\n'.format(v[0], v[1], v[2]))
+
+        # faces
+        f.write("s 1 \n")
+        f.write("g pMesh1\n")
+        f.write("usemtl defaultMat\n")
+
+        # Write faces
+        print("    writing %d faces" % len(t_pos_idx))
+        for i in range(len(t_pos_idx)):
+            f.write("f ")
+            for j in range(3):
+                f.write(' %s/%s/%s' % (str(t_pos_idx[i][j]+1), '' if v_tex is None else str(t_tex_idx[i][j]+1), '' if v_nrm is None else str(t_nrm_idx[i][j]+1)))
+            f.write("\n")
+
+
+def process_mesh(shape, name):
+    mesh = shape.clone()
+    output_glb = f'./{name}.glb'
+    output_obj = f'./{name}.obj'
+
+    # save the obj file for download
+    save_mesh(mesh, output_obj)
+
+    # save the glb for visualize
+    mesh_tri = trimesh.Trimesh(
+        vertices=mesh.v_pos[0].detach().cpu().numpy(), 
+        faces=mesh.t_pos_idx[0][..., [2,1,0]].detach().cpu().numpy(), 
+        process=False,  
+        maintain_order=True
+    )
+    mesh_tri.visual.vertex_colors = (mesh.v_nrm[0][..., [2,1,0]].detach().cpu().numpy() + 1.0) * 0.5 * 255.0
+    mesh_tri.export(file_obj=output_glb)
+
+    return output_glb, output_obj
+
+
+def run_pipeline(model_items, cfgs, input_img):
     epoch = 999
     total_iter = 999999
     model = model_items[0]
     memory_bank = model_items[1]
     memory_bank_keys = model_items[2]
 
+    device = f'cuda:{_GPU_ID}'
+
     input_image = torch.stack([torchvision.transforms.ToTensor()(input_img)], dim=0).to(device)
 
     with torch.no_grad():
@@ -455,7 +518,7 @@ def run_pipeline(model_items, cfgs, input_img, device):
         gray_light = FixedDirectionLight(direction=torch.FloatTensor([0, 0, 1]).to(device), amb=0.2, diff=0.7)
 
         image_pred, mask_pred, _, _, _, shading = model.render(
-            shape, texture_pred, mvp, w2c, campos, 256, background=model.background_mode, 
+            shape, texture_pred, mvp, w2c, campos, (256, 256), background=model.background_mode, 
             im_features=im_features, light=gray_light, prior_shape=prior_shape, render_mode='diffuse', 
             render_flow=False, dino_pred=None, im_features_map=im_features_map
         )
@@ -469,7 +532,7 @@ def run_pipeline(model_items, cfgs, input_img, device):
         nv_meshes = make_mesh(verts=bones_meshes.verts_padded(), faces=bones_meshes.faces_padded()[0:1],
                                 uvs=bones_meshes.textures.verts_uvs_padded(), uv_idx=bones_meshes.textures.faces_uvs_padded()[0:1],
                                 material=material_texture.Texture2D(bones_meshes.textures.maps_padded()))
-        buffers = render_mesh(dr.RasterizeGLContext(), nv_meshes, mvp, w2c, campos, nv_meshes.material, lgt=gray_light, feat=im_features, dino_pred=None, resolution=256, bsdf="diffuse")
+        buffers = render_mesh(dr.RasterizeGLContext(), nv_meshes, mvp, w2c, campos, nv_meshes.material, lgt=gray_light, feat=im_features, dino_pred=None, resolution=(256,256), bsdf="diffuse")
         
         shaded = buffers["shaded"].permute(0, 3, 1, 2)
         bone_image = shaded[:, :3, :, :]
@@ -481,20 +544,10 @@ def run_pipeline(model_items, cfgs, input_img, device):
         mesh_image = save_images(shading, mask_pred)
         mesh_bones_image = save_images(image_with_bones, mask_final)
 
-        final_shape = shape.clone()
-        prior_shape = prior_shape.clone()
-
-        final_mesh_tri = trimesh.Trimesh(
-            vertices=final_shape.v_pos[0].detach().cpu().numpy(), 
-            faces=final_shape.t_pos_idx[0].detach().cpu().numpy(), 
-            process=False,  
-            maintain_order=True)
-        prior_mesh_tri = trimesh.Trimesh(
-            vertices=prior_shape.v_pos[0].detach().cpu().numpy(), 
-            faces=prior_shape.t_pos_idx[0].detach().cpu().numpy(), 
-            process=False,  
-            maintain_order=True)
+        shape_glb, shape_obj = process_mesh(shape, 'reconstruced_shape')
+        base_shape_glb, base_shape_obj = process_mesh(prior_shape, 'reconstructed_base_shape')
 
+        return mesh_image, mesh_bones_image, shape_glb, shape_obj, base_shape_glb, base_shape_obj
 
 
 def run_demo():
@@ -582,7 +635,6 @@ def run_demo():
                         with gr.Column():
                             input_processing = gr.CheckboxGroup(['Use SAM to center animal'], 
                                                                 label='Input Image Preprocessing',
-                                                                 value=['Use SAM to center animal'],
                                                                  info='untick this, if animal is already centered, e.g. in example images')
                         # with gr.Column():
                         #     output_processing = gr.CheckboxGroup(['Background Removal'], label='Output Image Postprocessing', value=[]) 
@@ -599,23 +651,26 @@ def run_demo():
                     #     with gr.Column():
                     #         crop_size = gr.Number(192, label='Crop size')
                     # crop_size = 192
-                run_btn = gr.Button('Generate', variant='primary', interactive=True)
+                run_btn = gr.Button('Reconstruct', variant='primary', interactive=True)
         with gr.Row():
             view_1 = gr.Image(interactive=False, height=256, show_label=False)
             view_2 = gr.Image(interactive=False, height=256, show_label=False)
         with gr.Row():
-            shape_1 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  label="Reconstructed Model")
-            shape_2 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  label="Bank Base Shape Model")
+            shape_1 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  height=512, label="Reconstructed Model")
+            shape_1_download = gr.File(label="Download Full Reconstructed Model")
+        with gr.Row():
+            shape_2 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  height=512, label="Bank Base Shape Model")
+            shape_2_download = gr.File(label="Download Full Bank Base Shape Model")
 
         run_btn.click(fn=partial(preprocess, predictor), 
                         inputs=[input_image, input_processing], 
                         outputs=[processed_image_highres, processed_image], queue=True
             ).success(fn=partial(run_pipeline, model_items, model_cfgs), 
-                        inputs=[processed_image, device],
-                        outputs=[view_1, view_2, shape_1, shape_2]
+                        inputs=[processed_image],
+                        outputs=[view_1, view_2, shape_1, shape_1_download, shape_2, shape_2_download]
                         )
         demo.queue().launch(share=True, max_threads=80)
-        # _, local_url, share_url = demo.launch(share=True, server_name="0.0.0.0", server_port=23425)
+        # _, local_url, share_url = demo.queue().launch(share=True, server_name="0.0.0.0", server_port=23425)
         # print('local_url: ', local_url)