# -*- coding: utf-8 -*-
# CellViT Inference Method for Patch-Wise Inference on MoNuSeg dataset
#
# @ Fabian Hörst, fabian.hoerst@uk-essen.de
# Institute for Artifical Intelligence in Medicine,
# University Medicine Essen
import argparse
import inspect
import os
import sys
currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
parentdir = os.path.dirname(currentdir)
sys.path.insert(0, parentdir)
parentdir = os.path.dirname(parentdir)
sys.path.insert(0, parentdir)
from base_ml.base_experiment import BaseExperiment
BaseExperiment.seed_run(1232)
from pathlib import Path
from typing import List, Union, Tuple
import albumentations as A
import cv2 as cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import tqdm
from einops import rearrange
from matplotlib import pyplot as plt
from PIL import Image, ImageDraw
from skimage.color import rgba2rgb
from torch.utils.data import DataLoader
from torchmetrics.functional import dice
from torchmetrics.functional.classification import binary_jaccard_index
from torchvision import transforms
from cell_segmentation.datasets.monuseg import MoNuSegDataset
from cell_segmentation.inference.cell_detection import (
CellPostProcessor,
get_cell_position,
get_cell_position_marging,
get_edge_patch,
)
from cell_segmentation.utils.metrics import (
cell_detection_scores,
get_fast_pq,
remap_label,
)
from cell_segmentation.utils.post_proc_cellvit import calculate_instances
from cell_segmentation.utils.tools import pair_coordinates
from models.segmentation.cell_segmentation.cellvit import CellViT
from utils.logger import Logger
from utils.tools import unflatten_dict
class MoNuSegInference:
def __init__(
self,
model_path: Union[Path, str],
dataset_path: Union[Path, str],
outdir: Union[Path, str],
gpu: int,
patching: bool = False,
overlap: int = 0,
magnification: int = 40,
) -> None:
"""Cell Segmentation Inference class for MoNuSeg dataset
Args:
model_path (Union[Path, str]): Path to model checkpoint
dataset_path (Union[Path, str]): Path to dataset
outdir (Union[Path, str]): Output directory
gpu (int): CUDA GPU id to use
patching (bool, optional): If dataset should be pacthed to 256px. Defaults to False.
overlap (int, optional): If overlap should be used. Recommed (next to no overlap) is 64 px. Overlap in px.
If overlap is used, patching must be True. Defaults to 0.
magnification (int, optional): Dataset magnification. Defaults to 40.
"""
self.model_path = Path(model_path)
self.device = "cpu"
self.magnification = magnification
self.overlap = overlap
self.patching = patching
if overlap > 0:
assert patching, "Patching must be activated"
# self.__instantiate_logger()
self.__load_model()
# self.__load_inference_transforms()
self.__setup_amp()
def __instantiate_logger(self) -> None:
"""Instantiate logger
Logger is using no formatters. Logs are stored in the run directory under the filename: inference.log
"""
logger = Logger(
level="INFO",
log_dir=self.outdir,
comment="inference_monuseg",
use_timestamp=False,
formatter="%(message)s",
)
self.logger = logger.create_logger()
def __load_model(self) -> None:
"""Load model and checkpoint and load the state_dict"""
model_checkpoint = torch.load(self.model_path, map_location="cpu")
# unpack checkpoint
self.run_conf = unflatten_dict(model_checkpoint["config"], ".")
self.model = self.__get_model(model_type=model_checkpoint["arch"])
self.model.load_state_dict(model_checkpoint["model_state_dict"])
self.model.eval()
self.model.to(self.device)
def __get_model(
self, model_type: str
) -> Union[
CellViT,
]:
"""Return the trained model for inference
Args:
model_type (str): Name of the model. Must either be one of:
CellViT, CellViTShared, CellViT256, CellViT256Shared, CellViTSAM, CellViTSAMShared
Returns:
Union[CellViT, CellViTShared, CellViT256, CellViTShared, CellViTSAM, CellViTSAMShared]: Model
"""
implemented_models = [
"CellViT",
]
if model_type not in implemented_models:
raise NotImplementedError(
f"Unknown model type. Please select one of {implemented_models}"
)
if model_type in ["CellViT", "CellViTShared"]:
if model_type == "CellViT":
model_class = CellViT
model = model_class(
model256_path=self.run_conf["model"].get("pretrained_encoder"),
num_nuclei_classes=self.run_conf["data"]["num_nuclei_classes"],
num_tissue_classes=self.run_conf["data"]["num_tissue_classes"],
#embed_dim=self.run_conf["model"]["embed_dim"],
in_channels=self.run_conf["model"].get("input_channels", 3),
#depth=self.run_conf["model"]["depth"],
#num_heads=self.run_conf["model"]["num_heads"],
#extract_layers=self.run_conf["model"]["extract_layers"],
#regression_loss=self.run_conf["model"].get("regression_loss", False),
)
return model
def __load_inference_transforms(self) -> None:
"""Load the inference transformations from the run_configuration"""
self.logger.info("Loading inference transformations")
transform_settings = self.run_conf["transformations"]
if "normalize" in transform_settings:
mean = transform_settings["normalize"].get("mean", (0.5, 0.5, 0.5))
std = transform_settings["normalize"].get("std", (0.5, 0.5, 0.5))
else:
mean = (0.5, 0.5, 0.5)
std = (0.5, 0.5, 0.5)
self.inference_transforms = A.Compose([A.Normalize(mean=mean, std=std)])
def __setup_amp(self) -> None:
"""Setup automated mixed precision (amp) for inference."""
self.mixed_precision = self.run_conf["training"].get("mixed_precision", False)
def run_inference(self, generate_plots: bool = False) -> None:
"""Run inference
Args:
generate_plots (bool, optional): If plots should be generated. Defaults to False.
"""
self.model.eval()
# setup score tracker
image_names = [] # image names as str
binary_dice_scores = [] # binary dice scores per image
binary_jaccard_scores = [] # binary jaccard scores per image
pq_scores = [] # pq-scores per image
dq_scores = [] # dq-scores per image
sq_scores = [] # sq-scores per image
f1_ds = [] # f1-scores per image
prec_ds = [] # precision per image
rec_ds = [] # recall per image
inference_loop = tqdm.tqdm(
enumerate(self.inference_dataloader), total=len(self.inference_dataloader)
)
with torch.no_grad():
for image_idx, batch in inference_loop:
image_metrics = self.inference_step(
model=self.model, batch=batch, generate_plots=generate_plots
)
image_names.append(image_metrics["image_name"])
binary_dice_scores.append(image_metrics["binary_dice_score"])
binary_jaccard_scores.append(image_metrics["binary_jaccard_score"])
pq_scores.append(image_metrics["pq_score"])
dq_scores.append(image_metrics["dq_score"])
sq_scores.append(image_metrics["sq_score"])
f1_ds.append(image_metrics["f1_d"])
prec_ds.append(image_metrics["prec_d"])
rec_ds.append(image_metrics["rec_d"])
# average metrics for dataset
binary_dice_scores = np.array(binary_dice_scores)
binary_jaccard_scores = np.array(binary_jaccard_scores)
pq_scores = np.array(pq_scores)
dq_scores = np.array(dq_scores)
sq_scores = np.array(sq_scores)
f1_ds = np.array(f1_ds)
prec_ds = np.array(prec_ds)
rec_ds = np.array(rec_ds)
dataset_metrics = {
"Binary-Cell-Dice-Mean": float(np.nanmean(binary_dice_scores)),
"Binary-Cell-Jacard-Mean": float(np.nanmean(binary_jaccard_scores)),
"bPQ": float(np.nanmean(pq_scores)),
"bDQ": float(np.nanmean(dq_scores)),
"bSQ": float(np.nanmean(sq_scores)),
"f1_detection": float(np.nanmean(f1_ds)),
"precision_detection": float(np.nanmean(prec_ds)),
"recall_detection": float(np.nanmean(rec_ds)),
}
self.logger.info(f"{20*'*'} Binary Dataset metrics {20*'*'}")
[self.logger.info(f"{f'{k}:': <25} {v}") for k, v in dataset_metrics.items()]
def inference_step(
self, model: nn.Module, batch: object, generate_plots: bool = False
) -> dict:
"""Inference step
Args:
model (nn.Module): Training model, must return "nuclei_binary_map", "nuclei_type_map", "tissue_type" and "hv_map"
batch (object): Training batch, consisting of images ([0]), masks ([1]), tissue_types ([2]) and figure filenames ([3])
generate_plots (bool, optional): If plots should be generated. Defaults to False.
Returns:
Dict: Image_metrics with keys:
"""
img = batch[0].to(self.device)
if len(img.shape) > 4:
img = img[0]
img = rearrange(img, "c i j w h -> (i j) c w h")
mask = batch[1]
image_name = list(batch[2])
mask["instance_types"] = calculate_instances(
torch.unsqueeze(mask["nuclei_binary_map"], dim=0), mask["instance_map"]
)
model.zero_grad()
if self.mixed_precision:
with torch.autocast(device_type="cuda", dtype=torch.float16):
predictions_ = model.forward(img)
else:
predictions_ = model.forward(img)
if self.overlap == 0:
if self.patching:
predictions_ = self.post_process_patching(predictions_)
predictions = self.get_cell_predictions(predictions_)
image_metrics = self.calculate_step_metric(
predictions=predictions, gt=mask, image_name=image_name
)
elif self.patching and self.overlap != 0:
cell_list = self.post_process_patching_overlap(
predictions_, overlap=self.overlap
)
image_metrics, predictions = self.calculate_step_metric_overlap(
cell_list=cell_list, gt=mask, image_name=image_name
)
scores = [
float(image_metrics["binary_dice_score"].detach().cpu()),
float(image_metrics["binary_jaccard_score"].detach().cpu()),
image_metrics["pq_score"],
]
if generate_plots:
if self.overlap == 0 and self.patching:
batch_size = img.shape[0]
num_elems = int(np.sqrt(batch_size))
img = torch.permute(img, (0, 2, 3, 1))
img = rearrange(
img, "(i j) h w c -> (i h) (j w) c", i=num_elems, j=num_elems
)
img = torch.unsqueeze(img, dim=0)
img = torch.permute(img, (0, 3, 1, 2))
elif self.overlap != 0 and self.patching:
h, w = mask["nuclei_binary_map"].shape[1:]
total_img = torch.zeros((3, h, w))
decomposed_patch_num = int(np.sqrt(img.shape[0]))
for i in range(decomposed_patch_num):
for j in range(decomposed_patch_num):
x_global = i * 256 - i * self.overlap
y_global = j * 256 - j * self.overlap
total_img[
:, x_global : x_global + 256, y_global : y_global + 256
] = img[i * decomposed_patch_num + j]
img = total_img
img = img[None, :, :, :]
self.plot_results(
img=img,
predictions=predictions,
ground_truth=mask,
img_name=image_name[0],
scores=scores,
)
return image_metrics
def run_single_image_inference(self, model: nn.Module, image: np.ndarray, generate_plots: bool = True,
) -> dict:
"""Inference step
Args:
model (nn.Module): Training model, must return "nuclei_binary_map", "nuclei_type_map", "tissue_type" and "hv_map"
batch (object): Training batch, consisting of images ([0]), masks ([1]), tissue_types ([2]) and figure filenames ([3])
generate_plots (bool, optional): If plots should be generated. Defaults to False.
Returns:
Dict: Image_metrics with keys:
"""
# set image transforms
transform_settings = self.run_conf["transformations"]
if "normalize" in transform_settings:
mean = transform_settings["normalize"].get("mean", (0.5, 0.5, 0.5))
std = transform_settings["normalize"].get("std", (0.5, 0.5, 0.5))
else:
mean = (0.5, 0.5, 0.5)
std = (0.5, 0.5, 0.5)
transforms = A.Compose([A.Normalize(mean=mean, std=std)])
transformed_img = transforms(image=image)["image"]
image = torch.from_numpy(transformed_img).permute(2, 0, 1).unsqueeze(0).float()
img = image.to(self.device)
model.zero_grad()
predictions_ = model.forward(img)
if self.overlap == 0:
if self.patching:
predictions_ = self.post_process_patching(predictions_)
predictions = self.get_cell_predictions(predictions_)
image_output = self.plot_results(
img=img,
predictions=predictions
)
return image_output
def calculate_step_metric(
self, predictions: dict, gt: dict, image_name: List[str]
) -> dict:
"""Calculate step metric for one MoNuSeg image.
Args:
predictions (dict): Necssary keys:
* instance_map: Pixel-wise nuclear instance segmentation.
Each instance has its own integer, starting from 1. Shape: (1, H, W)
* nuclei_binary_map: Softmax output for binary nuclei branch. Shape: (1, 2, H, W)
* instance_types: Instance type prediction list.
Each list entry stands for one image. Each list entry is a dictionary with the following structure:
Main Key is the nuclei instance number (int), with a dict as value.
For each instance, the dictionary contains the keys: bbox (bounding box), centroid (centroid coordinates),
contour, type_prob (probability), type (nuclei type). Actually just one list entry, as we expecting batch-size=1 (one image)
gt (dict): Necessary keys:
* instance_map
* nuclei_binary_map
* instance_types
image_name (List[str]): Name of the image, list with [str]. List is necessary for backward compatibility
Returns:
dict: Image metrics for one MoNuSeg image. Keys are:
* image_name
* binary_dice_score
* binary_jaccard_score
* pq_score
* dq_score
* sq_score
* f1_d
* prec_d
* rec_d
"""
predictions["instance_map"] = predictions["instance_map"].detach().cpu()
instance_maps_gt = gt["instance_map"].detach().cpu()
pred_binary_map = torch.argmax(predictions["nuclei_binary_map"], dim=1)
target_binary_map = gt["nuclei_binary_map"].to(self.device)
cell_dice = (
dice(preds=pred_binary_map, target=target_binary_map, ignore_index=0)
.detach()
.cpu()
)
cell_jaccard = (
binary_jaccard_index(
preds=pred_binary_map,
target=target_binary_map,
)
.detach()
.cpu()
)
remapped_instance_pred = remap_label(predictions["instance_map"])
remapped_gt = remap_label(instance_maps_gt)
[dq, sq, pq], _ = get_fast_pq(true=remapped_gt, pred=remapped_instance_pred)
# detection scores
true_centroids = np.array(
[v["centroid"] for k, v in gt["instance_types"][0].items()]
)
pred_centroids = np.array(
[v["centroid"] for k, v in predictions["instance_types"].items()]
)
if true_centroids.shape[0] == 0:
true_centroids = np.array([[0, 0]])
if pred_centroids.shape[0] == 0:
pred_centroids = np.array([[0, 0]])
if self.magnification == 40:
pairing_radius = 12
else:
pairing_radius = 6
paired, unpaired_true, unpaired_pred = pair_coordinates(
true_centroids, pred_centroids, pairing_radius
)
f1_d, prec_d, rec_d = cell_detection_scores(
paired_true=paired[:, 0],
paired_pred=paired[:, 1],
unpaired_true=unpaired_true,
unpaired_pred=unpaired_pred,
)
image_metrics = {
"image_name": image_name,
"binary_dice_score": cell_dice,
"binary_jaccard_score": cell_jaccard,
"pq_score": pq,
"dq_score": dq,
"sq_score": sq,
"f1_d": f1_d,
"prec_d": prec_d,
"rec_d": rec_d,
}
return image_metrics
def convert_binary_type(self, instance_types: dict) -> dict:
"""Clean nuclei detection from type prediction to binary prediction
Args:
instance_types (dict): Dictionary with cells
Returns:
dict: Cleaned with just one class
"""
cleaned_instance_types = {}
for key, elem in instance_types.items():
if elem["type"] == 0:
continue
else:
elem["type"] = 0
cleaned_instance_types[key] = elem
return cleaned_instance_types
def get_cell_predictions(self, predictions: dict) -> dict:
"""Reshaping predictions and calculating instance maps and instance types
Args:
predictions (dict): Dictionary with the following keys:
* tissue_types: Logit tissue prediction output. Shape: (B, num_tissue_classes)
* nuclei_binary_map: Logit output for binary nuclei prediction branch. Shape: (B, H, W, 2)
* hv_map: Logit output for hv-prediction. Shape: (B, 2, H, W)
* nuclei_type_map: Logit output for nuclei instance-prediction. Shape: (B, num_nuclei_classes, H, W)
Returns:
dict:
* nuclei_binary_map: Softmax binary prediction. Shape: (B, 2, H, W
* nuclei_type_map: Softmax nuclei type map. Shape: (B, num_nuclei_classes, H, W)
* hv_map: Logit output for hv-prediction. Shape: (B, 2, H, W)
* tissue_types: Logit tissue prediction output. Shape: (B, num_tissue_classes)
* instance_map: Instance map, each instance has one integer. Shape: (B, H, W)
* instance_types: Instance type dict, cleaned. Keys:
'bbox', 'centroid', 'contour', 'type_prob', 'type'
"""
predictions["nuclei_binary_map"] = F.softmax(
predictions["nuclei_binary_map"], dim=1
)
predictions["nuclei_type_map"] = F.softmax(
predictions["nuclei_type_map"], dim=1
)
(
predictions["instance_map"],
predictions["instance_types"],
) = self.model.calculate_instance_map(
predictions, magnification=self.magnification
)
predictions["instance_types"] = self.convert_binary_type(
predictions["instance_types"][0]
)
return predictions
def post_process_patching(self, predictions: dict) -> dict:
"""Post-process patching by reassamble (without overlap) stitched predictions to one big image prediction
Args:
predictions (dict): Necessary keys:
* nuclei_binary_map: Logit binary prediction. Shape: (B, 2, 256, 256)
* hv_map: Logit output for hv-prediction. Shape: (B, 2, H, W)
* nuclei_type_map: Logit output for nuclei instance-prediction. Shape: (B, num_nuclei_classes, 256, 256)
Returns:
dict: Return elements that have been changed:
* nuclei_binary_map: Shape: (1, 2, H, W)
* hv_map: Shape: (1, 2, H, W)
* nuclei_type_map: (1, num_nuclei_classes, H, W)
"""
batch_size = predictions["nuclei_binary_map"].shape[0]
num_elems = int(np.sqrt(batch_size))
predictions["nuclei_binary_map"] = rearrange(
predictions["nuclei_binary_map"],
"(i j) d w h ->d (i w) (j h)",
i=num_elems,
j=num_elems,
)
predictions["hv_map"] = rearrange(
predictions["hv_map"],
"(i j) d w h -> d (i w) (j h)",
i=num_elems,
j=num_elems,
)
predictions["nuclei_type_map"] = rearrange(
predictions["nuclei_type_map"],
"(i j) d w h -> d (i w) (j h)",
i=num_elems,
j=num_elems,
)
predictions["nuclei_binary_map"] = torch.unsqueeze(
predictions["nuclei_binary_map"], dim=0
)
predictions["hv_map"] = torch.unsqueeze(predictions["hv_map"], dim=0)
predictions["nuclei_type_map"] = torch.unsqueeze(
predictions["nuclei_type_map"], dim=0
)
return predictions
def post_process_patching_overlap(self, predictions: dict, overlap: int) -> List:
"""Post processing overlapping cells by merging overlap. Use same merging strategy as for our
Args:
predictions (dict): Predictions with necessary keys:
* nuclei_binary_map: Binary nuclei prediction, Shape: (B, 2, H, W)
* nuclei_type_map: Nuclei type prediction, Shape: (B, num_nuclei_classes, H, W)
* hv_map: Binary HV Map predictions. Shape: (B, 2, H, W)
overlap (int): Used overlap as integer
Returns:
List: Cleaned (merged) cell list with each entry beeing one detected cell with dictionary as entries.
"""
predictions["nuclei_binary_map"] = F.softmax(
predictions["nuclei_binary_map"], dim=1
)
predictions["nuclei_type_map"] = F.softmax(
predictions["nuclei_type_map"], dim=1
)
(
predictions["instance_map"],
predictions["instance_types"],
) = self.model.calculate_instance_map(
predictions, magnification=self.magnification
)
predictions = self.merge_predictions(predictions, overlap)
return predictions
def merge_predictions(self, predictions: dict, overlap: int) -> list:
"""Merge overlapping cell predictions
Args:
predictions (dict): Predictions with necessary keys:
* nuclei_binary_map: Binary nuclei prediction, Shape: (B, 2, H, W)
* instance_types: Instance type dictionary with cell entries
overlap (int): Used overlap as integer
Returns:
list: Cleaned (merged) cell list with each entry beeing one detected cell with dictionary as entries.
"""
cell_list = []
decomposed_patch_num = int(np.sqrt(predictions["nuclei_binary_map"].shape[0]))
for i in range(decomposed_patch_num):
for j in range(decomposed_patch_num):
x_global = i * 256 - i * overlap
y_global = j * 256 - j * overlap
patch_instance_types = predictions["instance_types"][
i * decomposed_patch_num + j
]
for cell in patch_instance_types.values():
if cell["type"] == 0:
continue
offset_global = np.array([x_global, y_global])
centroid_global = cell["centroid"] + np.flip(offset_global)
contour_global = cell["contour"] + np.flip(offset_global)
bbox_global = cell["bbox"] + offset_global
cell_dict = {
"bbox": bbox_global.tolist(),
"centroid": centroid_global.tolist(),
"contour": contour_global.tolist(),
"type_prob": cell["type_prob"],
"type": cell["type"],
"patch_coordinates": [
i, # row
j, # col
],
"cell_status": get_cell_position_marging(cell["bbox"], 256, 64),
"offset_global": offset_global.tolist(),
}
if np.max(cell["bbox"]) == 256 or np.min(cell["bbox"]) == 0:
position = get_cell_position(cell["bbox"], 256)
cell_dict["edge_position"] = True
cell_dict["edge_information"] = {}
cell_dict["edge_information"]["position"] = position
cell_dict["edge_information"]["edge_patches"] = get_edge_patch(
position, i, j # row, col
)
else:
cell_dict["edge_position"] = False
cell_list.append(cell_dict)
self.logger.info(f"Detected cells before cleaning: {len(cell_list)}")
cell_processor = CellPostProcessor(cell_list, self.logger)
cleaned_cells = cell_processor.post_process_cells()
cell_list = [cell_list[idx_c] for idx_c in cleaned_cells.index.values]
self.logger.info(f"Detected cells after cleaning: {len(cell_list)}")
return cell_list
def calculate_step_metric_overlap(
self, cell_list: List[dict], gt: dict, image_name: List[str]
) -> Tuple[dict, dict]:
"""Calculate step metric and return merged predictions for plotting
Args:
cell_list (List[dict]): List with cell dicts
gt (dict): Ground-Truth dictionary
image_name (List[str]): Image Name as list with just one entry
Returns:
Tuple[dict, dict]:
dict: Image metrics for one MoNuSeg image. Keys are:
* image_name
* binary_dice_score
* binary_jaccard_score
* pq_score
* dq_score
* sq_score
* f1_d
* prec_d
* rec_d
dict: Predictions, reshaped for one image and for plotting
* nuclei_binary_map: Shape (1, 2, 1024, 1024) or (1, 2, 1024, 1024)
* instance_map: Shape (1, 1024, 1024) or or (1, 2, 512, 512)
* instance_types: Dict for each nuclei
"""
predictions = {}
h, w = gt["nuclei_binary_map"].shape[1:]
instance_type_map = np.zeros((h, w), dtype=np.int32)
for instance, cell in enumerate(cell_list):
contour = np.array(cell["contour"])[None, :, :]
cv2.fillPoly(instance_type_map, contour, instance)
predictions["instance_map"] = torch.Tensor(instance_type_map)
instance_maps_gt = gt["instance_map"].detach().cpu()
pred_arr = np.clip(instance_type_map, 0, 1)
target_binary_map = gt["nuclei_binary_map"].to(self.device).squeeze()
predictions["nuclei_binary_map"] = pred_arr
predictions["instance_types"] = cell_list
cell_dice = (
dice(
preds=torch.Tensor(pred_arr).to(self.device),
target=target_binary_map,
ignore_index=0,
)
.detach()
.cpu()
)
cell_jaccard = (
binary_jaccard_index(
preds=torch.Tensor(pred_arr).to(self.device),
target=target_binary_map,
)
.detach()
.cpu()
)
remapped_instance_pred = remap_label(predictions["instance_map"])[None, :, :]
remapped_gt = remap_label(instance_maps_gt)
[dq, sq, pq], _ = get_fast_pq(true=remapped_gt, pred=remapped_instance_pred)
# detection scores
true_centroids = np.array(
[v["centroid"] for k, v in gt["instance_types"][0].items()]
)
pred_centroids = np.array([v["centroid"] for v in cell_list])
if true_centroids.shape[0] == 0:
true_centroids = np.array([[0, 0]])
if pred_centroids.shape[0] == 0:
pred_centroids = np.array([[0, 0]])
if self.magnification == 40:
pairing_radius = 12
else:
pairing_radius = 6
paired, unpaired_true, unpaired_pred = pair_coordinates(
true_centroids, pred_centroids, pairing_radius
)
f1_d, prec_d, rec_d = cell_detection_scores(
paired_true=paired[:, 0],
paired_pred=paired[:, 1],
unpaired_true=unpaired_true,
unpaired_pred=unpaired_pred,
)
image_metrics = {
"image_name": image_name,
"binary_dice_score": cell_dice,
"binary_jaccard_score": cell_jaccard,
"pq_score": pq,
"dq_score": dq,
"sq_score": sq,
"f1_d": f1_d,
"prec_d": prec_d,
"rec_d": rec_d,
}
# align to common shapes
cleaned_instance_types = {
k + 1: v for k, v in enumerate(predictions["instance_types"])
}
for cell, results in cleaned_instance_types.items():
results["contour"] = np.array(results["contour"])
cleaned_instance_types[cell] = results
predictions["instance_types"] = cleaned_instance_types
predictions["instance_map"] = predictions["instance_map"][None, :, :]
predictions["nuclei_binary_map"] = F.one_hot(
torch.Tensor(predictions["nuclei_binary_map"]).type(torch.int64),
num_classes=2,
).permute(2, 0, 1)[None, :, :, :]
return image_metrics, predictions
def plot_results(
self,
img: torch.Tensor,
predictions: dict,
) -> None:
"""Plot MoNuSeg results
Args:
img (torch.Tensor): Image as torch.Tensor, with Shape (1, 3, 1024, 1024) or (1, 3, 512, 512)
predictions (dict): Prediction dictionary. Necessary keys:
* nuclei_binary_map: Shape (1, 2, 1024, 1024) or (1, 2, 512, 512)
* instance_map: Shape (1, 1024, 1024) or (1, 512, 512)
* instance_types: List[dict], but just one entry in list
ground_truth (dict): Ground-Truth dictionary. Necessary keys:
* nuclei_binary_map: (1, 1024, 1024) or or (1, 512, 512)
* instance_map: (1, 1024, 1024) or or (1, 512, 512)
* instance_types: List[dict], but just one entry in list
img_name (str): Image name as string
outdir (Path): Output directory for storing
scores (List[float]): Scores as list [Dice, Jaccard, bPQ]
"""
predictions["nuclei_binary_map"] = predictions["nuclei_binary_map"].permute(
0, 2, 3, 1
)
h = predictions["instance_map"].shape[1]
w = predictions["instance_map"].shape[2]
# process image and other maps
sample_image = img.permute(0, 2, 3, 1).contiguous().cpu().numpy()
pred_sample_binary_map = (
predictions["nuclei_binary_map"][:, :, :, 1].detach().cpu().numpy()
)[0]
pred_sample_instance_maps = (
predictions["instance_map"].detach().cpu().numpy()[0]
)
binary_cmap = plt.get_cmap("Greys_r")
instance_map = plt.get_cmap("viridis")
# invert the normalization of the sample images
transform_settings = self.run_conf["transformations"]
if "normalize" in transform_settings:
mean = transform_settings["normalize"].get("mean", (0.5, 0.5, 0.5))
std = transform_settings["normalize"].get("std", (0.5, 0.5, 0.5))
else:
mean = (0.5, 0.5, 0.5)
std = (0.5, 0.5, 0.5)
inv_normalize = transforms.Normalize(
mean=[-0.5 / mean[0], -0.5 / mean[1], -0.5 / mean[2]],
std=[1 / std[0], 1 / std[1], 1 / std[2]],
)
inv_samples = inv_normalize(torch.tensor(sample_image).permute(0, 3, 1, 2))
sample_image = inv_samples.permute(0, 2, 3, 1).detach().cpu().numpy()[0]
# start overlaying on image
placeholder = np.zeros(( h, 4 * w, 3))
# orig image
placeholder[:h, :w, :3] = sample_image
# binary prediction
placeholder[:h, w : 2 * w, :3] = rgba2rgb(
binary_cmap(pred_sample_binary_map * 255)
)
# instance_predictions
placeholder[:h, 2 * w : 3 * w, :3] = rgba2rgb(
instance_map(
(pred_sample_instance_maps - np.min(pred_sample_instance_maps))
/ (
np.max(pred_sample_instance_maps)
- np.min(pred_sample_instance_maps + 1e-10)
)
)
)
# pred
pred_contours_polygon = [
v["contour"] for v in predictions["instance_types"].values()
]
pred_contours_polygon = [
list(zip(poly[:, 0], poly[:, 1])) for poly in pred_contours_polygon
]
pred_contour_colors_polygon = [
"#70c6ff" for i in range(len(pred_contours_polygon))
]
pred_cell_image = Image.fromarray(
(sample_image * 255).astype(np.uint8)
).convert("RGB")
pred_drawing = ImageDraw.Draw(pred_cell_image)
add_patch = lambda poly, color: pred_drawing.polygon(
poly, outline=color, width=2
)
[
add_patch(poly, c)
for poly, c in zip(pred_contours_polygon, pred_contour_colors_polygon)
]
placeholder[: h, 3 * w : 4 * w, :3] = np.asarray(pred_cell_image) / 255
pred_cell_image.save("raw_pred.png")
# plotting
fig, axs = plt.subplots(figsize=(3, 2), dpi=1200)
axs.imshow(placeholder)
axs.set_xticks(np.arange(w / 2, 4 * w, w))
axs.set_xticklabels(
[
"Image",
"Binary-Cells",
"Instances",
"Countours",
],
fontsize=6,
)
axs.xaxis.tick_top()
axs.set_yticks([h / 2])
axs.set_yticklabels([ "Pred."], fontsize=6)
axs.tick_params(axis="both", which="both", length=0)
grid_x = np.arange(w, 3 * w, w)
grid_y = np.arange(h, 2 * h, h)
for x_seg in grid_x:
axs.axvline(x_seg, color="black")
for y_seg in grid_y:
axs.axhline(y_seg, color="black")
fig.suptitle(f"All Predictions for input image", fontsize=6)
fig.tight_layout()
fig.savefig("pred_img.png")
plt.close()
# CLI
class InferenceCellViTMoNuSegParser:
def __init__(self) -> None:
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
description="Perform CellViT inference for MoNuSeg dataset",
)
parser.add_argument(
"--model",
type=str,
help="Model checkpoint file that is used for inference",
default="./model_best.pth",
)
parser.add_argument(
"--dataset",
type=str,
help="Path to MoNuSeg dataset.",
default="/data/lunbinzeng/datasets/monuseg/testing/",
)
parser.add_argument(
"--outdir",
type=str,
help="Path to output directory to store results.",
default="/data/lunbinzeng/results/lkcell/small/2024-04-22T232903_CellViT-unireplknet-fold1-final/monuseg/inference/",
)
parser.add_argument(
"--gpu", type=int, help="Cuda-GPU ID for inference. Default: 0", default=0
)
parser.add_argument(
"--magnification",
type=int,
help="Dataset Magnification. Either 20 or 40. Default: 40",
choices=[20, 40],
default=20,
)
parser.add_argument(
"--patching",
type=bool,
help="Patch to 256px images. Default: False",
default=False,
)
parser.add_argument(
"--overlap",
type=int,
help="Patch overlap, just valid for patching",
default=0,
)
parser.add_argument(
"--plots",
type=bool,
help="Generate result plots. Default: False",
default=True,
)
self.parser = parser
def parse_arguments(self) -> dict:
opt = self.parser.parse_args()
return vars(opt)
if __name__ == "__main__":
configuration_parser = InferenceCellViTMoNuSegParser()
configuration = configuration_parser.parse_arguments()
print(configuration)
inf = MoNuSegInference(
model_path=configuration["model"],
dataset_path=configuration["dataset"],
outdir=configuration["outdir"],
gpu=configuration["gpu"],
patching=configuration["patching"],
magnification=configuration["magnification"],
overlap=configuration["overlap"],
)
inf.run_inference(generate_plots=configuration["plots"])