| from transformers import PretrainedConfig | |
| from typing import List | |
| class RRDBConfiguration(PretrainedConfig): | |
| model_type = "SRMRIModelRRDB" | |
| def __init__( | |
| self, | |
| nChannels=1, | |
| nDenseLayers=6, | |
| nInitFeat=6, | |
| GrowthRate=12, | |
| featureFusion=True, | |
| kernel_config=[3, 3, 3, 3], | |
| **kwargs): | |
| self.nChannels = nChannels | |
| self.nDenseLayers = nDenseLayers | |
| self.nInitFeat = nInitFeat | |
| self.GrowthRate = GrowthRate | |
| self.featureFusion = featureFusion | |
| self.kernel_config = kernel_config | |
| super().__init__(**kwargs) | |
| class UNetMSSConfiguration(PretrainedConfig): | |
| model_type = "SRMRIModelUNetMSS" | |
| def __init__( | |
| self, | |
| in_channels=1, | |
| n_classes=1, | |
| depth=3, | |
| wf=6, | |
| padding=True, | |
| batch_norm=False, | |
| up_mode='upconv', | |
| dropout=False, | |
| mss_level=2, | |
| mss_fromlatent=True, | |
| mss_up="trilinear", | |
| mss_interpb4=True, | |
| **kwargs): | |
| self.in_channels = in_channels | |
| self.n_classes = n_classes | |
| self.depth = depth | |
| self.wf = wf | |
| self.padding = padding | |
| self.batch_norm = batch_norm | |
| self.up_mode = up_mode | |
| self.dropout = dropout | |
| self.mss_level = mss_level | |
| self.mss_fromlatent = mss_fromlatent | |
| self.mss_up = mss_up | |
| self.mss_interpb4 = mss_interpb4 | |
| super().__init__(**kwargs) |