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)