"""Functions for training and running segmentation.""" import math import os import time import click import matplotlib.pyplot as plt import numpy as np import scipy.signal import skimage.draw import torch import torchvision import tqdm import echonet @click.command("segmentation") @click.option("--data_dir", type=click.Path(exists=True, file_okay=False), default=None) @click.option("--output", type=click.Path(file_okay=False), default=None) @click.option("--model_name", type=click.Choice( sorted(name for name in torchvision.models.segmentation.__dict__ if name.islower() and not name.startswith("__") and callable(torchvision.models.segmentation.__dict__[name]))), default="deeplabv3_resnet50") @click.option("--pretrained/--random", default=False) @click.option("--weights", type=click.Path(exists=True, dir_okay=False), default=None) @click.option("--run_test/--skip_test", default=False) @click.option("--save_video/--skip_video", default=False) @click.option("--num_epochs", type=int, default=50) @click.option("--lr", type=float, default=1e-5) @click.option("--weight_decay", type=float, default=0) @click.option("--lr_step_period", type=int, default=None) @click.option("--num_train_patients", type=int, default=None) @click.option("--num_workers", type=int, default=4) @click.option("--batch_size", type=int, default=20) @click.option("--device", type=str, default=None) @click.option("--seed", type=int, default=0) def run( data_dir=None, output=None, model_name="deeplabv3_resnet50", pretrained=False, weights=None, run_test=False, save_video=False, num_epochs=50, lr=1e-5, weight_decay=1e-5, lr_step_period=None, num_train_patients=None, num_workers=4, batch_size=20, device=None, seed=0, ): """Trains/tests segmentation model. Args: data_dir (str, optional): Directory containing dataset. Defaults to `echonet.config.DATA_DIR`. output (str, optional): Directory to place outputs. Defaults to output/segmentation/_/. model_name (str, optional): Name of segmentation model. One of ``deeplabv3_resnet50'', ``deeplabv3_resnet101'', ``fcn_resnet50'', or ``fcn_resnet101'' (options are torchvision.models.segmentation.) Defaults to ``deeplabv3_resnet50''. pretrained (bool, optional): Whether to use pretrained weights for model Defaults to False. weights (str, optional): Path to checkpoint containing weights to initialize model. Defaults to None. run_test (bool, optional): Whether or not to run on test. Defaults to False. save_video (bool, optional): Whether to save videos with segmentations. Defaults to False. num_epochs (int, optional): Number of epochs during training Defaults to 50. lr (float, optional): Learning rate for SGD Defaults to 1e-5. weight_decay (float, optional): Weight decay for SGD Defaults to 0. lr_step_period (int or None, optional): Period of learning rate decay (learning rate is decayed by a multiplicative factor of 0.1) Defaults to math.inf (never decay learning rate). num_train_patients (int or None, optional): Number of training patients for ablations. Defaults to all patients. num_workers (int, optional): Number of subprocesses to use for data loading. If 0, the data will be loaded in the main process. Defaults to 4. device (str or None, optional): Name of device to run on. Options from https://pytorch.org/docs/stable/tensor_attributes.html#torch.torch.device Defaults to ``cuda'' if available, and ``cpu'' otherwise. batch_size (int, optional): Number of samples to load per batch Defaults to 20. seed (int, optional): Seed for random number generator. Defaults to 0. """ # Seed RNGs np.random.seed(seed) torch.manual_seed(seed) # Set default output directory if output is None: output = os.path.join("output", "segmentation", "{}_{}".format(model_name, "pretrained" if pretrained else "random")) os.makedirs(output, exist_ok=True) # Set device for computations if device is None: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Set up model model = torchvision.models.segmentation.__dict__[model_name](pretrained=pretrained, aux_loss=False) model.classifier[-1] = torch.nn.Conv2d(model.classifier[-1].in_channels, 1, kernel_size=model.classifier[-1].kernel_size) # change number of outputs to 1 if device.type == "cuda": model = torch.nn.DataParallel(model) model.to(device) if weights is not None: checkpoint = torch.load(weights) model.load_state_dict(checkpoint['state_dict']) # Set up optimizer optim = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=weight_decay) if lr_step_period is None: lr_step_period = math.inf scheduler = torch.optim.lr_scheduler.StepLR(optim, lr_step_period) # Compute mean and std mean, std = echonet.utils.get_mean_and_std(echonet.datasets.Echo(root=data_dir, split="train")) tasks = ["LargeFrame", "SmallFrame", "LargeTrace", "SmallTrace"] kwargs = {"target_type": tasks, "mean": mean, "std": std } # Set up datasets and dataloaders dataset = {} dataset["train"] = echonet.datasets.Echo(root=data_dir, split="train", **kwargs) if num_train_patients is not None and len(dataset["train"]) > num_train_patients: # Subsample patients (used for ablation experiment) indices = np.random.choice(len(dataset["train"]), num_train_patients, replace=False) dataset["train"] = torch.utils.data.Subset(dataset["train"], indices) dataset["val"] = echonet.datasets.Echo(root=data_dir, split="val", **kwargs) # Run training and testing loops with open(os.path.join(output, "log.csv"), "a") as f: epoch_resume = 0 bestLoss = float("inf") try: # Attempt to load checkpoint checkpoint = torch.load(os.path.join(output, "checkpoint.pt")) model.load_state_dict(checkpoint['state_dict']) optim.load_state_dict(checkpoint['opt_dict']) scheduler.load_state_dict(checkpoint['scheduler_dict']) epoch_resume = checkpoint["epoch"] + 1 bestLoss = checkpoint["best_loss"] f.write("Resuming from epoch {}\n".format(epoch_resume)) except FileNotFoundError: f.write("Starting run from scratch\n") for epoch in range(epoch_resume, num_epochs): print("Epoch #{}".format(epoch), flush=True) for phase in ['train', 'val']: start_time = time.time() for i in range(torch.cuda.device_count()): torch.cuda.reset_peak_memory_stats(i) ds = dataset[phase] dataloader = torch.utils.data.DataLoader( ds, batch_size=batch_size, num_workers=num_workers, shuffle=True, pin_memory=(device.type == "cuda"), drop_last=(phase == "train")) loss, large_inter, large_union, small_inter, small_union = echonet.utils.segmentation.run_epoch(model, dataloader, phase == "train", optim, device) overall_dice = 2 * (large_inter.sum() + small_inter.sum()) / (large_union.sum() + large_inter.sum() + small_union.sum() + small_inter.sum()) large_dice = 2 * large_inter.sum() / (large_union.sum() + large_inter.sum()) small_dice = 2 * small_inter.sum() / (small_union.sum() + small_inter.sum()) f.write("{},{},{},{},{},{},{},{},{},{},{}\n".format(epoch, phase, loss, overall_dice, large_dice, small_dice, time.time() - start_time, large_inter.size, sum(torch.cuda.max_memory_allocated() for i in range(torch.cuda.device_count())), sum(torch.cuda.max_memory_reserved() for i in range(torch.cuda.device_count())), batch_size)) f.flush() scheduler.step() # Save checkpoint save = { 'epoch': epoch, 'state_dict': model.state_dict(), 'best_loss': bestLoss, 'loss': loss, 'opt_dict': optim.state_dict(), 'scheduler_dict': scheduler.state_dict(), } torch.save(save, os.path.join(output, "checkpoint.pt")) if loss < bestLoss: torch.save(save, os.path.join(output, "best.pt")) bestLoss = loss # Load best weights if num_epochs != 0: checkpoint = torch.load(os.path.join(output, "best.pt")) model.load_state_dict(checkpoint['state_dict']) f.write("Best validation loss {} from epoch {}\n".format(checkpoint["loss"], checkpoint["epoch"])) if run_test: # Run on validation and test for split in ["val", "test"]: dataset = echonet.datasets.Echo(root=data_dir, split=split, **kwargs) dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False, pin_memory=(device.type == "cuda")) loss, large_inter, large_union, small_inter, small_union = echonet.utils.segmentation.run_epoch(model, dataloader, False, None, device) overall_dice = 2 * (large_inter + small_inter) / (large_union + large_inter + small_union + small_inter) large_dice = 2 * large_inter / (large_union + large_inter) small_dice = 2 * small_inter / (small_union + small_inter) with open(os.path.join(output, "{}_dice.csv".format(split)), "w") as g: g.write("Filename, Overall, Large, Small\n") for (filename, overall, large, small) in zip(dataset.fnames, overall_dice, large_dice, small_dice): g.write("{},{},{},{}\n".format(filename, overall, large, small)) f.write("{} dice (overall): {:.4f} ({:.4f} - {:.4f})\n".format(split, *echonet.utils.bootstrap(np.concatenate((large_inter, small_inter)), np.concatenate((large_union, small_union)), echonet.utils.dice_similarity_coefficient))) f.write("{} dice (large): {:.4f} ({:.4f} - {:.4f})\n".format(split, *echonet.utils.bootstrap(large_inter, large_union, echonet.utils.dice_similarity_coefficient))) f.write("{} dice (small): {:.4f} ({:.4f} - {:.4f})\n".format(split, *echonet.utils.bootstrap(small_inter, small_union, echonet.utils.dice_similarity_coefficient))) f.flush() # Saving videos with segmentations dataset = echonet.datasets.Echo(root=data_dir, split="test", target_type=["Filename", "LargeIndex", "SmallIndex"], # Need filename for saving, and human-selected frames to annotate mean=mean, std=std, # Normalization length=None, max_length=None, period=1 # Take all frames ) dataloader = torch.utils.data.DataLoader(dataset, batch_size=10, num_workers=num_workers, shuffle=False, pin_memory=False, collate_fn=_video_collate_fn) # Save videos with segmentation if save_video and not all(os.path.isfile(os.path.join(output, "videos", f)) for f in dataloader.dataset.fnames): # Only run if missing videos model.eval() os.makedirs(os.path.join(output, "videos"), exist_ok=True) os.makedirs(os.path.join(output, "size"), exist_ok=True) echonet.utils.latexify() with torch.no_grad(): with open(os.path.join(output, "size.csv"), "w") as g: g.write("Filename,Frame,Size,HumanLarge,HumanSmall,ComputerSmall\n") for (x, (filenames, large_index, small_index), length) in tqdm.tqdm(dataloader): # Run segmentation model on blocks of frames one-by-one # The whole concatenated video may be too long to run together y = np.concatenate([model(x[i:(i + batch_size), :, :, :].to(device))["out"].detach().cpu().numpy() for i in range(0, x.shape[0], batch_size)]) start = 0 x = x.numpy() for (i, (filename, offset)) in enumerate(zip(filenames, length)): # Extract one video and segmentation predictions video = x[start:(start + offset), ...] logit = y[start:(start + offset), 0, :, :] # Un-normalize video video *= std.reshape(1, 3, 1, 1) video += mean.reshape(1, 3, 1, 1) # Get frames, channels, height, and width f, c, h, w = video.shape # pylint: disable=W0612 assert c == 3 # Put two copies of the video side by side video = np.concatenate((video, video), 3) # If a pixel is in the segmentation, saturate blue channel # Leave alone otherwise video[:, 0, :, w:] = np.maximum(255. * (logit > 0), video[:, 0, :, w:]) # pylint: disable=E1111 # Add blank canvas under pair of videos video = np.concatenate((video, np.zeros_like(video)), 2) # Compute size of segmentation per frame size = (logit > 0).sum((1, 2)) # Identify systole frames with peak detection trim_min = sorted(size)[round(len(size) ** 0.05)] trim_max = sorted(size)[round(len(size) ** 0.95)] trim_range = trim_max - trim_min systole = set(scipy.signal.find_peaks(-size, distance=20, prominence=(0.50 * trim_range))[0]) # Write sizes and frames to file for (frame, s) in enumerate(size): g.write("{},{},{},{},{},{}\n".format(filename, frame, s, 1 if frame == large_index[i] else 0, 1 if frame == small_index[i] else 0, 1 if frame in systole else 0)) # Plot sizes fig = plt.figure(figsize=(size.shape[0] / 50 * 1.5, 3)) plt.scatter(np.arange(size.shape[0]) / 50, size, s=1) ylim = plt.ylim() for s in systole: plt.plot(np.array([s, s]) / 50, ylim, linewidth=1) plt.ylim(ylim) plt.title(os.path.splitext(filename)[0]) plt.xlabel("Seconds") plt.ylabel("Size (pixels)") plt.tight_layout() plt.savefig(os.path.join(output, "size", os.path.splitext(filename)[0] + ".pdf")) plt.close(fig) # Normalize size to [0, 1] size -= size.min() size = size / size.max() size = 1 - size # Iterate the frames in this video for (f, s) in enumerate(size): # On all frames, mark a pixel for the size of the frame video[:, :, int(round(115 + 100 * s)), int(round(f / len(size) * 200 + 10))] = 255. if f in systole: # If frame is computer-selected systole, mark with a line video[:, :, 115:224, int(round(f / len(size) * 200 + 10))] = 255. def dash(start, stop, on=10, off=10): buf = [] x = start while x < stop: buf.extend(range(x, x + on)) x += on x += off buf = np.array(buf) buf = buf[buf < stop] return buf d = dash(115, 224) if f == large_index[i]: # If frame is human-selected diastole, mark with green dashed line on all frames video[:, :, d, int(round(f / len(size) * 200 + 10))] = np.array([0, 225, 0]).reshape((1, 3, 1)) if f == small_index[i]: # If frame is human-selected systole, mark with red dashed line on all frames video[:, :, d, int(round(f / len(size) * 200 + 10))] = np.array([0, 0, 225]).reshape((1, 3, 1)) # Get pixels for a circle centered on the pixel r, c = skimage.draw.disk((int(round(115 + 100 * s)), int(round(f / len(size) * 200 + 10))), 4.1) # On the frame that's being shown, put a circle over the pixel video[f, :, r, c] = 255. # Rearrange dimensions and save video = video.transpose(1, 0, 2, 3) video = video.astype(np.uint8) echonet.utils.savevideo(os.path.join(output, "videos", filename), video, 50) # Move to next video start += offset def run_epoch(model, dataloader, train, optim, device): """Run one epoch of training/evaluation for segmentation. Args: model (torch.nn.Module): Model to train/evaulate. dataloder (torch.utils.data.DataLoader): Dataloader for dataset. train (bool): Whether or not to train model. optim (torch.optim.Optimizer): Optimizer device (torch.device): Device to run on """ total = 0. n = 0 pos = 0 neg = 0 pos_pix = 0 neg_pix = 0 model.train(train) large_inter = 0 large_union = 0 small_inter = 0 small_union = 0 large_inter_list = [] large_union_list = [] small_inter_list = [] small_union_list = [] with torch.set_grad_enabled(train): with tqdm.tqdm(total=len(dataloader)) as pbar: for (_, (large_frame, small_frame, large_trace, small_trace)) in dataloader: # Count number of pixels in/out of human segmentation pos += (large_trace == 1).sum().item() pos += (small_trace == 1).sum().item() neg += (large_trace == 0).sum().item() neg += (small_trace == 0).sum().item() # Count number of pixels in/out of computer segmentation pos_pix += (large_trace == 1).sum(0).to("cpu").detach().numpy() pos_pix += (small_trace == 1).sum(0).to("cpu").detach().numpy() neg_pix += (large_trace == 0).sum(0).to("cpu").detach().numpy() neg_pix += (small_trace == 0).sum(0).to("cpu").detach().numpy() # Run prediction for diastolic frames and compute loss large_frame = large_frame.to(device) large_trace = large_trace.to(device) y_large = model(large_frame)["out"] loss_large = torch.nn.functional.binary_cross_entropy_with_logits(y_large[:, 0, :, :], large_trace, reduction="sum") # Compute pixel intersection and union between human and computer segmentations large_inter += np.logical_and(y_large[:, 0, :, :].detach().cpu().numpy() > 0., large_trace[:, :, :].detach().cpu().numpy() > 0.).sum() large_union += np.logical_or(y_large[:, 0, :, :].detach().cpu().numpy() > 0., large_trace[:, :, :].detach().cpu().numpy() > 0.).sum() large_inter_list.extend(np.logical_and(y_large[:, 0, :, :].detach().cpu().numpy() > 0., large_trace[:, :, :].detach().cpu().numpy() > 0.).sum((1, 2))) large_union_list.extend(np.logical_or(y_large[:, 0, :, :].detach().cpu().numpy() > 0., large_trace[:, :, :].detach().cpu().numpy() > 0.).sum((1, 2))) # Run prediction for systolic frames and compute loss small_frame = small_frame.to(device) small_trace = small_trace.to(device) y_small = model(small_frame)["out"] loss_small = torch.nn.functional.binary_cross_entropy_with_logits(y_small[:, 0, :, :], small_trace, reduction="sum") # Compute pixel intersection and union between human and computer segmentations small_inter += np.logical_and(y_small[:, 0, :, :].detach().cpu().numpy() > 0., small_trace[:, :, :].detach().cpu().numpy() > 0.).sum() small_union += np.logical_or(y_small[:, 0, :, :].detach().cpu().numpy() > 0., small_trace[:, :, :].detach().cpu().numpy() > 0.).sum() small_inter_list.extend(np.logical_and(y_small[:, 0, :, :].detach().cpu().numpy() > 0., small_trace[:, :, :].detach().cpu().numpy() > 0.).sum((1, 2))) small_union_list.extend(np.logical_or(y_small[:, 0, :, :].detach().cpu().numpy() > 0., small_trace[:, :, :].detach().cpu().numpy() > 0.).sum((1, 2))) # Take gradient step if training loss = (loss_large + loss_small) / 2 if train: optim.zero_grad() loss.backward() optim.step() # Accumulate losses and compute baselines total += loss.item() n += large_trace.size(0) p = pos / (pos + neg) p_pix = (pos_pix + 1) / (pos_pix + neg_pix + 2) # Show info on process bar pbar.set_postfix_str("{:.4f} ({:.4f}) / {:.4f} {:.4f}, {:.4f}, {:.4f}".format(total / n / 112 / 112, loss.item() / large_trace.size(0) / 112 / 112, -p * math.log(p) - (1 - p) * math.log(1 - p), (-p_pix * np.log(p_pix) - (1 - p_pix) * np.log(1 - p_pix)).mean(), 2 * large_inter / (large_union + large_inter), 2 * small_inter / (small_union + small_inter))) pbar.update() large_inter_list = np.array(large_inter_list) large_union_list = np.array(large_union_list) small_inter_list = np.array(small_inter_list) small_union_list = np.array(small_union_list) return (total / n / 112 / 112, large_inter_list, large_union_list, small_inter_list, small_union_list, ) def _video_collate_fn(x): """Collate function for Pytorch dataloader to merge multiple videos. This function should be used in a dataloader for a dataset that returns a video as the first element, along with some (non-zero) tuple of targets. Then, the input x is a list of tuples: - x[i][0] is the i-th video in the batch - x[i][1] are the targets for the i-th video This function returns a 3-tuple: - The first element is the videos concatenated along the frames dimension. This is done so that videos of different lengths can be processed together (tensors cannot be "jagged", so we cannot have a dimension for video, and another for frames). - The second element is contains the targets with no modification. - The third element is a list of the lengths of the videos in frames. """ video, target = zip(*x) # Extract the videos and targets # ``video'' is a tuple of length ``batch_size'' # Each element has shape (channels=3, frames, height, width) # height and width are expected to be the same across videos, but # frames can be different. # ``target'' is also a tuple of length ``batch_size'' # Each element is a tuple of the targets for the item. i = list(map(lambda t: t.shape[1], video)) # Extract lengths of videos in frames # This contatenates the videos along the the frames dimension (basically # playing the videos one after another). The frames dimension is then # moved to be first. # Resulting shape is (total frames, channels=3, height, width) video = torch.as_tensor(np.swapaxes(np.concatenate(video, 1), 0, 1)) # Swap dimensions (approximately a transpose) # Before: target[i][j] is the j-th target of element i # After: target[i][j] is the i-th target of element j target = zip(*target) return video, target, i