|
import itertools |
|
import os |
|
import random |
|
import shutil |
|
from math import ceil |
|
from typing import Dict, List |
|
|
|
import numpy as np |
|
import prettytable as pt |
|
import torch |
|
import torch.nn as nn |
|
from fast_pytorch_kmeans import KMeans |
|
from pathlib import Path |
|
from scipy.stats import hmean |
|
from sklearn import metrics |
|
from termcolor import cprint |
|
from torchvision.utils import draw_segmentation_masks, make_grid, save_image |
|
|
|
import utils.misc as misc |
|
from losses import get_spixel_tgt_map, get_volume_seg_map |
|
from utils.convcrf import convcrf |
|
from utils.crf import DenseCRF |
|
|
|
|
|
def train( |
|
model: nn.Module, |
|
dataloader, |
|
dataset_title: str, |
|
optimizer_dict: Dict, |
|
criterion, |
|
epoch: int, |
|
writer, |
|
suffix: str, |
|
opt, |
|
): |
|
|
|
metric_logger = misc.MetricLogger(writer=writer, suffix=suffix) |
|
cprint("{}-th epoch training on {}".format(epoch, dataset_title), "blue") |
|
model.train() |
|
roc_auc_elements = { |
|
modality: {"map_scores": [], "vol_scores": []} |
|
for modality in itertools.chain(opt.modality, ["ensemble"]) |
|
} |
|
roc_auc_elements["labels"] = [] |
|
|
|
for i, data in metric_logger.log_every( |
|
dataloader, print_freq=opt.print_freq, header=f"[{suffix} {epoch}]" |
|
): |
|
if (opt.debug or opt.wholetest) and i > 50: |
|
break |
|
|
|
for modality, optimizer in optimizer_dict.items(): |
|
optimizer.zero_grad() |
|
|
|
image = data["image"].to(opt.device) |
|
unnormalized_image = data["unnormalized_image"].to(opt.device) |
|
label = data["label"].to(opt.device) |
|
mask = data["mask"].to(opt.device) |
|
spixel = data["spixel"].to(opt.device) if opt.mvc_spixel else None |
|
|
|
outputs = model( |
|
image, |
|
seg_size=None |
|
if opt.loss_on_mid_map |
|
else [image.shape[-2], image.shape[-1]], |
|
) |
|
|
|
losses = criterion( |
|
outputs, |
|
label, |
|
mask, |
|
epoch=epoch, |
|
max_epoch=opt.epochs, |
|
spixel=spixel, |
|
raw_image=unnormalized_image, |
|
) |
|
total_loss = losses["total_loss"] |
|
total_loss.backward() |
|
|
|
for modality in opt.modality: |
|
if opt.grad_clip > 0.0: |
|
grad_norm = nn.utils.clip_grad_norm_( |
|
model.sub_models[modality].parameters(), opt.grad_clip |
|
) |
|
metric_logger.update(**{f"grad_norm/{modality}": grad_norm}) |
|
|
|
optimizer_dict[modality].step() |
|
|
|
|
|
roc_auc_elements["labels"].extend(label.tolist()) |
|
for modality in itertools.chain(opt.modality, ["ensemble"]): |
|
roc_auc_elements[modality]["map_scores"].extend( |
|
outputs[modality]["map_pred"].tolist() |
|
) |
|
roc_auc_elements[modality]["vol_scores"].extend( |
|
(outputs[modality]["vol_pred"]).tolist() |
|
) |
|
|
|
metric_logger.update(**losses) |
|
|
|
image_metrics = update_image_roc_auc_metric( |
|
opt.modality + ["ensemble"], roc_auc_elements, None |
|
) |
|
metric_logger.update(**image_metrics) |
|
|
|
metric_logger.write_tensorboard(epoch) |
|
print("Average status:") |
|
print(metric_logger.stat_table()) |
|
|
|
|
|
def bundled_evaluate( |
|
model: nn.Module, dataloaders: Dict, criterion, epoch, writer, suffix, opt |
|
): |
|
|
|
metric_logger = misc.MetricLogger(writer=writer, suffix=suffix + "_avg") |
|
for dataset, dataloader in dataloaders.items(): |
|
outputs = evaluate( |
|
model, |
|
dataloader, |
|
criterion, |
|
dataset, |
|
epoch, |
|
writer, |
|
suffix + f"_{dataset}", |
|
opt, |
|
) |
|
old_keys = list(outputs.keys()) |
|
for k in old_keys: |
|
outputs[k.replace(dataset.upper(), "AVG")] = outputs[k] |
|
for k in old_keys: |
|
del outputs[k] |
|
|
|
metric_logger.update(**outputs) |
|
|
|
metric_logger.write_tensorboard(epoch) |
|
print("Average status:") |
|
print(metric_logger.stat_table()) |
|
return metric_logger.get_meters() |
|
|
|
|
|
def evaluate( |
|
model: nn.Module, |
|
dataloader, |
|
criterion, |
|
dataset_title: str, |
|
epoch: int, |
|
writer, |
|
suffix: str, |
|
opt, |
|
): |
|
|
|
metric_logger = misc.MetricLogger(writer=writer, suffix=suffix) |
|
cprint("{}-th epoch evaluation on {}".format(epoch, dataset_title.upper()), "blue") |
|
|
|
model.eval() |
|
|
|
if opt.crf_postproc: |
|
postprocess = DenseCRF( |
|
iter_max=opt.crf_iter_max, |
|
pos_w=opt.crf_pos_w, |
|
pos_xy_std=opt.crf_pos_xy_std, |
|
bi_w=opt.crf_bi_w, |
|
bi_xy_std=opt.crf_bi_xy_std, |
|
bi_rgb_std=opt.crf_bi_rgb_std, |
|
) |
|
elif opt.convcrf_postproc: |
|
convcrf_config = convcrf.default_conf |
|
convcrf_config["skip_init_softmax"] = True |
|
convcrf_config["final_softmax"] = True |
|
shape = [opt.convcrf_shape, opt.convcrf_shape] |
|
postprocess = convcrf.GaussCRF( |
|
conf=convcrf_config, shape=shape, nclasses=2, use_gpu=True |
|
).to(opt.device) |
|
|
|
figure_path = opt.figure_path + f"_{dataset_title.upper()}" |
|
if opt.save_figure: |
|
if os.path.exists(figure_path): |
|
shutil.rmtree(figure_path) |
|
os.mkdir(figure_path) |
|
cprint("Saving figures to {}".format(figure_path), "blue") |
|
|
|
if opt.max_pool_postproc > 1: |
|
max_pool = nn.MaxPool2d( |
|
kernel_size=opt.max_pool_postproc, |
|
stride=1, |
|
padding=(opt.max_pool_postproc - 1) // 2, |
|
).to(opt.device) |
|
else: |
|
max_pool = nn.Identity().to(opt.device) |
|
|
|
roc_auc_elements = { |
|
modality: {"map_scores": [], "vol_scores": []} |
|
for modality in itertools.chain(opt.modality, ["ensemble"]) |
|
} |
|
roc_auc_elements["labels"] = [] |
|
with torch.no_grad(): |
|
for i, data in metric_logger.log_every( |
|
dataloader, print_freq=opt.print_freq, header=f"[{suffix} {epoch}]" |
|
): |
|
if (opt.debug or opt.wholetest) and i > 50: |
|
break |
|
|
|
image_size = data["image"].shape[-2:] |
|
label = data["label"] |
|
mask = data["mask"] |
|
if opt.crf_postproc or opt.spixel_postproc or opt.convcrf_postproc: |
|
spixel = data["spixel"].to(opt.device) |
|
if max(image_size) > opt.tile_size and opt.large_image_strategy == "slide": |
|
outputs = sliding_predict( |
|
model, data, opt.tile_size, opt.tile_overlap, opt |
|
) |
|
else: |
|
image = data["image"].to(opt.device) |
|
outputs = model(image, seg_size=image.shape[-2:]) |
|
|
|
if opt.max_pool_postproc > 1: |
|
for modality in itertools.chain(opt.modality, ["ensemble"]): |
|
outputs[modality]["out_map"] = max_pool( |
|
outputs[modality]["out_map"] |
|
) |
|
|
|
if opt.crf_postproc: |
|
raw_prob = outputs["ensemble"]["out_map"] |
|
image = data["unnormalized_image"] * 255.0 |
|
if opt.crf_downsample > 1: |
|
image = ( |
|
torch.nn.functional.interpolate( |
|
image, |
|
size=( |
|
image_size[0] // opt.crf_downsample, |
|
image_size[1] // opt.crf_downsample, |
|
), |
|
mode="bilinear", |
|
align_corners=False, |
|
) |
|
.clamp(0, 255) |
|
.int() |
|
) |
|
image = image.squeeze(0).numpy().astype(np.uint8).transpose(1, 2, 0) |
|
for modality in itertools.chain(opt.modality, ["ensemble"]): |
|
prob = outputs[modality]["out_map"].squeeze(1) |
|
if opt.crf_downsample > 1: |
|
prob = ( |
|
torch.nn.functional.interpolate( |
|
prob, |
|
size=( |
|
image_size[0] // opt.crf_downsample, |
|
image_size[1] // opt.crf_downsample, |
|
), |
|
mode="bilinear", |
|
align_corners=False, |
|
) |
|
.clamp(0, 1) |
|
.squeeze(0) |
|
) |
|
prob = torch.cat([prob, 1 - prob], dim=0).detach().cpu().numpy() |
|
prob = postprocess(image, prob) |
|
prob = prob[None, 0, ...] |
|
prob = torch.tensor(prob, device=opt.device).unsqueeze(0) |
|
if opt.crf_downsample > 1: |
|
prob = torch.nn.functional.interpolate( |
|
prob, size=image_size, mode="bilinear", align_corners=False |
|
).clamp(0, 1) |
|
outputs[modality]["out_map"] = prob |
|
outputs[modality]["map_pred"] = ( |
|
outputs[modality]["out_map"].max().unsqueeze(0) |
|
) |
|
elif opt.convcrf_postproc: |
|
raw_prob = outputs["ensemble"]["out_map"] |
|
image = data["unnormalized_image"].to(opt.device) * 255.0 |
|
image = ( |
|
torch.nn.functional.interpolate( |
|
image, |
|
size=(opt.convcrf_shape, opt.convcrf_shape), |
|
mode="bilinear", |
|
align_corners=False, |
|
) |
|
.clamp(0, 255) |
|
.int() |
|
) |
|
for modality in itertools.chain(opt.modality, ["ensemble"]): |
|
prob = outputs[modality]["out_map"] |
|
prob = torch.cat([prob, 1 - prob], dim=1) |
|
prob = torch.nn.functional.interpolate( |
|
prob, |
|
size=(opt.convcrf_shape, opt.convcrf_shape), |
|
mode="bilinear", |
|
align_corners=False, |
|
).clamp(0, 1) |
|
prob = postprocess(unary=prob, img=image) |
|
prob = torch.nn.functional.interpolate( |
|
prob, size=image_size, mode="bilinear", align_corners=False |
|
).clamp(0, 1) |
|
outputs[modality]["out_map"] = prob[:, 0, None, ...] |
|
outputs[modality]["map_pred"] = ( |
|
outputs[modality]["out_map"].max().unsqueeze(0) |
|
) |
|
elif opt.spixel_postproc: |
|
raw_prob = outputs["ensemble"]["out_map"] |
|
for modality in itertools.chain(opt.modality, ["ensemble"]): |
|
outputs[modality]["out_map"] = get_spixel_tgt_map( |
|
outputs[modality]["out_map"], spixel |
|
) |
|
|
|
|
|
roc_auc_elements["labels"].extend(label.detach().cpu().tolist()) |
|
for modality in itertools.chain(opt.modality, ["ensemble"]): |
|
roc_auc_elements[modality]["map_scores"].extend( |
|
outputs[modality]["map_pred"].detach().cpu().tolist() |
|
) |
|
roc_auc_elements[modality]["vol_scores"].extend( |
|
(outputs[modality]["vol_pred"]).detach().cpu().tolist() |
|
) |
|
|
|
|
|
out_map = { |
|
modality: outputs[modality]["out_map"] > opt.mask_threshold |
|
for modality in itertools.chain(opt.modality, ["ensemble"]) |
|
} |
|
|
|
|
|
if label.item() == 1.0: |
|
for modality in itertools.chain(opt.modality, ["ensemble"]): |
|
pixel_metrics = misc.calculate_pixel_f1( |
|
out_map[modality].float().detach().cpu().numpy().flatten(), |
|
mask.detach().cpu().numpy().flatten(), |
|
suffix=f"/{modality}", |
|
) |
|
metric_logger.update(**pixel_metrics) |
|
|
|
|
|
if opt.save_figure: |
|
unnormalized_image = data["unnormalized_image"] |
|
|
|
image_id = Path(data["id"][0]).stem |
|
save_image( |
|
( |
|
outputs["ensemble"]["out_map"][0, ...] > opt.mask_threshold |
|
).float() |
|
* 255, |
|
os.path.join(figure_path, f"{image_id}_ensemble_map.png"), |
|
) |
|
|
|
image_metrics = update_image_roc_auc_metric( |
|
opt.modality + ["ensemble"], |
|
roc_auc_elements, |
|
{ |
|
modality: metric_logger.meters[f"pixel_f1/{modality}"].avg |
|
for modality in itertools.chain(opt.modality, ["ensemble"]) |
|
}, |
|
) |
|
metric_logger.update(**image_metrics) |
|
|
|
metric_logger.prepend_subprefix(f"{dataset_title.upper()}_") |
|
metric_logger.write_tensorboard(epoch) |
|
print("Average status:") |
|
print(metric_logger.stat_table()) |
|
|
|
return metric_logger.get_meters() |
|
|
|
|
|
def update_image_roc_auc_metric(modalities: List, roc_auc_elements, pixel_f1=None): |
|
|
|
result = {} |
|
for modality in modalities: |
|
image_metrics = misc.calculate_img_score( |
|
np.array(roc_auc_elements[modality]["map_scores"]) > 0.5, |
|
(np.array(roc_auc_elements["labels"]) > 0).astype(np.int), |
|
suffix=f"/{modality}", |
|
) |
|
if pixel_f1 is not None: |
|
image_f1 = image_metrics[f"image_f1/{modality}"] |
|
combined_f1 = hmean([image_f1, pixel_f1[modality]]) |
|
image_metrics[f"comb_f1/{modality}"] = float(combined_f1) |
|
if 0.0 in roc_auc_elements["labels"] and 1.0 in roc_auc_elements["labels"]: |
|
image_auc = metrics.roc_auc_score( |
|
roc_auc_elements["labels"], roc_auc_elements[modality]["map_scores"] |
|
) |
|
image_metrics[f"image_auc/{modality}"] = image_auc |
|
result.update(image_metrics) |
|
|
|
return result |
|
|
|
|
|
def pad_image(image, target_size): |
|
image_size = image.shape[-2:] |
|
if image_size != target_size: |
|
row_missing = target_size[0] - image_size[0] |
|
col_missing = target_size[1] - image_size[1] |
|
image = nn.functional.pad( |
|
image, (0, row_missing, 0, col_missing), "constant", 0 |
|
) |
|
return image |
|
|
|
|
|
def sliding_predict(model: nn.Module, data, tile_size, tile_overlap, opt): |
|
image = data["image"] |
|
mask = data["mask"] |
|
image = image.to(opt.device) |
|
image_size = image.shape[-2:] |
|
stride = ceil(tile_size * (1 - tile_overlap)) |
|
tile_rows = int(ceil((image_size[0] - tile_size) / stride) + 1) |
|
tile_cols = int(ceil((image_size[1] - tile_size) / stride) + 1) |
|
result = {} |
|
for modality in itertools.chain(opt.modality, ["ensemble"]): |
|
result[modality] = { |
|
"out_map": torch.zeros_like( |
|
mask, requires_grad=False, dtype=torch.float32, device=opt.device |
|
), |
|
"out_vol_map": torch.zeros_like( |
|
mask, requires_grad=False, dtype=torch.float32, device=opt.device |
|
), |
|
} |
|
map_counter = torch.zeros_like( |
|
mask, requires_grad=False, dtype=torch.float32, device=opt.device |
|
) |
|
|
|
with torch.no_grad(): |
|
for row in range(tile_rows): |
|
for col in range(tile_cols): |
|
x1 = int(col * stride) |
|
y1 = int(row * stride) |
|
x2 = min(x1 + tile_size, image_size[1]) |
|
y2 = min(y1 + tile_size, image_size[0]) |
|
x1 = max(int(x2 - tile_size), 0) |
|
y1 = max(int(y2 - tile_size), 0) |
|
|
|
image_tile = image[:, :, y1:y2, x1:x2] |
|
image_tile = pad_image(image_tile, [opt.tile_size, opt.tile_size]) |
|
tile_outputs = model(image_tile, seg_size=(image_tile.shape[-2:])) |
|
for modality in itertools.chain(opt.modality, ["ensemble"]): |
|
result[modality]["out_map"][:, :, y1:y2, x1:x2] += tile_outputs[ |
|
modality |
|
]["out_map"][:, :, : y2 - y1, : x2 - x1] |
|
out_vol_map = get_volume_seg_map( |
|
tile_outputs[modality]["out_vol"], |
|
size=image_tile.shape[-2:], |
|
label=data["label"], |
|
kmeans=KMeans(2) if opt.consistency_kmeans else None, |
|
)[:, :, : y2 - y1, : x2 - x1] |
|
result[modality]["out_vol_map"][:, :, y1:y2, x1:x2] += out_vol_map |
|
map_counter[:, :, y1:y2, x1:x2] += 1 |
|
|
|
for modality in itertools.chain(opt.modality, ["ensemble"]): |
|
result[modality]["out_map"] /= map_counter |
|
result[modality]["out_vol_map"] /= map_counter |
|
result[modality]["map_pred"] = ( |
|
result[modality]["out_map"].max().unsqueeze(0) |
|
) |
|
result[modality]["vol_pred"] = ( |
|
result[modality]["out_vol_map"].max().unsqueeze(0) |
|
) |
|
|
|
return result |
|
|
