Upload folder using huggingface_hub

third_party/dust3r_utils.py

@@ -135,7 +135,6 @@ def infer_match(images, model, vis=False, niter=300, lr=0.01, schedule='cosine',
             plt.plot([x0, x1 + W0], [y0, y1], '-+', color=cmap(i / (n_viz - 1)), scalex=False, scaley=False)
         plt.show(block=True)
 
-    matches_im0 = remap_points(images[0].shape, matches_im0)
     matches_im1 = remap_points(images[1].shape, matches_im1)
     return matches_im0, matches_im1
 

third_party/gen_baking.py

@@ -55,27 +55,17 @@ def load_objs_as_meshes_fast(
     path_manager: Optional[PathManager] = None,
 ):
     tex = None
-    if create_texture_atlas:
-        # TexturesAtlas type
-        tex = TexturesAtlas(atlas=[aux.texture_atlas.to(device)])
-    else:
-        # TexturesUV type
-        tex_maps = aux.texture_images
-        if tex_maps is not None and len(tex_maps) > 0:
-            verts_uvs = aux.verts_uvs.to(device)  # (V, 2)
-            faces_uvs = faces.textures_idx.to(device)  # (F, 3)
-            image = list(tex_maps.values())[0].to(device)[None]
-            tex = TexturesUV(verts_uvs=[verts_uvs], faces_uvs=[faces_uvs], maps=image)
+    tex_maps = aux.texture_images
+    if tex_maps is not None and len(tex_maps) > 0:
+        verts_uvs = aux.verts_uvs.to(device)  # (V, 2)
+        faces_uvs = faces.textures_idx.to(device)  # (F, 3)
+        image = list(tex_maps.values())[0].to(device)[None]
+        tex = TexturesUV(verts_uvs=[verts_uvs], faces_uvs=[faces_uvs], maps=image)
     mesh = Meshes( verts=[verts.to(device)], faces=[faces.verts_idx.to(device)], textures=tex)
     return mesh
 
 
 def get_triangle_to_triangle(tri_1, tri_2, img_refined):
-    '''
-        args:
-            tri_1: 
-            tri_2: 
-    '''
     r1 = cv2.boundingRect(tri_1)
     r2 = cv2.boundingRect(tri_2)
     
@@ -115,11 +105,8 @@ def back_projection(
     erode_scale=1/100.0, 
     device="cuda:0"
 ):
-    # obj_file: 带有uv的obj
-    # init_texture_file： 初始展开的uv贴图
-    # render_resolution 正面视角渲染分辨率
-    # uv_resolution 贴图分辨率
-    # thres：normal threshold
+    # obj_file: obj with uv
+    # init_texture_file： uv texture image
 
     os.makedirs(dst_dir, exist_ok=True)
 
@@ -151,8 +138,6 @@ def back_projection(
     ], dim=-1)
     triangle_uvs = np.clip(np.round(np.float32(triangle_uvs.cpu())).astype(np.int64), 0, uv_resolution-1)
 
-    # import ipdb;ipdb.set_trace()
-
     
     R0, T0 = look_at_view_transform(camera_dist, views[0][0], views[0][1])
 
@@ -197,7 +182,7 @@ def back_projection(
  
     # seamlessClone
     try:
-        images = cv2.seamlessClone(src, dst, src_mask, center, cv2.NORMAL_CLONE) # more qingxi
+        images = cv2.seamlessClone(src, dst, src_mask, center, cv2.NORMAL_CLONE) # better 
         # images = cv2.seamlessClone(src, dst, src_mask, center, cv2.MIXED_CLONE)
     except Exception as err:
         print(f"\n\n Warning seamlessClone error: {err} \n\n")
@@ -225,7 +210,6 @@ def back_projection(
     
     t0 = time.time()
     
-    
     SOFT_NORM = True # process big angle-diff face, true:flase? coeff:skip
     
     for k in faces_covered:

third_party/mesh_baker.py

@@ -25,8 +25,7 @@ class MeshBaker:
         align_model = "third_party/weights/DUSt3R_ViTLarge_BaseDecoder_512_dpt",
         device = "cuda:0", 
         align_times = 1,
-        iou_thresh = 0.8, 
-        force_baking_ele_list = None,
+        iou_thresh = 0.8,
         save_memory = False
     ):
         self.device = device
@@ -36,8 +35,6 @@ class MeshBaker:
         self.align_times = align_times
         self.align_model.eval()
         self.iou_thresh = iou_thresh
-        self.force_baking_ele_list = [] if force_baking_ele_list is None else force_baking_ele_list
-        self.force_baking_ele_list = [int(_) for _ in self.force_baking_ele_list]
         set_parameter_grad_false(self.align_model)
         print('baking align model', get_parameter_number(self.align_model))
     
@@ -75,7 +72,7 @@ class MeshBaker:
         torch.cuda.empty_cache()
         return res
     
-    def call(self, save_folder):
+    def call(self, save_folder, force=False, front='auto', others=['180°'], align_times=3, seed=0):
         obj_path         = os.path.join(save_folder, "mesh.obj")
         raw_texture_path = os.path.join(save_folder, "texture.png")
         views_pil        = os.path.join(save_folder, "views.jpg")
@@ -93,46 +90,76 @@ class MeshBaker:
             cond_pil, views = views[0], views[1:]
         else:
             raise FileNotFoundError("views file not found")
-                
-        rendered_views = render_func(obj_path, elev=0, n_views=2)
         
-        for ele_idx, ele in enumerate([0, 180]):
+        others = [int(x.replace("°", "")) for x in others]
+        
+        if len(others)==0:
+            rendered_views = render_func(obj_path, elev=0, n_views=1)
+        elif len(others)==1 and others[0]==180:
+            rendered_views = render_func(obj_path, elev=0, n_views=2)
+        else:
+            rendered_views = render_func(obj_path, elev=0, n_views=6)
+            
+        print(f"Need baking views are {others}")
+        others = [0] + others
+        
+        seed_everything(seed)
+        
+        for ele_idx, ele in enumerate([0, 60, 120, 180, 240, 300]):
+            
+            if ele not in others: continue
+            
+            print(f"\n Baking view ele_{ele} ...")
             
             if ele == 0:
-                aligned_cond, cond_info, _ = self.align_and_check(cond_pil, rendered_views[0], align_times=self.align_times)
-                aligned_cond.save(save_folder + f'/aligned_cond.jpg')
+                if front == 'multi-view front view' or front == 'auto':
+                    aligned_cond, cond_info, _ = self.align_and_check(cond_pil, rendered_views[0], align_times=self.align_times)
+                    if cond_info is None: continue
+                    aligned_cond.convert("RGB").save(save_folder + f'/aligned_cond.jpg')
+                    if front == 'multi-view front view':
+                        aligned_img, info = aligned_cond, cond_info
+                        print("Using Cond Image to bake front view")
         
-                aligned_img, info, _ = self.align_and_check(views[0], rendered_views[0], align_times=self.align_times)
-                aligned_img.save(save_folder + f'/aligned_{ele}.jpg')
+                if front == 'input image' or front == 'auto':
+                    aligned_img, info, _ = self.align_and_check(views[0], rendered_views[0], align_times=self.align_times)
+                    if info is None: continue
+                    aligned_img.save(save_folder + f'/aligned_{ele}.jpg')
+                    print("Using Input Image to bake front view")
                 
-                if info['mask_iou'] < cond_info['mask_iou']:
-                    print("Using Cond Image to bake front view")
-                    aligned_img = aligned_cond
-                    info = cond_info
+                if front == 'auto' and info['mask_iou'] < cond_info['mask_iou']:
+                    print("Auto using Cond Image to bake front view")
+                    aligned_img, info = aligned_cond, cond_info
+                    
                 need_baking = info['mask_iou'] > self.iou_thresh
+                
             else:
-                aligned_img, info, need_baking = self.align_and_check(views[ele//60], rendered_views[ele_idx])
+                aligned_img, info, need_baking = self.align_and_check(views[ele//60], rendered_views[min(ele//60, len(others)-1)])
+                if info is None: continue
                 aligned_img.save(save_folder + f'/aligned_{ele}.jpg')
 
-            if need_baking or (ele in self.force_baking_ele_list):
-                st = time.time()
-                view1_res = back_projection(
-                    obj_file = obj_path,
-                    init_texture_file = raw_texture_path,
-                    front_view_file = aligned_img,
-                    dst_dir = os.path.join(save_folder, f"view_{ele_idx}"),
-                    render_resolution = aligned_img.size[0], 
-                    uv_resolution = 1024,
-                    views = [[0, ele]],
-                    device = self.device
-                )
-                print(f"view_{ele_idx} elevation_{ele} baking finished at {time.time() - st}")
-                obj_path = os.path.join(save_folder, f"view_{ele_idx}/bake/mesh.obj")
-                raw_texture_path = os.path.join(save_folder, f"view_{ele_idx}/bake/texture.png")
-            else:
-                print(f"Skip view_{ele_idx} elevation_{ele} baking")
+            try:
+                if need_baking or force:
+                    st = time.time()
+                    view1_res = back_projection(
+                        obj_file = obj_path,
+                        init_texture_file = raw_texture_path,
+                        front_view_file = aligned_img,
+                        dst_dir = os.path.join(save_folder, f"view_{ele_idx}"),
+                        render_resolution = aligned_img.size[0], 
+                        uv_resolution = 1024,
+                        views = [[0, ele]],
+                        device = self.device
+                    )
+                    print(f"view_{ele_idx} elevation_{ele} baking finished at {time.time() - st}")
+                    obj_path = os.path.join(save_folder, f"view_{ele_idx}/bake/mesh.obj")
+                    raw_texture_path = os.path.join(save_folder, f"view_{ele_idx}/bake/texture.png")
+                else:
+                    print(f"Skip view_{ele_idx} elevation_{ele} baking")
+            except Exception as err:
+                print(err)
+                continue
 
-        print("Baking Finished")
+        print("\nBaking Finished\n")
         return obj_path