init project

modules/pe3r/models.py

@@ -2,7 +2,7 @@ import os
 import sys
 sys.path.append(os.path.abspath('./modules/ultralytics'))
 
-from transformers import AutoTokenizer, AutoModel, AutoProcessor
+from transformers import AutoTokenizer, AutoModel, AutoProcessor, SamModel
 from modules.mast3r.model import AsymmetricMASt3R
 
 # from modules.sam2.build_sam import build_sam2_video_predictor
@@ -32,10 +32,13 @@ class Models:
 
         # -- mobilesamv2 & sam1 --
         # SAM1_ENCODER_CKP = './checkpoints/sam_vit_h.pt'
-        SAM1_ENCODER_CKP = 'facebook/sam-vit-huge/model.safetensors'
+        # SAM1_ENCODER_CKP = 'facebook/sam-vit-huge/model.safetensors'
         SAM1_DECODER_CKP = './checkpoints/Prompt_guided_Mask_Decoder.pt'
         self.mobilesamv2 = sam_model_registry['sam_vit_h'](None)
-        image_encoder=sam_model_registry['sam_vit_h_encoder'](SAM1_ENCODER_CKP)
+        # image_encoder=sam_model_registry['sam_vit_h_encoder'](SAM1_ENCODER_CKP)
+        sam1 = SamModel.from_pretrained("facebook/sam-vit-huge", device=device)
+        image_encoder = sam1.image_encoder
+
         prompt_encoder, mask_decoder = sam_model_registry['prompt_guided_decoder'](SAM1_DECODER_CKP)
         self.mobilesamv2.prompt_encoder = prompt_encoder
         self.mobilesamv2.mask_decoder = mask_decoder