init project

modules/pe3r/demo.py

@@ -236,17 +236,17 @@ def slerp_multiple(vectors, t_values):
     return interpolated_vector
 
 @torch.no_grad
-def get_mask_from_img_sam1(mobilesamv2, yolov8, sam1_image, yolov8_image, original_size, input_size, transform):
+def get_mask_from_img_sam1(mobilesamv2, yolov8, sam1_image, yolov8_image, original_size, input_size, transform, device):
     sam_mask=[]
     img_area = original_size[0] * original_size[1]
 
-    obj_results = yolov8(yolov8_image,device='cuda',retina_masks=False,imgsz=1024,conf=0.25,iou=0.95,verbose=False)
+    obj_results = yolov8(yolov8_image,device=device,retina_masks=False,imgsz=1024,conf=0.25,iou=0.95,verbose=False)
     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()
     
-    # obj_results = yolov8(yolov8_image,device='cuda',retina_masks=False,imgsz=512,conf=0.25,iou=0.9,verbose=False)
+    # 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)
@@ -289,7 +289,7 @@ def get_mask_from_img_sam1(mobilesamv2, yolov8, sam1_image, yolov8_image, origin
     return ret_mask
 
 @torch.no_grad
-def get_cog_feats(images, pe3r):
+def get_cog_feats(images, pe3r, device):
     cog_seg_maps = []
     rev_cog_seg_maps = []
     inference_state = pe3r.sam2.init_state(images=images.sam2_images, video_height=images.sam2_video_size[0], video_width=images.sam2_video_size[1])
@@ -300,7 +300,7 @@ def get_cog_feats(images, pe3r):
     np_images = images.np_images
     np_images_size = images.np_images_size
     
-    sam1_masks = get_mask_from_img_sam1(pe3r.mobilesamv2, pe3r.yolov8, sam1_images[0], np_images[0], np_images_size[0], sam1_images_size[0], images.sam1_transform)
+    sam1_masks = get_mask_from_img_sam1(pe3r.mobilesamv2, pe3r.yolov8, sam1_images[0], np_images[0], np_images_size[0], sam1_images_size[0], images.sam1_transform, device)
     for mask in sam1_masks:
         _, _, _ = pe3r.sam2.add_new_mask(
             inference_state=inference_state,
@@ -438,7 +438,7 @@ def get_reconstructed_scene(outdir, pe3r, device, silent, filelist, schedule, ni
     images = Images(filelist=filelist, device=device)
     
     try:
-        cog_seg_maps, rev_cog_seg_maps, cog_feats = get_cog_feats(images, pe3r)
+        cog_seg_maps, rev_cog_seg_maps, cog_feats = get_cog_feats(images, pe3r, device)
         imgs = load_images(images, rev_cog_seg_maps, size=512, verbose=not silent)
     except Exception as e:
         rev_cog_seg_maps = []