File size: 1,201 Bytes
3818098 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
from transformers import PreTrainedModel
from .w_net_3d import WNet3dUNet, WNet3dAttUNet, WNet3dUNetMSS
from .WNetConfigs import WNet3DConfig, AttWNet3DConfig, WNetMSS3DConfig
class WNet3D(PreTrainedModel):
config_class = WNet3DConfig
def __init__(self, config):
super().__init__(config)
self.model = WNet3dUNet(
in_ch=config.in_ch,
out_ch=config.out_ch,
init_features=config.init_features)
def forward(self, x):
return self.model(x)
class AttWNet3D(PreTrainedModel):
config_class = AttWNet3DConfig
def __init__(self, config):
super().__init__(config)
self.model = WNet3dAttUNet(
in_ch=config.in_ch,
out_ch=config.out_ch,
init_features=config.init_features)
def forward(self, x):
return self.model(x)
class WNetMSS3D(PreTrainedModel):
config_class = WNetMSS3DConfig
def __init__(self, config):
super().__init__(config)
self.model = WNet3dUNetMSS(
in_ch=config.in_ch,
out_ch=config.out_ch,
init_features=config.init_features)
def forward(self, x):
return self.model(x) |