init project

app.py

@@ -34,14 +34,26 @@ from typing import Any, Dict, Generator,List
 import matplotlib.pyplot as pl
 
 from modules.mobilesamv2.utils.transforms import ResizeLongestSide
-from modules.pe3r.models import Models
+# from modules.pe3r.models import Models
 import torchvision.transforms as tvf
 
-from transformers import AutoTokenizer, AutoModel, AutoProcessor
+
+
+sys.path.append(os.path.abspath('./modules/ultralytics'))
+
+from transformers import AutoTokenizer, AutoModel, AutoProcessor, SamModel
+from modules.mast3r.model import AsymmetricMASt3R
+
+# from modules.sam2.build_sam import build_sam2_video_predictor
+from modules.mobilesamv2.promt_mobilesamv2 import ObjectAwareModel
+from modules.mobilesamv2 import sam_model_registry
+
+from sam2.sam2_video_predictor import SAM2VideoPredictor
+
 
 silent = False
 # device = 'cuda' if torch.cuda.is_available() else 'cpu' #'cpu' #
-pe3r = Models('cpu') #
+# pe3r = Models('cpu') #
 # print(device)
 
 def _convert_scene_output_to_glb(outdir, imgs, pts3d, mask, focals, cams2world, cam_size=0.05,
@@ -248,13 +260,25 @@ def slerp_multiple(vectors, t_values):
 def get_mask_from_img_sam1(sam1_image, yolov8_image, original_size, input_size, transform):
 
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    pe3r.yolov8.to(device)
-    pe3r.mobilesamv2.to(device)
+    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]
 
-    obj_results = pe3r.yolov8(yolov8_image,device=device,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)
@@ -268,21 +292,21 @@ def get_mask_from_img_sam1(sam1_image, yolov8_image, original_size, input_size,
 
     # input_boxes = torch.cat((input_boxes1, input_boxes2), dim=0)
 
-    input_image = pe3r.mobilesamv2.preprocess(sam1_image)
-    image_embedding = pe3r.mobilesamv2.image_encoder(input_image)['last_hidden_state']
+    input_image = mobilesamv2.preprocess(sam1_image)
+    image_embedding = mobilesamv2.image_encoder(input_image)['last_hidden_state']
 
     image_embedding=torch.repeat_interleave(image_embedding, 320, dim=0)
-    prompt_embedding=pe3r.mobilesamv2.prompt_encoder.get_dense_pe()
+    prompt_embedding=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 = pe3r.mobilesamv2.prompt_encoder(
+            sparse_embeddings, dense_embeddings = mobilesamv2.prompt_encoder(
                 points=None,
                 boxes=boxes,
                 masks=None,)
-            low_res_masks, _ = pe3r.mobilesamv2.mask_decoder(
+            low_res_masks, _ = mobilesamv2.mask_decoder(
                 image_embeddings=image_embedding,
                 image_pe=prompt_embedding,
                 sparse_prompt_embeddings=sparse_embeddings,
@@ -290,8 +314,8 @@ def get_mask_from_img_sam1(sam1_image, yolov8_image, original_size, input_size,
                 multimask_output=False,
                 simple_type=True,
             )
-            low_res_masks=pe3r.mobilesamv2.postprocess_masks(low_res_masks, input_size, original_size)
-            sam_mask_pre = (low_res_masks > pe3r.mobilesamv2.mask_threshold)
+            low_res_masks=mobilesamv2.postprocess_masks(low_res_masks, input_size, original_size)
+            sam_mask_pre = (low_res_masks > mobilesamv2.mask_threshold)
             for mask in sam_mask_pre:
                 if mask.sum() / img_area > 0.002:
                     sam_mask.append(mask.squeeze(1))
@@ -306,13 +330,15 @@ def get_mask_from_img_sam1(sam1_image, yolov8_image, original_size, input_size,
 def get_cog_feats(images):
 
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    pe3r.sam2.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")
 
     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])
+    inference_state = sam2.init_state(images=images.sam2_images, video_height=images.sam2_video_size[0], video_width=images.sam2_video_size[1])
     mask_num = 0
 
     sam1_images = images.sam1_images
@@ -322,7 +348,7 @@ def get_cog_feats(images):
     
     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(
+        _, _, _ = sam2.add_new_mask(
             inference_state=inference_state,
             frame_idx=0,
             obj_id=mask_num,
@@ -331,7 +357,7 @@ def get_cog_feats(images):
         mask_num += 1
 
     video_segments = {}  # video_segments contains the per-frame segmentation results
-    for out_frame_idx, out_obj_ids, out_mask_logits in pe3r.sam2.propagate_in_video(inference_state):
+    for out_frame_idx, out_obj_ids, out_mask_logits in sam2.propagate_in_video(inference_state):
         sam2_masks = (out_mask_logits > 0.0).squeeze(1)
 
         video_segments[out_frame_idx] = {
@@ -455,8 +481,8 @@ def get_reconstructed_scene(outdir, filelist, schedule, niter, min_conf_thr,
     """
 
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
-
-    pe3r.mast3r.to(device)
+    MAST3R_CKP = 'naver/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric'
+    mast3r = AsymmetricMASt3R.from_pretrained(MAST3R_CKP).to(device)
 
     if len(filelist) < 2:
         raise gradio.Error("Please input at least 2 images.")
@@ -485,7 +511,7 @@ def get_reconstructed_scene(outdir, filelist, schedule, niter, min_conf_thr,
         scenegraph_type = scenegraph_type + "-" + str(refid)
 
     pairs = make_pairs(imgs, scene_graph=scenegraph_type, prefilter=None, symmetrize=True)
-    output = inference(pairs, pe3r.mast3r, device, batch_size=1, verbose=not silent)
+    output = inference(pairs, mast3r, device, batch_size=1, verbose=not silent)
     mode = GlobalAlignerMode.PointCloudOptimizer if len(imgs) > 2 else GlobalAlignerMode.PairViewer
     scene_1 = global_aligner(output, cog_seg_maps, rev_cog_seg_maps, cog_feats, device=device, mode=mode, verbose=not silent)
     lr = 0.01
@@ -498,7 +524,7 @@ def get_reconstructed_scene(outdir, filelist, schedule, niter, min_conf_thr,
             # print(imgs[i]['img'].shape, scene.imgs[i].shape, ImgNorm(scene.imgs[i])[None])
             imgs[i]['img'] = ImgNorm(scene_1.imgs[i])[None]
         pairs = make_pairs(imgs, scene_graph=scenegraph_type, prefilter=None, symmetrize=True)
-        output = inference(pairs, pe3r.mast3r, device, batch_size=1, verbose=not silent)
+        output = inference(pairs, mast3r, device, batch_size=1, verbose=not silent)
         mode = GlobalAlignerMode.PointCloudOptimizer if len(imgs) > 2 else GlobalAlignerMode.PairViewer
         scene = global_aligner(output, cog_seg_maps, rev_cog_seg_maps, cog_feats, device=device, mode=mode, verbose=not silent)
         ori_imgs = scene.ori_imgs