def test_shape(self, input_param, input_shape, expected_shape):
device = "cuda" if torch.cuda.is_available() else "cpu"
print(input_param)
net = BasicUNet(**input_param).to(device)
with eval_mode(net):
result = net(torch.randn(input_shape).to(device))
self.assertEqual(result.shape, expected_shape)
def get_network(network, channels, dimensions):
if network == 'unet':
if channels == 16:
features = (16, 32, 64, 128, 256)
elif channels == 32:
features = (32, 64, 128, 256, 512)
else:
features = (64, 128, 256, 512, 1024)
logging.info('Using Unet with features: {}'.format(features))
network = UNet(dimensions=dimensions,
in_channels=3,
out_channels=1,
channels=features,
strides=[2, 2, 2, 2],
norm=Norm.BATCH)
else:
if channels == 16:
features = (16, 32, 64, 128, 256, 16)
elif channels == 32:
features = (32, 64, 128, 256, 512, 32)
else:
features = (64, 128, 256, 512, 1024, 64)
logging.info('Using BasicUnet with features: {}'.format(features))
network = BasicUNet(dimensions=dimensions,
in_channels=3,
out_channels=1,
features=features)
return network
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",
"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 run_test(net_name="basicunet",
batch_size=64,
train_steps=100,
device="cuda:0"):
class _TestBatch(Dataset):
def __getitem__(self, _unused_id):
im, seg = create_test_image_2d(128,
128,
noise_max=1,
num_objs=4,
num_seg_classes=1)
return im[None], seg[None].astype(np.float32)
def __len__(self):
return train_steps
if net_name == "basicunet":
net = BasicUNet(spatial_dims=2,
in_channels=1,
out_channels=1,
features=(4, 8, 8, 16, 16, 32))
elif net_name == "unet":
net = UNet(spatial_dims=2,
in_channels=1,
out_channels=1,
channels=(4, 8, 16, 32),
strides=(2, 2, 2),
num_res_units=2)
net.to(device)
loss = DiceLoss(sigmoid=True)
opt = torch.optim.Adam(net.parameters(), 1e-4)
src = DataLoader(_TestBatch(), batch_size=batch_size)
trainer = create_supervised_trainer(net, opt, loss, device, False)
trainer.run(src, 1)
loss = trainer.state.output
return loss
def test_script(self):
net = BasicUNet(spatial_dims=2, in_channels=1, out_channels=3)
test_data = torch.randn(16, 1, 32, 32)
test_script_save(net, test_data)
def test_script(self):
net = BasicUNet(dimensions=2, in_channels=1, out_channels=3)
test_data = torch.randn(16, 1, 32, 32)
out_orig, out_reloaded = test_script_save(net, test_data)
assert torch.allclose(out_orig, out_reloaded)