def main():
opt = get_args()
print(opt)
device = get_device()
print("Device count", torch.cuda.device_count())
if opt.seed > -1:
set_seed(opt.seed)
output_path = Path("results") / "resnet" / opt.name
output_path.mkdir(parents=True, exist_ok=True)
# Save options.
with open(output_path / "opt.json", "w") as f:
json.dump(vars(opt), f, indent=4)
logger = ResNetLogger(opt.name, opt.n_epochs, opt.tensorboard)
use_pretrained = False
feature_extract = False
n_classes = 5
# Freeze layers if we're only using it for feature extraction.
model = create_small_retina_model(use_pretrained,
feature_extract,
n_classes,
load_name=opt.load_name)
model = model.to(device)
train(
model=model,
num_epochs=opt.n_epochs,
log_interval=opt.log_interval,
val_interval=opt.val_interval,
batch_size=opt.batch_size,
img_size=opt.img_size,
lr=opt.lr,
logger=logger,
device=device,
feature_extract=feature_extract,
use_hdf5=opt.use_hdf5,
synthetic_name=opt.synthetic_name,
n_synthetic=opt.n_synthetic,
use_real=opt.use_real,
)
checkpoint_path = output_path / "checkpoints"
checkpoint_path.mkdir(parents=True, exist_ok=True)
torch.save(model.state_dict(), checkpoint_path / "model_latest.pth")
print("Finished!")
def main():
opt = get_args()
# TODO(sonjoonho): Save options.
if opt.seed > 0:
set_seed(opt.seed)
output_path = Path("results") / "resnet_labels" / opt.name
output_path.mkdir(parents=True, exist_ok=True)
logger = ResNetLogger(opt.name, opt.n_epochs, opt.tensorboard)
use_pretrained = False
feature_extract = False
# Freeze layers if we're only using it for feature extraction.
device = get_device()
model = create_label_model(use_pretrained, feature_extract)
model = model.to(device)
train(
model=model,
num_epochs=opt.n_epochs,
log_interval=opt.log_interval,
val_interval=opt.val_interval,
batch_size=opt.batch_size,
img_size=opt.img_size,
lr=opt.lr,
logger=logger,
device=device,
feature_extract=feature_extract,
use_synthetic=opt.use_synthetic,
)
checkpoint_path = output_path / "checkpoints"
checkpoint_path.mkdir(parents=True, exist_ok=True)
torch.save(model.state_dict(), checkpoint_path / "model_latest.pth")
def main():
opt = get_args()
path = Path("results") / "progan" / opt.name
checkpoint_path = path / "checkpoints"
with open(path / "opt.json", "r") as f:
opt_train = json.load(f, object_hook=lambda d: SimpleNamespace(**d))
n_classes = 5
which_labels = sorted([Labels[l] for l in opt_train.lesions], key=lambda x: x.value)
n_channels = len(which_labels) + 1
out_path = Path(opt.out_dir) / "progan" / opt.name / "test"
label_out_path = out_path / "label"
inst_out_path = out_path / "inst"
label_out_path.mkdir(exist_ok=True, parents=True)
inst_out_path.mkdir(exist_ok=True, parents=True)
if opt.seed > -1:
set_seed(opt.seed)
device = torch.device("cuda")
checkpoint_suffix = "final"
if opt.epoch:
checkpoint_suffix = opt.epoch
generator = load_generator(
checkpoint_path / f"generator_{checkpoint_suffix}.pth",
n_channels,
opt_train.img_size,
n_classes,
opt_train.latent_dim,
)
generator.to(device)
batch_sizes = [opt.batch_size for _ in range(opt.n_samples // opt.batch_size)]
last_batch = opt.n_samples % opt.batch_size
if last_batch > 0:
batch_sizes += [last_batch]
print(f"Batch sizes: {batch_sizes}")
total_idx = 0
for i, bs in enumerate(batch_sizes):
if n_classes is None:
gen_label = None
z = torch.randn((bs, opt_train.latent_dim), device=device)
else:
z = torch.randn((bs, opt_train.latent_dim - n_classes), device=device)
gen_label = torch.randint(n_classes, (bs,), device=device)
outputs = generator(z, labels=gen_label, depth=9, alpha=1.0).detach()
if opt.upsample_factor > 1:
# Interestingly, transforms.Resize now gives a warning when not using the
# `InterpolationMode` enum, but the same is not true for
# `functional.interpolate`.
outputs = F.interpolate(
outputs,
scale_factor=opt.upsample_factor,
mode="bilinear",
align_corners=False,
)
outputs = torch.where(outputs > 0.5, 1.0, 0.0)
labels, inst = split_channels(outputs)
inst = torch.ones_like(inst) - inst
inst *= 255
# At this point, `labels` contains values in the range [0, 255]. Pytorch's
# `save_image` function expects values between [0, 1]. `colour_labels` does this
# scaling for us, otherwise we do it here.
labels = torch.argmax(labels, dim=1, keepdim=True).float()
if opt.mask_retina:
_, height, width = labels[0].shape
image = Labels.BG.value * np.ones((height, width))
center = (height // 2, width // 2)
radius = height // 2 - 1
colour = Labels.RETINA.value
circle = cv2.circle(image, center, radius, colour, thickness=cv2.FILLED)
circle_tensor = torch.from_numpy(circle)
circle_tensor = circle_tensor.to(device).float()
labels = torch.where(labels == Labels.RETINA.value, circle_tensor, labels)
labels = torch.where(labels == Labels.BG.value, circle_tensor, labels)
if opt.colour:
labels = colour_labels_flat(labels) * 255.0
else:
labels[labels == Labels.BG.value] = 255
# We permute instead of squeezing since the labels may have colour channels.
labels = labels.permute(0, 2, 3, 1).cpu().numpy()
inst = inst.permute(0, 2, 3, 1).cpu().numpy()
for j in range(bs):
total_idx += 1
if gen_label is None:
dr_grade = ""
else:
dr_grade = gen_label[j].item()
# Save as PNG to avoid compression artifacts.
filename = f"test_{dr_grade}_{total_idx:05}.png"
label_path = str(out_path / "label" / filename)
if opt.colour:
label = cv2.cvtColor(labels[j], cv2.COLOR_BGR2RGB)
else:
label = labels[j]
cv2.imwrite(label_path, label)
inst_path = str(out_path / "inst" / filename)
cv2.imwrite(
inst_path,
inst[j],
)
print(f"Generated images of shape {labels.shape}")
def main():
opt = get_args()
print(opt)
set_seed(213)
if opt.conditional:
n_classes = 5
else:
n_classes = None
device = get_device()
print(f"Device count: {torch.cuda.device_count()}")
output_path = Path(opt.output_dir) / "progan" / opt.name
output_path.mkdir(parents=True, exist_ok=True)
# Save options.
with open(output_path / "opt.json", "w") as f:
json.dump(vars(opt), f, indent=4)
n_channels = len(opt.lesions) + 1
# The initial resolution will be 2**2 = 4.
# The final resolution will be 2**9 = 512.
start_depth = 2
final_depth = 9
generator = Generator(
depth=final_depth,
n_channels=n_channels,
latent_size=opt.latent_dim,
n_classes=n_classes,
)
discriminator = Discriminator(
depth=final_depth,
num_channels=n_channels,
latent_size=opt.latent_dim,
n_classes=n_classes,
)
generator = DataParallel(generator)
discriminator = DataParallel(discriminator)
generator.to(device)
discriminator.to(device)
transform = transforms.Compose(
[
transforms.Resize(opt.img_size, InterpolationMode.NEAREST),
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
transforms.ToTensor(),
],
)
dataset = CombinedDataset(
return_inst=False,
return_image=False,
return_transformed=False,
label_transform=transform,
)
if opt.use_copypaste:
synthetic_dataset = CopyPasteDataset(label_transform=transform)
dataset = ConcatDataset((dataset, synthetic_dataset))
depth_epochs = [20, 40, 60, 80, 100, 120, 140, 160]
batch_sizes = [512, 256, 128, 64, 32, 16, 8, 4]
n_stages = final_depth - start_depth + 1
assert len(depth_epochs) == n_stages
assert len(batch_sizes) == n_stages
total_epochs = sum(depth_epochs)
logger = ProGANLogger(opt.name, total_epochs, opt.tensorboard)
start_time = time.time()
train(
generator,
discriminator,
dataset,
batch_sizes,
device,
output_path,
opt.lr_g,
opt.lr_d,
depth_epochs,
opt.n_critic,
opt.n_gen,
opt.clip_gradient,
opt.label_smoothing,
opt.latent_dim,
n_classes,
opt.sample_interval,
opt.chkpt_interval,
logger,
opt.log_step,
opt.lesions,
opt.use_ada,
start_depth,
final_depth,
)
logger.close()
end_time = time.time()
# Throw away fractional seconds since we don't need that level of precision.
execution_time = int(end_time - start_time)
print(f"Finished in {format_seconds(execution_time)}")
def main():
opt = get_args()
print(opt)
device = get_device()
print("Device count", torch.cuda.device_count())
if opt.seed > -1:
set_seed(opt.seed)
output_path = Path("results") / "unet" / opt.name
checkpoint_path = output_path / "checkpoints"
checkpoint_path.mkdir(parents=True, exist_ok=True)
# Save options.
with open(output_path / "opt.json", "w") as f:
json.dump(vars(opt), f, indent=4)
which_labels = sorted([Labels[l] for l in opt.lesions],
key=lambda x: x.value)
n_classes = len(which_labels) + 1
model = create_model(opt.load_name, n_classes)
model = model.to(device=device)
train_loader, val_loader = make_dataloaders(
opt.img_size,
opt.n_synthetic,
opt.batch_size,
opt.synthetic_name,
)
n_train = len(train_loader.dataset)
n_val = len(val_loader.dataset)
print(f"""
Name: {opt.name}
Epochs: {opt.n_epochs}
Training size: {n_train}
Validation size: {n_val}
Real size: {opt.n_real}
Synthetic size: {opt.n_synthetic}
Labels: {which_labels}
""")
logger = UNetLogger(opt.name, opt.n_epochs, opt.tensorboard)
train(
model=model,
epochs=opt.n_epochs,
lr=opt.lr,
device=device,
log_interval=opt.log_interval,
val_interval=opt.val_interval,
checkpoint_path=checkpoint_path,
train_loader=train_loader,
val_loader=val_loader,
logger=logger,
labels=which_labels,
)
logger.close()
print("Finished!")
def main():
opt = get_args()
print(opt)
set_seed(213)
n_classes = 5
device = get_device()
output_path = Path(opt.output_dir) / "acgan" / opt.name
output_path.mkdir(parents=True, exist_ok=True)
# Save options.
with open(output_path / "opt.json", "w") as f:
json.dump(vars(opt), f, indent=4)
n_channels = len(opt.lesions) + 1
generator = Generator(n_channels, opt.img_size, n_classes, opt.latent_dim)
discriminator = Discriminator(n_channels, opt.img_size, n_classes)
generator.to(device)
discriminator.to(device)
# Initialize weights.
generator.apply(weights_init_normal)
discriminator.apply(weights_init_normal)
transform = transforms.Compose([
transforms.Resize(opt.img_size, InterpolationMode.NEAREST),
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
transforms.ToTensor(),
], )
dataset = CombinedDataset(
return_inst=False,
return_image=False,
return_transformed=False,
return_filename=False,
label_transform=transform,
)
if opt.filter_dataset:
dataset.df = dataset.df[dataset.df["Source"] == opt.filter_dataset]
if opt.use_copypaste:
synthetic_dataset = CopyPasteDataset(label_transform=transform)
dataset = ConcatDataset((dataset, synthetic_dataset))
dataloader = DataLoader(
dataset,
batch_size=opt.batch_size,
shuffle=True,
drop_last=True,
)
logger = ACGANLogger(opt.name, opt.n_epochs, opt.tensorboard)
start_time = time.time()
print(f"Train size: {len(dataset)}")
train(
generator,
discriminator,
dataloader,
device,
output_path,
opt.lr_g,
opt.lr_d,
opt.n_epochs,
opt.n_critic,
opt.n_gen,
opt.clip_gradient,
opt.label_smoothing,
opt.latent_dim,
n_classes,
opt.sample_interval,
opt.val_interval,
opt.chkpt_interval,
logger,
opt.log_step,
opt.lesions,
opt.use_ada,
)
logger.close()
end_time = time.time()
# Throw away fractional seconds since we don't need that level of precision.
execution_time = int(end_time - start_time)
print(f"Finished in {format_seconds(execution_time)}")
def main():
opt = get_args()
out_path = Path(opt.out_dir) / "copypaste"
label_path = out_path / "label"
label_path.mkdir(parents=True, exist_ok=True)
inst_path = out_path / "inst"
inst_path.mkdir(parents=True, exist_ok=True)
if opt.seed > -1:
set_seed(opt.seed)
nc = 8 + 1
transform = T.Compose([
T.Resize(opt.img_size, interpolation=T.InterpolationMode.NEAREST),
T.ToTensor(),
])
samples_per_grade = opt.n_samples // 5
ticker = 0
for dr_grade in range(5):
dataset = CombinedDataset(
return_image=False,
return_transformed=False,
label_transform=transform,
)
dataset.df = dataset.df[dataset.df["Grade"] == dr_grade]
dataloader = DataLoader(dataset,
batch_size=nc,
shuffle=True,
drop_last=True)
infinite_loader = infinite(dataloader)
for _ in tqdm(range(samples_per_grade)):
batch = next(infinite_loader)
source_label = batch["label"]
assert len(source_label) == nc
source_label = get_label_semantics(source_label)
_, _, height, width = source_label.shape
retina = make_circle(height, width)
od = source_label[[1], Labels.OD.value, :, :]
ma = source_label[[2], Labels.MA.value, :, :]
he = source_label[[3], Labels.HE.value, :, :]
ex = source_label[[4], Labels.EX.value, :, :]
se = source_label[[5], Labels.SE.value, :, :]
nv = source_label[[6], Labels.NV.value, :, :]
irma = source_label[[7], Labels.IRMA.value, :, :]
retina = retina - od - ma - he - ex - se - nv - irma
# Create background by subtracting everything else.
bg = torch.ones_like(
retina) - od - ma - he - ex - se - nv - irma - retina
bg = torch.clamp(bg, 0, 1)
combined = torch.cat([retina, od, ma, he, ex, se, nv, irma, bg],
dim=0)
new_label = torch.argmax(combined, dim=0).float().numpy()
if opt.colour:
new_label = colour_labels_numpy(new_label)
else:
# Set background to 255.
new_label[new_label == (nc - 1)] = 255.0
# TODO(sonjoonho): Generate grade more intelligently. This could be done just by
# sampling from existing images with the desired grade.
# End-point is inclusive.
filename = f"copypaste_{dr_grade}_{ticker:05}.png"
cv2.imwrite(str(label_path / filename), new_label)
inst = od.squeeze().numpy()
inst = np.ones_like(inst) - inst
inst *= 255.0
cv2.imwrite(str(inst_path / filename), inst)
ticker += 1
def generate_listwise_cookpad(_ctx):
set_seed()
cookpad.generate_n_listwise.generate()
def generate_listwise_msmarco(_ctx):
set_seed()
msmarco.generate_n_listwise.generate()