Update models/ISOMER/scripts/project_mesh.py

models/ISOMER/scripts/project_mesh.py

@@ -91,7 +91,7 @@ class Pix2FacesRenderer:
         vertices = vertices * torch.tensor([-1, -1, 1]).to(vertices)
         vertices = torch.cat([vertices, torch.ones_like(vertices[..., :1])], dim=-1).to(torch.float32)
         return vertices
-
+    @spaces.GPU
     def render_pix2faces_nvdiff(self, meshes: Meshes, cameras: CamerasBase, H=512, W=512):
         meshes = meshes.to(self.device)
         cameras = cameras.to(self.device)
@@ -110,7 +110,7 @@ def get_visible_faces(meshes: Meshes, cameras: CamerasBase, resolution=1024):
     unique_faces = torch.unique(pix_to_face.flatten())
     unique_faces = unique_faces[unique_faces != -1]
     return unique_faces
-
+@spaces.GPU
 def project_color(meshes: Meshes, cameras: CamerasBase, pil_image: Image.Image, use_alpha=True, eps=0.05, resolution=1024, device="cuda") -> dict:
     """
     Projects color from a given image onto a 3D mesh.
@@ -180,6 +180,7 @@ def project_color(meshes: Meshes, cameras: CamerasBase, pil_image: Image.Image,
         "cos_angles": cos_angles,
     }
 
+@spaces.GPU
 def complete_unseen_vertex_color(meshes: Meshes, valid_index: torch.Tensor) -> dict:
     """
     meshes: the mesh with vertex color to be completed.
@@ -228,6 +229,7 @@ def complete_unseen_vertex_color(meshes: Meshes, valid_index: torch.Tensor) -> d
     meshes.textures = TexturesVertex(verts_features=[colors])
     return meshes
 
+@spaces.GPU
 def load_glb_mesh(glb_path, device="cuda"):
     meshes = load_objs_as_meshes([glb_path], device=device)
     return meshes
@@ -240,6 +242,7 @@ def get_separated_images_from_img_grid(img_grid_path, image_num):
         img_list.append(grid.crop((i*h, 0, i*h + h, h)))
     return img_list
 
+@spaces.GPU
 def get_fov_camera_(azimuth, elevation, fovy, radius, mesh, auto_center, scale_factor, device='cuda'):
     if auto_center:
         verts = mesh.verts_packed()
@@ -255,6 +258,7 @@ def get_fov_camera_(azimuth, elevation, fovy, radius, mesh, auto_center, scale_f
     cameras = FoVPerspectiveCameras(device=device, R=R, T=T, fov=fovy)
     return cameras
 
+@spaces.GPU
 def multiview_color_projection(meshes: Meshes, image_list: List[Image.Image], cameras_list: List[CamerasBase], weights=None, eps=0.05, resolution=1024, device="cuda", reweight_with_cosangle="square", use_alpha=True, confidence_threshold=0.1, complete_unseen=False, below_confidence_strategy="smooth") -> Meshes:
     """
     Projects color from a given image onto a 3D mesh.
@@ -327,6 +331,7 @@ def multiview_color_projection(meshes: Meshes, image_list: List[Image.Image], ca
     del meshes
     return ret_mesh
 
+@spaces.GPU
 def get_cameras_list(azim_list, device, elevation, fov_in_degrees=None, focal=2/1.35, dist=1.1, cam_type='orthographic'):
     ret = []
     for azim in azim_list:
@@ -336,6 +341,7 @@ def get_cameras_list(azim_list, device, elevation, fov_in_degrees=None, focal=2/
         ret.append(cameras)
     return ret
 
+@spaces.GPU
 def get_cameras_list_azi_ele(azim_list, elev_list, device, fov_in_degrees=None, focal=2/1.35, dist=1.1, cam_type='orthographic'):
     ret = []
     for i in range(len(azim_list)):
@@ -345,18 +351,23 @@ def get_cameras_list_azi_ele(azim_list, elev_list, device, fov_in_degrees=None,
         ret.append(cameras)
     return ret
 
+@spaces.GPU
 def get_8view_cameras(device, focal=2/1.35):
     return get_cameras_list(azim_list = [180, 225, 270, 315, 0, 45, 90, 135], elevation=0, device=device, focal=focal)
 
+@spaces.GPU
 def get_6view_cameras(device, focal=2/1.35):
     return get_cameras_list(azim_list = [180, 225, 270, 0, 90, 135], elevation=0, device=device, focal=focal)
 
+@spaces.GPU
 def get_4view_cameras(device, focal=2/1.35):
     return get_cameras_list(azim_list = [180, 270, 0, 90], elevation=0, device=device, focal=focal)
 
+@spaces.GPU
 def get_2view_cameras(device, focal=2/1.35):
     return get_cameras_list(azim_list = [180, 0], elevation=0, device=device, focal=focal)
 
+@spaces.GPU
 def get_multiple_view_cameras(device, focal=2/1.35, offset=180, num_views=8, dist=1.1):
     return get_cameras_list(azim_list = (np.linspace(0, 360, num_views+1)[:-1] + offset) % 360, elevation=0, device=device, focal=focal, dist=dist)