add config

config.json

@@ -6,7 +6,6 @@
     "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

@@ -9,13 +9,13 @@ class SegVolConfig(PretrainedConfig):
     def __init__(
         self,
         test_mode=True,
-        clip_model='openai/clip-vit-base-patch32',
+        # clip_model='.',
         **kwargs,
     ):
         self.spatial_size = [32, 256, 256]
         self.patch_size = [4, 16, 16]
         self.test_mode = test_mode
-        self.clip_model = clip_model
+        # self.clip_model = clip_model
         super().__init__(**kwargs)
 
 class SegVolModel(PreTrainedModel):
@@ -36,7 +36,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,
+            # clip_model=self.config.clip_model,
             test_mode=self.config.test_mode,
             )
         
@@ -118,7 +118,6 @@ class SegVolModel(PreTrainedModel):
         return logits_global_single
 
     def forward_train(self, image, train_organs, train_labels):
-        print('in forward_train')
         loss = self.model(image, text=None, boxes=None, points=None,
                  train_organs=train_organs,
                  train_labels=train_labels)
@@ -318,7 +317,6 @@ def generate_box(pred_pre, bbox_shift=None):
     ones_idx = (meaning_post_label > 0).nonzero(as_tuple=True)
     if all(tensor.nelement() == 0 for tensor in ones_idx):
         bboxes = torch.tensor([-1,-1,-1,-1,-1,-1])
-        # print(bboxes, bboxes.shape)
         return bboxes
     min_coords = [dim.min() for dim in ones_idx]    # [x_min, y_min, z_min]
     max_coords = [dim.max() for dim in ones_idx]    # [x_max, y_max, z_max]
@@ -395,8 +393,6 @@ def select_points(preds, num_positive_extra=4, num_negative_extra=0, fix_extra_p
         extra_negative_points = torch.tensor(extra_negative_points).reshape(-1, 3)
         points = torch.cat((points, extra_negative_points), dim=0)
         labels = torch.cat((labels, torch.zeros((extra_negative_points.shape[0]))))
-        # print('extra_negative_points ', extra_negative_points, extra_negative_points.shape)
-        # print('==> points ', points.shape, labels)
     
     if fix_extra_point_num is None:
         left_point_num = num_positive_extra + num_negative_extra + 1 - labels.shape[0]
@@ -415,7 +411,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
-from transformers import AutoTokenizer, CLIPTextModel, CLIPTextConfig
+from transformers import CLIPTextModel, CLIPTextConfig
 import random
 
 #%% set up model
@@ -426,7 +422,7 @@ class SegVol(nn.Module):
                 prompt_encoder,
                 roi_size,
                 patch_size,
-                clip_model,
+                # clip_model,
                 test_mode=False,
                 ):
         super().__init__()
@@ -434,7 +430,7 @@ class SegVol(nn.Module):
         self.image_encoder = image_encoder
         self.mask_decoder = mask_decoder
         self.prompt_encoder = prompt_encoder
-        self.text_encoder = TextEncoder(clip_model)
+        self.text_encoder = TextEncoder()
         self.feat_shape = np.array(roi_size)/np.array(patch_size)
         self.test_mode = test_mode
         self.dice_loss = BinaryDiceLoss().to(self.custom_device)
@@ -453,7 +449,6 @@ class SegVol(nn.Module):
         
         # train mode
         ## sl
-        print('supervised_forward ready')
         sl_loss = self.supervised_forward(image, image_embedding, img_shape, kwargs['train_organs'], kwargs['train_labels'])
         ## ssl
         # ssl_loss = self.unsupervised_forward(image, image_embedding, kwargs['pseudo_seg_cleaned'], img_shape)
@@ -612,12 +607,12 @@ class SegVol(nn.Module):
     #     return pseudo_labels, bboxes
     
 class TextEncoder(nn.Module):
-    def __init__(self, clip_model):
+    def __init__(self):
         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(clip_model)
+        self.tokenizer = None
         self.dim_align = nn.Linear(512, 768)
         # freeze text encoder
         for param in self.clip_text_model.parameters():

special_tokens_map.json

@@ -0,0 +1 @@
+{"bos_token": {"content": "<|startoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "eos_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "unk_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "pad_token": "<|endoftext|>"}

tokenizer_config.json

@@ -0,0 +1 @@
+{"unk_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "bos_token": {"content": "<|startoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "eos_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "pad_token": "<|endoftext|>", "add_prefix_space": false, "errors": "replace", "do_lower_case": true, "name_or_path": "/home/yuxin/BAAI/code_release/segvol_transformers/config/clip", "special_tokens_map_file": "/home/yuxin/BAAI/code_release/segvol_transformers/config/clip/special_tokens_map.json", "tokenizer_class": "CLIPTokenizer"}