init project

modules/pe3r/demo.py

@@ -244,13 +244,13 @@ def get_mask_from_img_sam1(mobilesamv2, yolov8, sam1_image, yolov8_image, origin
     input_boxes1 = obj_results[0].boxes.xyxy
     input_boxes1 = input_boxes1.cpu().numpy()
     input_boxes1 = transform.apply_boxes(input_boxes1, original_size)
-    input_boxes = torch.from_numpy(input_boxes1).cuda()
+    input_boxes = torch.from_numpy(input_boxes1).to(device)
     
     # obj_results = yolov8(yolov8_image,device=device,retina_masks=False,imgsz=512,conf=0.25,iou=0.9,verbose=False)
     # input_boxes2 = obj_results[0].boxes.xyxy
     # input_boxes2 = input_boxes2.cpu().numpy()
     # input_boxes2 = transform.apply_boxes(input_boxes2, original_size)
-    # input_boxes2 = torch.from_numpy(input_boxes2).cuda()
+    # input_boxes2 = torch.from_numpy(input_boxes2).to(device)
 
     # input_boxes = torch.cat((input_boxes1, input_boxes2), dim=0)
 
@@ -383,7 +383,7 @@ def get_cog_feats(images, pe3r, device):
         seg_imgs = torch.from_numpy(seg_imgs).permute(0,3,1,2) # / 255.0
         
         inputs = pe3r.siglip_processor(images=seg_imgs, return_tensors="pt")
-        inputs = {key: value.to("cuda") for key, value in inputs.items()}
+        inputs = {key: value.to(device) for key, value in inputs.items()}
         
         image_features = pe3r.siglip.get_image_features(**inputs)
         image_features = image_features / image_features.norm(dim=-1, keepdim=True)
@@ -512,12 +512,12 @@ def get_reconstructed_scene(outdir, pe3r, device, silent, filelist, schedule, ni
     return scene, outfile, imgs
 
 @spaces.GPU(duration=180)
-def get_3D_object_from_scene(outdir, pe3r, silent, text, threshold, scene, min_conf_thr, as_pointcloud, 
+def get_3D_object_from_scene(outdir, pe3r, silent, device, text, threshold, scene, min_conf_thr, as_pointcloud, 
                  mask_sky, clean_depth, transparent_cams, cam_size):
     
     texts = [text]
     inputs = pe3r.siglip_tokenizer(text=texts, padding="max_length", return_tensors="pt")
-    inputs = {key: value.to("cuda") for key, value in inputs.items()}
+    inputs = {key: value.to(device) for key, value in inputs.items()}
     with torch.no_grad():
         text_feats =pe3r.siglip.get_text_features(**inputs)
         text_feats = text_feats / text_feats.norm(dim=-1, keepdim=True)
@@ -563,7 +563,7 @@ def main_demo(tmpdirname, server_name, server_port, silent=False):
     
     recon_fun = functools.partial(get_reconstructed_scene, tmpdirname, pe3r, device, silent)
     model_from_scene_fun = functools.partial(get_3D_model_from_scene, tmpdirname, silent)
-    get_3D_object_from_scene_fun = functools.partial(get_3D_object_from_scene, tmpdirname, pe3r, silent)
+    get_3D_object_from_scene_fun = functools.partial(get_3D_object_from_scene, tmpdirname, pe3r, silent, device)
 
     with gradio.Blocks(css=""".gradio-container {margin: 0 !important; min-width: 100%};""", title="PE3R Demo") as demo:
         # scene state is save so that you can change conf_thr, cam_size... without rerunning the inference