|
|
|
|
|
|
|
|
|
|
|
|
| from __future__ import annotations
|
| import os
|
| import pickle
|
| import numpy as np
|
| import re
|
| from scipy.ndimage import distance_transform_edt as distance
|
|
|
| import torch
|
| from torch.utils.data.sampler import Sampler
|
| import torch.distributed as dist
|
| import cv2
|
| from numpy import ndarray
|
|
|
|
|
|
|
|
|
| def load_model(path):
|
| """Loads model and return it without DataParallel table."""
|
| if os.path.isfile(path):
|
| print("=> loading checkpoint '{}'".format(path))
|
| checkpoint = torch.load(path)
|
|
|
| for key in checkpoint["state_dict"]:
|
| print(key)
|
|
|
|
|
| N = checkpoint["state_dict"]["decoder.out_conv.bias"].size()
|
|
|
|
|
| sob = "sobel.0.weight" in checkpoint["state_dict"].keys()
|
| model = models.__dict__[checkpoint["arch"]](sobel=sob, out=int(N[0]))
|
|
|
|
|
| def rename_key(key):
|
| if not "module" in key:
|
| return key
|
| return "".join(key.split(".module"))
|
|
|
| checkpoint["state_dict"] = {
|
| rename_key(key): val for key, val in checkpoint["state_dict"].items()
|
| }
|
|
|
|
|
| model.load_state_dict(checkpoint["state_dict"])
|
| print("Loaded")
|
| else:
|
| model = None
|
| print("=> no checkpoint found at '{}'".format(path))
|
| return model
|
|
|
|
|
| def load_checkpoint(path, model, optimizer, from_ddp=False):
|
| """loads previous checkpoint
|
|
|
| Args:
|
| path (str): path to checkpoint
|
| model (model): model to restore checkpoint to
|
| optimizer (optimizer): torch optimizer to load optimizer state_dict to
|
| from_ddp (bool, optional): load DistributedDataParallel checkpoint to regular model. Defaults to False.
|
|
|
| Returns:
|
| model, optimizer, epoch_num, loss
|
| """
|
|
|
| checkpoint = torch.load(path)
|
|
|
| model.load_state_dict(checkpoint["state_dict"])
|
|
|
| optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
|
|
|
| loss = checkpoint["loss"]
|
| return model, optimizer, checkpoint["epoch"], loss.item()
|
|
|
|
|
| def restore_model(logger, snapshot_path, model_num=None):
|
| """wrapper function to read log dir and load restore a previous checkpoint
|
|
|
| Args:
|
| logger (Logger): logger object (for info output to console)
|
| snapshot_path (str): path to checkpoint directory
|
|
|
| Returns:
|
| model, optimizer, start_epoch, performance
|
| """
|
| try:
|
|
|
| logger.info(f"Snapshot path: {snapshot_path}")
|
| iter_num = []
|
| name = "model_iter"
|
| if model_num:
|
| name = model_num
|
| for filename in os.listdir(snapshot_path):
|
| if name in filename:
|
| basename, extension = os.path.splitext(filename)
|
| iter_num.append(int(basename.split("_")[2]))
|
| iter_num = max(iter_num)
|
| for filename in os.listdir(snapshot_path):
|
| if name in filename and str(iter_num) in filename:
|
| model_checkpoint = filename
|
| except Exception as e:
|
| logger.warning(f"Error finding previous checkpoints: {e}")
|
|
|
| try:
|
| logger.info(f"Restoring model checkpoint: {model_checkpoint}")
|
| model, optimizer, start_epoch, performance = load_checkpoint(
|
| snapshot_path + "/" + model_checkpoint, model, optimizer
|
| )
|
| logger.info(f"Models restored from iteration {iter_num}")
|
| return model, optimizer, start_epoch, performance
|
| except Exception as e:
|
| logger.warning(f"Unable to restore model checkpoint: {e}, using new model")
|
|
|
|
|
| def save_checkpoint(epoch, model, optimizer, loss, path):
|
| """Saves model as checkpoint"""
|
| torch.save(
|
| {
|
| "epoch": epoch,
|
| "state_dict": model.state_dict(),
|
| "optimizer_state_dict": optimizer.state_dict(),
|
| "loss": loss,
|
| },
|
| path,
|
| )
|
|
|
|
|
| class UnifLabelSampler(Sampler):
|
| """Samples elements uniformely accross pseudolabels.
|
| Args:
|
| N (int): size of returned iterator.
|
| images_lists: dict of key (target), value (list of data with this target)
|
| """
|
|
|
| def __init__(self, N, images_lists):
|
| self.N = N
|
| self.images_lists = images_lists
|
| self.indexes = self.generate_indexes_epoch()
|
|
|
| def generate_indexes_epoch(self):
|
| size_per_pseudolabel = int(self.N / len(self.images_lists)) + 1
|
| res = np.zeros(size_per_pseudolabel * len(self.images_lists))
|
|
|
| for i in range(len(self.images_lists)):
|
| indexes = np.random.choice(
|
| self.images_lists[i],
|
| size_per_pseudolabel,
|
| replace=(len(self.images_lists[i]) <= size_per_pseudolabel),
|
| )
|
| res[i * size_per_pseudolabel : (i + 1) * size_per_pseudolabel] = indexes
|
|
|
| np.random.shuffle(res)
|
| return res[: self.N].astype("int")
|
|
|
| def __iter__(self):
|
| return iter(self.indexes)
|
|
|
| def __len__(self):
|
| return self.N
|
|
|
|
|
| class AverageMeter(object):
|
| """Computes and stores the average and current value"""
|
|
|
| def __init__(self):
|
| self.reset()
|
|
|
| def reset(self):
|
| self.val = 0
|
| self.avg = 0
|
| self.sum = 0
|
| self.count = 0
|
|
|
| def update(self, val, n=1):
|
| self.val = val
|
| self.sum += val * n
|
| self.count += n
|
| self.avg = self.sum / self.count
|
|
|
|
|
| def learning_rate_decay(optimizer, t, lr_0):
|
| for param_group in optimizer.param_groups:
|
| lr = lr_0 / np.sqrt(1 + lr_0 * param_group["weight_decay"] * t)
|
| param_group["lr"] = lr
|
|
|
|
|
| class Logger:
|
| """Class to update every epoch to keep trace of the results
|
| Methods:
|
| - log() log and save
|
| """
|
|
|
| def __init__(self, path):
|
| self.path = path
|
| self.data = []
|
|
|
| def log(self, train_point):
|
| self.data.append(train_point)
|
| with open(os.path.join(self.path), "wb") as fp:
|
| pickle.dump(self.data, fp, -1)
|
|
|
|
|
| def compute_sdf(img_gt, out_shape):
|
| """
|
| compute the signed distance map of binary mask
|
| input: segmentation, shape = (batch_size, x, y, z)
|
| output: the Signed Distance Map (SDM)
|
| sdf(x) = 0; x in segmentation boundary
|
| -inf|x-y|; x in segmentation
|
| +inf|x-y|; x out of segmentation
|
| normalize sdf to [-1,1]
|
| """
|
|
|
| img_gt = img_gt.astype(np.uint8)
|
| normalized_sdf = np.zeros(out_shape)
|
|
|
| for b in range(out_shape[0]):
|
| posmask = img_gt[b].astype(np.bool)
|
| if posmask.any():
|
| negmask = ~posmask
|
| posdis = distance(posmask)
|
| negdis = distance(negmask)
|
| boundary = skimage_seg.find_boundaries(posmask, mode="inner").astype(
|
| np.uint8
|
| )
|
| sdf = (negdis - np.min(negdis)) / (np.max(negdis) - np.min(negdis)) - (
|
| posdis - np.min(posdis)
|
| ) / (np.max(posdis) - np.min(posdis))
|
| sdf[boundary == 1] = 0
|
| normalized_sdf[b] = sdf
|
|
|
|
|
|
|
| return normalized_sdf
|
|
|
|
|
|
|
| def distributed_setup(rank, world_size):
|
| os.environ["MASTER_ADDR"] = "localhost"
|
| os.environ["MASTER_PORT"] = "12355"
|
| print("setting up dist process group now")
|
| dist.init_process_group("nccl", rank=rank, world_size=world_size)
|
|
|
|
|
| def load_ddp_to_nddp(state_dict):
|
| pattern = re.compile("module")
|
| for k, v in state_dict.items():
|
| if re.search("module", k):
|
| model_dict[re.sub(pattern, "", k)] = v
|
| else:
|
| model_dict = state_dict
|
| return model_dict
|
|
|
| def save_osmancheckpoint(epoch, ema_model, model, optimizer, best_dice, best_iter, stu_best_dice, stu_best_iter, path):
|
| """Saves model as checkpoint"""
|
| torch.save(
|
| {
|
| "epoch": epoch,
|
| "ema_state_dict": ema_model.state_dict(),
|
| "state_dict": model.state_dict(),
|
| "optimizer_state_dict": optimizer.state_dict(),
|
| "best_dice": best_dice,
|
| "best_iter": best_iter,
|
| "stu_best_dice": stu_best_dice,
|
| "stu_best_iter": stu_best_iter
|
| },
|
| path,
|
| )
|
|
|
| def load_osmancheckpoint(path, ema_model, model, optimizer, from_ddp=False):
|
| """loads previous checkpoint
|
|
|
| Args:
|
| path (str): path to checkpoint
|
| model (model): model to restore checkpoint to
|
| optimizer (optimizer): torch optimizer to load optimizer state_dict to
|
| from_ddp (bool, optional): load DistributedDataParallel checkpoint to regular model. Defaults to False.
|
|
|
| Returns:
|
| model, optimizer, epoch_num, loss
|
| """
|
|
|
| checkpoint = torch.load(path)
|
|
|
| ema_model.load_state_dict(checkpoint["ema_state_dict"])
|
| model.load_state_dict(checkpoint["state_dict"])
|
|
|
| optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
|
|
|
| best_dice, best_iter = checkpoint['best_dice'], checkpoint['best_iter']
|
| stu_best_dice, stu_best_iter = checkpoint['stu_best_dice'], checkpoint['stu_best_iter']
|
| return checkpoint["epoch"], ema_model, model, optimizer, best_dice, best_iter, stu_best_dice, stu_best_iter
|
|
|
| def draw_contour(
|
| image: ndarray,
|
| label: ndarray,
|
| color: tuple[int, int, int],
|
| ) -> ndarray:
|
| """Draw contour of label on image with color.
|
|
|
| Uses `cv2.dilate` to find contour.
|
|
|
| Args:
|
| image: ndarray of shape (h, w, 3)
|
| label: ndarray of shape (h, w) where non-zero values are foreground
|
| color: color of contour
|
| """
|
| binary = (label > 0).astype('uint8') * 255
|
| dilated = cv2.dilate(binary, np.ones((3, 3), dtype='uint8'), iterations=1)
|
| contour_mask = (dilated - binary) > 0
|
| image[contour_mask] = color
|
| return image
|
|
|
|
|
| def make_prediction(
|
| image: ndarray,
|
| *predictions: tuple[ndarray, tuple[int, int, int]],
|
| ) -> ndarray:
|
| """Draw contour of predictions on image.
|
|
|
| Args:
|
| image: ndarray of shape (h, w) or (h, w, 3)
|
| if image is (h, w), it will be repeated to (h, w, 3)
|
| predictions: each prediction is a tuple of (label, color)
|
| label: ndarray of shape (c, h, w) or (h, w)
|
| where c is number of classes or 1
|
| color: color of contour
|
| """
|
| if image.ndim == 2:
|
|
|
| image = np.repeat(image[..., np.newaxis], 3, axis=-1)
|
|
|
| for prediction, color in predictions:
|
| if prediction.ndim == 2:
|
|
|
| prediction = prediction[np.newaxis, ...]
|
| for class_map in prediction:
|
| image = draw_contour(image, class_map, color)
|
| return image
|
|
|
| def draw_contour_and_save(img, pred, mask, save_path='./img/1/example.png'):
|
| color_gt = (255,0,0)
|
| color_pred_list = [(0,255,0),(0,0,255),(255,255,0),(255,0,255)]
|
| if len(pred.shape) == 2:
|
| pred = pred.unsqueeze(0)
|
| if len(mask.shape) == 2:
|
| mask = mask.unsqueeze(0)
|
| if len(img.shape) == 2:
|
| img = img.unsqueeze(0)
|
| if len(img) == 1:
|
| img = img.repeat(3,1,1)
|
| img = (img-img.min())/(img.max()-img.min()) * 255
|
| img = img.cpu().numpy().transpose(1,2,0)
|
| classes = len(pred)
|
| for i in range(classes):
|
| predi = pred[i].float().cpu().numpy()
|
| maski = mask[i].float().cpu().numpy()
|
| seg_res = make_prediction(img, (predi, color_pred_list[i]), (maski, color_gt))
|
| seg_res = cv2.cvtColor(seg_res, cv2.COLOR_RGB2BGR)
|
| cv2.imwrite(save_path, seg_res)
|
|
|
| def draw_mask_and_save(img, pred, save_path='./img/1/example.png'):
|
| color_list = [(255,0,0),(0,255,0),(0,0,255),(255,255,0),(255,0,255)]
|
| if len(pred.shape) == 2:
|
| pred = pred.unsqueeze(0)
|
| if len(img.shape) == 2:
|
| img = img.unsqueeze(0)
|
| if len(img) == 1:
|
| img = img.repeat(3,1,1)
|
| d, h, w = img.shape
|
| rgb = np.zeros((h, w, 3))
|
| mask = np.ones((h, w, 3))
|
| if img.min() < -0.5:
|
| img = (img+1)*127.5
|
| elif img.max() < 1.5:
|
| img = img*255
|
| img = img.cpu().numpy().transpose(1,2,0)
|
| classes = len(pred)
|
| for i in reversed(range(classes)):
|
| mask[pred[i] == 1] = (0.5, 0.5, 0.5)
|
| rgb[pred[i] == 1] = color_list[i]
|
| res = (img+rgb)*mask
|
| res = res.astype(np.uint8)
|
| res = cv2.cvtColor(res, cv2.COLOR_RGB2BGR)
|
| cv2.imwrite(save_path, res)
|
|
|
