from transformers import PreTrainedModel
from .config import MonoSceneConfig
from monoscene.monoscene import MonoScene


class MonoSceneModel(PreTrainedModel):
    config_class = MonoSceneConfig

    def __init__(self, config):
        super().__init__(config)
        self.model = MonoScene(
            dataset=config.dataset,
            n_classes=config.n_classes,
            feature=config.feature,
            project_res=['1', '2', '4', '8'],
            project_scale=config.project_scale,
            full_scene_size=config.full_scene_size
        )
     

    def forward(self, tensor):
        return self.model.forward(tensor)