add model

config.json

@@ -0,0 +1,23 @@
+{
+  "architectures": [
+    "SegVolModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "config_segvol.SegVolConfig",
+    "AutoModel": "model_segvol.SegVolModel"
+  },
+  "model_type": "segvol",
+  "patch_size": [
+    4,
+    16,
+    16
+  ],
+  "spatial_size": [
+    32,
+    256,
+    256
+  ],
+  "test_mode": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.18.0"
+}

config_segvol.py

@@ -0,0 +1,13 @@
+from transformers import PretrainedConfig
+
+class SegVolConfig(PretrainedConfig):
+    model_type = "segvol"
+
+    def __init__(
+        self,
+        **kwargs,
+    ):
+        self.spatial_size = [32, 256, 256]
+        self.patch_size = [4, 16, 16]
+        self.test_mode = True
+        super().__init__(**kwargs)

model_segvol.py

@@ -0,0 +1,24 @@
+from transformers import PreTrainedModel
+from config_segvol import SegVolConfig
+from network.model import SegVol
+from segment_anything_volumetric import sam_model_registry
+
+
+class SegVolModel(PreTrainedModel):
+    config_class = SegVolConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        sam_model = sam_model_registry['vit'](self.config.patch_size, self.config.spatial_size)
+        self.model = SegVol(
+            image_encoder=sam_model.image_encoder, 
+            mask_decoder=sam_model.mask_decoder,
+            prompt_encoder=sam_model.prompt_encoder,
+            roi_size=self.config.spatial_size,
+            patch_size=self.config.patch_size,
+            test_mode=self.config.test_mode,
+            )
+
+    def forward(self, image, text=None, boxes=None, points=None, **kwargs):
+        return self.model.forward(image, text=text, boxes=boxes, points=points, **kwargs)
+

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:500f2758a8f989339b2b2baf09a819169bc87549795193d3cfe505726ac0b399
+size 723726667