File size: 1,436 Bytes
8966f37 badc3e1 8966f37 badc3e1 8966f37 | 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 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | 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}
|