add model

config.json

@@ -6,6 +6,7 @@
     "AutoConfig": "model_segvol_single.SegVolConfig",
     "AutoModel": "model_segvol_single.SegVolModel"
   },
+  "clip_model": "openai/clip-vit-base-patch32",
   "model_type": "segvol",
   "patch_size": [
     4,

model_segvol_single.py

@@ -8,11 +8,13 @@ class SegVolConfig(PretrainedConfig):
     def __init__(
         self,
         test_mode=True,
+        clip_model='openai/clip-vit-base-patch32',
         **kwargs,
     ):
         self.spatial_size = [32, 256, 256]
         self.patch_size = [4, 16, 16]
         self.test_mode = test_mode
+        self.clip_model = clip_model
         super().__init__(**kwargs)
 
 class SegVolModel(PreTrainedModel):
@@ -33,6 +35,7 @@ class SegVolModel(PreTrainedModel):
             prompt_encoder=sam_model.prompt_encoder,
             roi_size=self.config.spatial_size,
             patch_size=self.config.patch_size,
+            clip_model=self.config.clip_model,
             test_mode=self.config.test_mode,
             )
         
@@ -48,7 +51,7 @@ class SegVolModel(PreTrainedModel):
                 point_prompt_group=None, 
                 use_zoom=True):
         assert image.shape[0] == 1 and zoomed_image.shape[0] == 1, 'batch size should be 1'
-        assert not (text_prompt is None and bbox_prompt_group is None and point_prompt_group is None)
+        assert not (text_prompt is None and bbox_prompt_group is None and point_prompt_group is None), 'Drive SegVol using at least one type of prompt'
         bbox_prompt, bbox_prompt_map, point_prompt, point_prompt_map=None, None, None, None
         if bbox_prompt_group is not None:
             bbox_prompt, bbox_prompt_map = bbox_prompt_group
@@ -363,6 +366,7 @@ class SegVol(nn.Module):
                 prompt_encoder,
                 roi_size,
                 patch_size,
+                clip_model,
                 test_mode=False,
                 ):
         super().__init__()
@@ -370,7 +374,7 @@ class SegVol(nn.Module):
         self.image_encoder = image_encoder
         self.mask_decoder = mask_decoder
         self.prompt_encoder = prompt_encoder
-        self.text_encoder = TextEncoder()
+        self.text_encoder = TextEncoder(clip_model)
         self.feat_shape = np.array(roi_size)/np.array(patch_size)
         self.test_mode = test_mode
         self.dice_loss = BinaryDiceLoss().to(self.custom_device)
@@ -547,12 +551,12 @@ class SegVol(nn.Module):
         return pseudo_labels, bboxes
     
 class TextEncoder(nn.Module):
-    def __init__(self):
+    def __init__(self, clip_model):
         super().__init__()
         self.custom_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         config = CLIPTextConfig()
         self.clip_text_model = CLIPTextModel(config)
-        self.tokenizer = AutoTokenizer.from_pretrained('openai/clip-vit-base-patch32')
+        self.tokenizer = AutoTokenizer.from_pretrained(clip_model)
         self.dim_align = nn.Linear(512, 768)
         # freeze text encoder
         for param in self.clip_text_model.parameters():