| from transformers import PretrainedConfig | |
| class CondUNetConfig(PretrainedConfig): | |
| model_type = "cond_unet" | |
| def __init__( | |
| self, | |
| in_channels=3, | |
| num_labels=1, | |
| n_organs=10, | |
| size=32, | |
| depth=5, | |
| attn_start=0, | |
| use_attn=True, | |
| image_size=512, | |
| patch_size=8, | |
| emb_dim=768, | |
| n_heads=8, | |
| use_dwt=False, | |
| wavelet="haar", | |
| dwt_bands=None, | |
| use_shape=False, | |
| shape_res=32, | |
| unknown_organ_id=-1, | |
| keep_aspect_ratio=True, | |
| self_normalize=True, | |
| **kwargs, | |
| ): | |
| super().__init__(**kwargs) | |
| self.in_channels = in_channels | |
| self.num_labels = num_labels | |
| self.n_organs = n_organs | |
| self.size = size | |
| self.depth = depth | |
| self.attn_start = attn_start | |
| self.use_attn = use_attn | |
| self.image_size = image_size | |
| self.patch_size = patch_size | |
| self.emb_dim = emb_dim | |
| self.n_heads = n_heads | |
| self.use_dwt = use_dwt | |
| self.wavelet = wavelet | |
| self.dwt_bands = dwt_bands or ["LL", "LH", "HL", "HH"] | |
| self.use_shape = use_shape | |
| self.shape_res = shape_res | |
| self.unknown_organ_id = unknown_organ_id | |
| self.keep_aspect_ratio = keep_aspect_ratio | |
| self.self_normalize = self_normalize | |
| self.id2label = {0: "foreground"} | |
| self.label2id = {"foreground": 0} | |