def init(self, name: str, model_dir: str, conf: Dict[str, str],
planner: Any, **kwargs):
super().init(name, model_dir, conf, planner, **kwargs)
# Labels
self.labels = ["Nuclei"]
self.label_colors = {"Nuclei": (0, 255, 255)}
# Model Files
self.path = [
os.path.join(self.model_dir,
f"pretrained_{name}.pt"), # pretrained
os.path.join(self.model_dir, f"{name}.pt"), # published
]
# Download PreTrained Model
if strtobool(self.conf.get("use_pretrained_model", "true")):
url = f"{self.PRE_TRAINED_PATH}/pathology_deepedit_nuclei.pt"
download_file(url, self.path[0])
# Network
self.network = BasicUNet(
spatial_dims=2,
in_channels=5,
out_channels=1,
features=(32, 64, 128, 256, 512, 32),
)
def init(self, name: str, model_dir: str, conf: Dict[str, str],
planner: Any, **kwargs):
super().init(name, model_dir, conf, planner, **kwargs)
# Labels
self.labels = {
"spleen": 1,
"right kidney": 2,
"left kidney": 3,
"gallbladder": 4,
"esophagus": 5,
"liver": 6,
"stomach": 7,
"aorta": 8,
"inferior vena cava": 9,
"portal vein and splenic vein": 10,
"pancreas": 11,
"right adrenal gland": 12,
"left adrenal gland": 13,
}
# Model Files
self.path = [
os.path.join(self.model_dir,
f"pretrained_{name}.pt"), # pretrained
os.path.join(self.model_dir, f"{name}.pt"), # published
]
# Download PreTrained Model
if strtobool(self.conf.get("use_pretrained_model", "true")):
url = f"{self.PRE_TRAINED_PATH}/segmentation_unet_multilabel.pt"
download_file(url, self.path[0])
self.target_spacing = (1.0, 1.0, 1.0) # target space for image
self.spatial_size = (96, 96, 96) # train input size
self.roi_size = (128, 128, 128) # sliding window size for infer
# Network
self.network = UNet(
spatial_dims=3,
in_channels=1,
out_channels=len(self.labels.keys()) +
1, # All labels plus background
channels=[16, 32, 64, 128, 256],
strides=[2, 2, 2, 2],
num_res_units=2,
norm="batch",
)
def init(self, name: str, model_dir: str, conf: Dict[str, str],
planner: Any, **kwargs):
super().init(name, model_dir, conf, planner, **kwargs)
# Labels
self.labels = ["spleen"]
# Model Files
self.path = [
os.path.join(self.model_dir,
f"pretrained_{name}.pt"), # pretrained
os.path.join(self.model_dir, f"{name}.pt"), # published
]
# Download PreTrained Model
if strtobool(self.conf.get("use_pretrained_model", "true")):
url = f"{self.NGC_PATH}/clara_pt_spleen_ct_segmentation/versions/1/files/models/model.pt"
download_file(url, self.path[0])
# Network
self.network = UNet(
spatial_dims=3,
in_channels=1,
out_channels=2,
channels=[16, 32, 64, 128, 256],
strides=[2, 2, 2, 2],
num_res_units=2,
norm="batch",
)
# Others
self.epistemic_enabled = strtobool(
conf.get("epistemic_enabled", "false"))
self.epistemic_samples = int(conf.get("epistemic_samples", "5"))
logger.info(
f"EPISTEMIC Enabled: {self.epistemic_enabled}; Samples: {self.epistemic_samples}"
)
self.tta_enabled = strtobool(conf.get("tta_enabled", "false"))
self.tta_samples = int(conf.get("tta_samples", "5"))
logger.info(
f"TTA Enabled: {self.tta_enabled}; Samples: {self.tta_samples}")
def init(self, name: str, model_dir: str, conf: Dict[str, str],
planner: Any, **kwargs):
super().init(name, model_dir, conf, planner, **kwargs)
# Labels
self.labels = [
"spleen",
"right kidney",
"left kidney",
"gallbladder",
"esophagus",
"liver",
"stomach",
"aorta",
"inferior vena cava",
"portal vein and splenic vein",
"pancreas",
"right adrenal gland",
"left adrenal gland",
]
# Model Files
self.path = [
os.path.join(self.model_dir,
f"pretrained_{name}.pt"), # pretrained
os.path.join(self.model_dir, f"{name}.pt"), # published
]
# Download PreTrained Model
if strtobool(self.conf.get("use_pretrained_model", "true")):
url = f"{self.NGC_PATH}/clara_pt_deepgrow_3d_annotation/versions/1/files/models/model.pt"
download_file(url, self.path[0])
# Network
self.network = BasicUNet(
spatial_dims=3,
in_channels=3,
out_channels=1,
features=(32, 64, 128, 256, 512, 32),
)
def init(self, name: str, model_dir: str, conf: Dict[str, str],
planner: Any, **kwargs):
super().init(name, model_dir, conf, planner, **kwargs)
# Labels
self.labels = {
"Neoplastic cells": 1,
"Inflammatory": 2,
"Connective/Soft tissue cells": 3,
"Dead Cells": 4,
"Epithelial": 5,
}
self.label_colors = {
"Neoplastic cells": (255, 0, 0),
"Inflammatory": (255, 255, 0),
"Connective/Soft tissue cells": (0, 255, 0),
"Dead Cells": (0, 0, 0),
"Epithelial": (0, 0, 255),
}
# Model Files
self.path = [
os.path.join(self.model_dir,
f"pretrained_{name}.pt"), # pretrained
os.path.join(self.model_dir, f"{name}.pt"), # published
]
# Download PreTrained Model
if strtobool(self.conf.get("use_pretrained_model", "true")):
url = f"{self.PRE_TRAINED_PATH}/pathology_segmentation_nuclei.pt"
download_file(url, self.path[0])
# Network
self.network = BasicUNet(
spatial_dims=2,
in_channels=3,
out_channels=6,
features=(32, 64, 128, 256, 512, 32),
)
def init(self, name: str, model_dir: str, conf: Dict[str, str],
planner: Any, **kwargs):
super().init(name, model_dir, conf, planner, **kwargs)
# Labels
self.labels = ["Nuclei"]
self.label_colors = {"Nuclei": (0, 255, 255)}
# Model Files
self.path = [
os.path.join(self.model_dir,
f"pretrained_{name}.pt"), # pretrained
os.path.join(self.model_dir, f"{name}.pt"), # published
]
# Download PreTrained Model
if strtobool(self.conf.get("use_pretrained_model", "true")):
url = f"{self.PRE_TRAINED_PATH}/NuClick_UNet_40xAll.pth"
download_file(url, self.path[0])
# Network
self.network = UNet(n_channels=5, n_classes=1)
def init(self, name: str, model_dir: str, conf: Dict[str, str],
planner: Any, **kwargs):
super().init(name, model_dir, conf, planner, **kwargs)
# Multilabel
# self.labels = {
# "spleen": 1,
# "right kidney": 2,
# "left kidney": 3,
# "liver": 6,
# "stomach": 7,
# "aorta": 8,
# "inferior vena cava": 9,
# "background": 0,
# }
# Single label
self.labels = {
"spleen": 1,
"background": 0,
}
# Number of input channels - 4 for BRATS and 1 for spleen
self.number_intensity_ch = 1
network = self.conf.get("network", "dynunet")
# Model Files
self.path = [
os.path.join(self.model_dir,
f"pretrained_{self.name}_{network}.pt"), # pretrained
os.path.join(self.model_dir,
f"{self.name}_{network}.pt"), # published
]
# Download PreTrained Model
if strtobool(self.conf.get("use_pretrained_model", "true")):
url = f"{self.PRE_TRAINED_PATH}/deepedit_{network}_singlelabel.pt"
download_file(url, self.path[0])
self.target_spacing = (1.0, 1.0, 1.0) # target space for image
self.spatial_size = (128, 128, 128) # train input size
# Network
if network == "unetr":
self.network = UNETR(
spatial_dims=3,
in_channels=len(self.labels) + self.number_intensity_ch,
out_channels=len(self.labels),
img_size=self.spatial_size,
feature_size=64,
hidden_size=1536,
mlp_dim=3072,
num_heads=48,
pos_embed="conv",
norm_name="instance",
res_block=True,
)
else:
self.network = DynUNet(
spatial_dims=3,
in_channels=len(self.labels) + self.number_intensity_ch,
out_channels=len(self.labels),
kernel_size=[[3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3],
[3, 3, 3], [3, 3, 3]],
strides=[[1, 1, 1], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2],
[2, 2, 1]],
upsample_kernel_size=[[2, 2, 2], [2, 2, 2], [2, 2, 2],
[2, 2, 2], [2, 2, 1]],
norm_name="instance",
deep_supervision=False,
res_block=True,
)
