from transformers import PreTrainedModel
from .unet3D import UNet
from .UNetFTSRConfig import UNetFTSRConfig

class UNetFTSR(PreTrainedModel):
    config_class = UNetFTSRConfig
    def __init__(self, config):
        super().__init__(config)
        self.model = UNet(
            in_channels=config.in_channels, 
            n_classes=config.n_classes, 
            depth=config.depth, 
            wf=config.wf,
            padding=config.padding,
            batch_norm=config.batch_norm, 
            up_mode=config.up_mode, 
            dropout=config.dropout)
    def forward(self, x):
        return self.model(x)