init project

app.py

@@ -257,23 +257,10 @@ def slerp_multiple(vectors, t_values):
     return interpolated_vector
 
 # @torch.no_grad
-def get_mask_from_img_sam1(sam1_image, yolov8_image, original_size, input_size, transform):
+def get_mask_from_img_sam1(yolov8, mobilesamv2, sam1_image, yolov8_image, original_size, input_size, transform):
 
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    SAM1_DECODER_CKP = './checkpoints/Prompt_guided_Mask_Decoder.pt'
-    mobilesamv2 = sam_model_registry['sam_vit_h'](None)
-    sam1 = SamModel.from_pretrained('facebook/sam-vit-huge')
-    image_encoder = sam1.vision_encoder
-
-    prompt_encoder, mask_decoder = sam_model_registry['prompt_guided_decoder'](SAM1_DECODER_CKP)
-    mobilesamv2.prompt_encoder = prompt_encoder
-    mobilesamv2.mask_decoder = mask_decoder
-    mobilesamv2.image_encoder=image_encoder
-    mobilesamv2.to(device=device)
-    mobilesamv2.eval()
     
-    YOLO8_CKP='./checkpoints/ObjectAwareModel.pt'
-    yolov8 = ObjectAwareModel(YOLO8_CKP)
 
     sam_mask=[]
     img_area = original_size[0] * original_size[1]
@@ -327,15 +314,10 @@ def get_mask_from_img_sam1(sam1_image, yolov8_image, original_size, input_size,
     return ret_mask
 
 # @torch.no_grad
-def get_cog_feats(images):
+def get_cog_feats(images, sam2, siglip, siglip_processor, yolov8, mobilesamv2):
 
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-    sam2 = SAM2VideoPredictor.from_pretrained('facebook/sam2.1-hiera-large', device=device)
-
-    siglip = AutoModel.from_pretrained("google/siglip-large-patch16-256", device_map=device)
-    siglip_processor = AutoProcessor.from_pretrained("google/siglip-large-patch16-256")
-
     cog_seg_maps = []
     rev_cog_seg_maps = []
     inference_state = sam2.init_state(images=images.sam2_images, video_height=images.sam2_video_size[0], video_width=images.sam2_video_size[1])
@@ -346,7 +328,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(sam1_images[0], np_images[0], np_images_size[0], sam1_images_size[0], images.sam1_transform)
+    sam1_masks = get_mask_from_img_sam1(yolov8, mobilesamv2, sam1_images[0], np_images[0], np_images_size[0], sam1_images_size[0], images.sam1_transform)
     for mask in sam1_masks:
         _, _, _ = sam2.add_new_mask(
             inference_state=inference_state,
@@ -368,7 +350,7 @@ def get_cog_feats(images):
         if out_frame_idx == 0:
             continue
 
-        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)
+        sam1_masks = get_mask_from_img_sam1(yolov8, mobilesamv2, 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
@@ -484,13 +466,33 @@ def get_reconstructed_scene(outdir, filelist, schedule, niter, min_conf_thr,
     MAST3R_CKP = 'naver/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric'
     mast3r = AsymmetricMASt3R.from_pretrained(MAST3R_CKP).to(device)
 
+    sam2 = SAM2VideoPredictor.from_pretrained('facebook/sam2.1-hiera-large', device=device)
+
+    siglip = AutoModel.from_pretrained("google/siglip-large-patch16-256", device_map=device)
+    siglip_processor = AutoProcessor.from_pretrained("google/siglip-large-patch16-256")
+
+    SAM1_DECODER_CKP = './checkpoints/Prompt_guided_Mask_Decoder.pt'
+    mobilesamv2 = sam_model_registry['sam_vit_h'](None)
+    sam1 = SamModel.from_pretrained('facebook/sam-vit-huge')
+    image_encoder = sam1.vision_encoder
+
+    prompt_encoder, mask_decoder = sam_model_registry['prompt_guided_decoder'](SAM1_DECODER_CKP)
+    mobilesamv2.prompt_encoder = prompt_encoder
+    mobilesamv2.mask_decoder = mask_decoder
+    mobilesamv2.image_encoder=image_encoder
+    mobilesamv2.to(device=device)
+    mobilesamv2.eval()
+    
+    YOLO8_CKP='./checkpoints/ObjectAwareModel.pt'
+    yolov8 = ObjectAwareModel(YOLO8_CKP)
+
     if len(filelist) < 2:
         raise gradio.Error("Please input at least 2 images.")
 
     images = Images(filelist=filelist, device=device)
     
     # try:
-    cog_seg_maps, rev_cog_seg_maps, cog_feats = get_cog_feats(images)
+    cog_seg_maps, rev_cog_seg_maps, cog_feats = get_cog_feats(images, sam2, siglip, siglip_processor, yolov8, mobilesamv2)
     imgs = load_images(images, rev_cog_seg_maps, size=512, verbose=not silent)
     # except Exception as e:
     #     rev_cog_seg_maps = []