Update app.py

app.py

@@ -17,10 +17,10 @@ import trimesh
 from moge.model.v1 import MoGeModel
 from utils.geometry import compute_pointmap
 import cv2
-# from huggingface_hub import hf_hub_download
-# from PIL import Image
-# import matplotlib.pyplot as plt
-# from eval_wrapper.eval import EvalWrapper, eval_scene
+from huggingface_hub import hf_hub_download
+from PIL import Image
+import matplotlib.pyplot as plt
+from eval_wrapper.eval import EvalWrapper, eval_scene
 
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -28,21 +28,27 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 outdir = "/tmp/rayst3r"
 
 moge_model = MoGeModel.from_pretrained("Ruicheng/moge-vitl").to(device)
+dino_model = torch.hub.load('facebookresearch/dinov2', "dinov2_vitl14_reg")
+dino_model.eval()
+dino_model.to(device)
 
-
-# # loading all necessary models
+print("Loading RaySt3R model")
+rayst3r_checkpoint = hf_hub_download("bartduis/rayst3r", "rayst3r.pth")
+rayst3r_model = EvalWrapper(rayst3r_checkpoint,device='cpu')
+rayst3r_model = rayst3r_model.to(device)
+print("Loaded all models")
 
 # print("Loading MoGe model")
 # # Load the model from huggingface hub (or load from local).
 
 
-# def depth2uint16(depth):
-#     return depth * torch.iinfo(torch.uint16).max / 10.0 # threshold is in m, convert to uint16 value
+def depth2uint16(depth):
+    return depth * torch.iinfo(torch.uint16).max / 10.0 # threshold is in m, convert to uint16 value
 
-# def save_tensor_as_png(tensor: torch.Tensor, path: str, dtype: torch.dtype | None = None):
-#     if dtype is None:
-#         dtype = tensor.dtype
-#     Image.fromarray(tensor.to(dtype).cpu().numpy()).save(path)
+def save_tensor_as_png(tensor: torch.Tensor, path: str, dtype: torch.dtype | None = None):
+    if dtype is None:
+        dtype = tensor.dtype
+    Image.fromarray(tensor.to(dtype).cpu().numpy()).save(path)
 
 # def colorize_points_with_turbo_all_dims(points, method='norm',cmap='turbo'):
 #     """
@@ -78,62 +84,56 @@ moge_model = MoGeModel.from_pretrained("Ruicheng/moge-vitl").to(device)
 
 #     return colors
 
-# def prep_for_rayst3r(img,depth_dict,mask):
-#     H, W = img.shape[:2]
-#     intrinsics = depth_dict["intrinsics"].detach().cpu()
-#     intrinsics[0] *= W
-#     intrinsics[1] *= H
-
-#     input_dir = os.path.join(outdir, "input")
-#     if os.path.exists(input_dir):
-#         shutil.rmtree(input_dir)
-#     os.makedirs(input_dir, exist_ok=True)
-#     # save intrinsics
-#     torch.save(intrinsics, os.path.join(input_dir, "intrinsics.pt"))
-
-#     # save depth
-#     depth = depth_dict["depth"].cpu()
-#     depth = depth2uint16(depth)
-#     save_tensor_as_png(depth, os.path.join(input_dir, "depth.png"),dtype=torch.uint16)
-
-#     # save mask as bool 
-#     save_tensor_as_png(torch.from_numpy(mask).bool(), os.path.join(input_dir, "mask.png"),dtype=torch.bool)
-#     # save image
-#     save_tensor_as_png(torch.from_numpy(img), os.path.join(input_dir, "rgb.png"))
-
-# @GPU(duration = 180)
-# def rayst3r_to_glb(img,depth_dict,mask,max_total_points=10e6,rotated=False):
-#     prep_for_rayst3r(img,depth_dict,mask)
+def prep_for_rayst3r(img,depth_dict,mask):
+    H, W = img.shape[:2]
+    intrinsics = depth_dict["intrinsics"].detach().cpu()
+    intrinsics[0] *= W
+    intrinsics[1] *= H
+
+    input_dir = os.path.join(outdir, "input")
+    if os.path.exists(input_dir):
+        shutil.rmtree(input_dir)
+    os.makedirs(input_dir, exist_ok=True)
+    # save intrinsics
+    torch.save(intrinsics, os.path.join(input_dir, "intrinsics.pt"))
+
+    # save depth
+    depth = depth_dict["depth"].cpu()
+    depth = depth2uint16(depth)
+    save_tensor_as_png(depth, os.path.join(input_dir, "depth.png"),dtype=torch.uint16)
+
+    # save mask as bool 
+    save_tensor_as_png(torch.from_numpy(mask).bool(), os.path.join(input_dir, "mask.png"),dtype=torch.bool)
+    # save image
+    save_tensor_as_png(torch.from_numpy(img), os.path.join(input_dir, "rgb.png"))
+
+@GPU(duration = 180)
+def rayst3r_to_glb(img,depth_dict,mask,max_total_points=10e6,rotated=False):
+    prep_for_rayst3r(img,depth_dict,mask)
+    print('Doneneee')
     
-#     dino_model = torch.hub.load('facebookresearch/dinov2', "dinov2_vitl14_reg")
-#     dino_model.eval()
-#     dino_model.to(device)
     
-#     print("Loading RaySt3R model")
-#     rayst3r_checkpoint = hf_hub_download("bartduis/rayst3r", "rayst3r.pth")
-#     rayst3r_model = EvalWrapper(rayst3r_checkpoint,device='cpu')
-#     rayst3r_model = rayst3r_model.to(device)
     
-#     rayst3r_points = eval_scene(rayst3r_model,os.path.join(outdir, "input"),do_filter_all_masks=True,dino_model=dino_model, device = device).cpu()
+    # rayst3r_points = eval_scene(rayst3r_model,os.path.join(outdir, "input"),do_filter_all_masks=True,dino_model=dino_model, device = device).cpu()
 
-#     # subsample points
-#     n_points = min(max_total_points,rayst3r_points.shape[0])
-#     rayst3r_points = rayst3r_points[torch.randperm(rayst3r_points.shape[0])[:n_points]].numpy()
+    # # subsample points
+    # n_points = min(max_total_points,rayst3r_points.shape[0])
+    # rayst3r_points = rayst3r_points[torch.randperm(rayst3r_points.shape[0])[:n_points]].numpy()
     
-#     rayst3r_points[:,1] = -rayst3r_points[:,1]
-#     rayst3r_points[:,2] = -rayst3r_points[:,2]
+    # rayst3r_points[:,1] = -rayst3r_points[:,1]
+    # rayst3r_points[:,2] = -rayst3r_points[:,2]
     
-#     # make all points red
-#     colors = colorize_points_with_turbo_all_dims(rayst3r_points)
+    # # make all points red
+    # colors = colorize_points_with_turbo_all_dims(rayst3r_points)
 
-#     # load the input glb
-#     scene = trimesh.Scene()
-#     pct = trimesh.PointCloud(rayst3r_points, colors=colors, radius=0.01)
-#     scene.add_geometry(pct)
+    # # load the input glb
+    # scene = trimesh.Scene()
+    # pct = trimesh.PointCloud(rayst3r_points, colors=colors, radius=0.01)
+    # scene.add_geometry(pct)
     
-#     outfile = os.path.join(outdir, "rayst3r.glb")
-#     scene.export(outfile)
-#     return outfile
+    # outfile = os.path.join(outdir, "rayst3r.glb")
+    # scene.export(outfile)
+    # return outfile
 
 
 def input_to_glb(outdir,img,depth_dict,mask,rotated=False):
@@ -202,6 +202,7 @@ def process_image(input_img):
         shutil.rmtree(outdir)
     os.makedirs(outdir)
     input_glb = input_to_glb(outdir,input_img,depth_dict,mask,rotated=rotated)
+    inference_glb = rayst3r_to_glb(input_img,depth_dict,mask,rotated=rotated)
     print(input_glb)
     return input_img, input_img