init project

app.py

@@ -38,7 +38,8 @@ from modules.pe3r.models import Models
 import torchvision.transforms as tvf
 
 silent = False
-pe3r = Models('cpu') #
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+pe3r = Models(device) #
 
 def _convert_scene_output_to_glb(outdir, imgs, pts3d, mask, focals, cams2world, cam_size=0.05,
                                  cam_color=None, as_pointcloud=False,
@@ -241,14 +242,14 @@ 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(sam1_image, yolov8_image, original_size, input_size, transform):
 
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    # device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
     sam_mask=[]
     img_area = original_size[0] * original_size[1]
 
-    obj_results = yolov8(yolov8_image,device=device,retina_masks=False,imgsz=1024,conf=0.25,iou=0.95,verbose=False)
+    obj_results = pe3r.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)
@@ -262,21 +263,21 @@ def get_mask_from_img_sam1(mobilesamv2, yolov8, sam1_image, yolov8_image, origin
 
     # input_boxes = torch.cat((input_boxes1, input_boxes2), dim=0)
 
-    input_image = mobilesamv2.preprocess(sam1_image)
-    image_embedding = mobilesamv2.image_encoder(input_image)['last_hidden_state']
+    input_image = pe3r.mobilesamv2.preprocess(sam1_image)
+    image_embedding = pe3r.mobilesamv2.image_encoder(input_image)['last_hidden_state']
 
     image_embedding=torch.repeat_interleave(image_embedding, 320, dim=0)
-    prompt_embedding=mobilesamv2.prompt_encoder.get_dense_pe()
+    prompt_embedding=pe3r.mobilesamv2.prompt_encoder.get_dense_pe()
     prompt_embedding=torch.repeat_interleave(prompt_embedding, 320, dim=0)
     for (boxes,) in batch_iterator(320, input_boxes):
         with torch.no_grad():
             image_embedding=image_embedding[0:boxes.shape[0],:,:,:]
             prompt_embedding=prompt_embedding[0:boxes.shape[0],:,:,:]
-            sparse_embeddings, dense_embeddings = mobilesamv2.prompt_encoder(
+            sparse_embeddings, dense_embeddings = pe3r.mobilesamv2.prompt_encoder(
                 points=None,
                 boxes=boxes,
                 masks=None,)
-            low_res_masks, _ = mobilesamv2.mask_decoder(
+            low_res_masks, _ = pe3r.mobilesamv2.mask_decoder(
                 image_embeddings=image_embedding,
                 image_pe=prompt_embedding,
                 sparse_prompt_embeddings=sparse_embeddings,
@@ -284,8 +285,8 @@ def get_mask_from_img_sam1(mobilesamv2, yolov8, sam1_image, yolov8_image, origin
                 multimask_output=False,
                 simple_type=True,
             )
-            low_res_masks=mobilesamv2.postprocess_masks(low_res_masks, input_size, original_size)
-            sam_mask_pre = (low_res_masks > mobilesamv2.mask_threshold)
+            low_res_masks=pe3r.mobilesamv2.postprocess_masks(low_res_masks, input_size, original_size)
+            sam_mask_pre = (low_res_masks > pe3r.mobilesamv2.mask_threshold)
             for mask in sam_mask_pre:
                 if mask.sum() / img_area > 0.002:
                     sam_mask.append(mask.squeeze(1))
@@ -299,7 +300,7 @@ def get_mask_from_img_sam1(mobilesamv2, yolov8, sam1_image, yolov8_image, origin
 @torch.no_grad
 def get_cog_feats(images):
 
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    # device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
     cog_seg_maps = []
     rev_cog_seg_maps = []
@@ -311,7 +312,7 @@ def get_cog_feats(images):
     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(sam1_images[0], np_images[0], np_images_size[0], sam1_images_size[0], images.sam1_transform)
     for mask in sam1_masks:
         _, _, _ = pe3r.sam2.add_new_mask(
             inference_state=inference_state,
@@ -333,7 +334,7 @@ def get_cog_feats(images):
         if out_frame_idx == 0:
             continue
 
-        sam1_masks = get_mask_from_img_sam1(pe3r.mobilesamv2, pe3r.yolov8, sam1_images[out_frame_idx], np_images[out_frame_idx], np_images_size[out_frame_idx], sam1_images_size[out_frame_idx], images.sam1_transform)
+        sam1_masks = get_mask_from_img_sam1(sam1_images[out_frame_idx], np_images[out_frame_idx], np_images_size[out_frame_idx], sam1_images_size[out_frame_idx], images.sam1_transform)
 
         for sam1_mask in sam1_masks:
             flg = 1
@@ -445,7 +446,7 @@ def get_reconstructed_scene(outdir, filelist, schedule, niter, min_conf_thr,
     then run get_3D_model_from_scene
     """
 
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    # device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
     if len(filelist) < 2:
         raise gradio.Error("Please input at least 2 images.")
@@ -510,7 +511,7 @@ def get_reconstructed_scene(outdir, filelist, schedule, niter, min_conf_thr,
 def get_3D_object_from_scene(outdir, text, threshold, scene, min_conf_thr, as_pointcloud, 
                  mask_sky, clean_depth, transparent_cams, cam_size):
     
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    # device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
     texts = [text]
     inputs = pe3r.siglip_tokenizer(text=texts, padding="max_length", return_tensors="pt")