|
|
import os
|
|
|
import torch
|
|
|
import numpy as np
|
|
|
|
|
|
|
|
|
class EarlyStopping():
|
|
|
def __init__(self, patience=5, min_delta=0.0):
|
|
|
self.patience = patience
|
|
|
self.min_delta = min_delta
|
|
|
self.best_metric = None
|
|
|
self.best_index = 0
|
|
|
|
|
|
|
|
|
def check_patience(self, metric_history):
|
|
|
if self.patience == 0:
|
|
|
return False
|
|
|
best_metric = min(metric_history)
|
|
|
best_index = metric_history.index(best_metric)
|
|
|
|
|
|
if len(metric_history) - best_index >= self.patience:
|
|
|
return True
|
|
|
|
|
|
return False
|
|
|
|
|
|
def check_improvement(self, metric_history):
|
|
|
if self.min_delta == 0.0:
|
|
|
return False
|
|
|
|
|
|
if len(metric_history) < 2:
|
|
|
return False
|
|
|
|
|
|
if metric_history[-2] - metric_history[-1] >= self.min_delta:
|
|
|
return True
|
|
|
|
|
|
return False
|
|
|
|
|
|
def __call__(self, metric_history):
|
|
|
if self.min_delta != 0.0:
|
|
|
return self.check_patience(metric_history)
|
|
|
|
|
|
if self.patience != 0:
|
|
|
return self.check_improvement(metric_history)
|
|
|
|
|
|
return False
|
|
|
|
|
|
class SaveBestModel():
|
|
|
def __init__(self, folder="./", mode='min'):
|
|
|
self.best_metric = None
|
|
|
self.folder = folder
|
|
|
self.mode = mode
|
|
|
|
|
|
def __call__(self, model, current_metric, model_name="best.pth"):
|
|
|
if self.best_metric is None:
|
|
|
self.best_metric = current_metric
|
|
|
save_path = os.path.join(self.folder, model_name)
|
|
|
torch.save(model.state_dict(), save_path)
|
|
|
else:
|
|
|
if (self.mode == 'min' and current_metric < self.best_metric) or \
|
|
|
(self.mode == 'max' and current_metric > self.best_metric):
|
|
|
self.best_metric = current_metric
|
|
|
save_path = os.path.join(self.folder, model_name)
|
|
|
torch.save(model.state_dict(), save_path)
|
|
|
|
|
|
|
|
|
class ModelLoss():
|
|
|
def __init__(self, task='segmentation', loss='focal', focal_alpha=0.25, focal_gamma=2.0):
|
|
|
if task not in ['segmentation', 'mae']:
|
|
|
raise ValueError(f"Unsupported task: {task}")
|
|
|
|
|
|
self.task = task
|
|
|
self.loss = loss
|
|
|
self.focal_alpha = focal_alpha
|
|
|
self.focal_gamma = focal_gamma
|
|
|
if self.task == 'segmentation' and self.loss not in ['focal', 'cross_entropy']:
|
|
|
raise ValueError(f"Unsupported loss for segmentation task: {self.loss}")
|
|
|
|
|
|
@staticmethod
|
|
|
def l1(input, target):
|
|
|
return torch.nn.functional.l1_loss(input, target)
|
|
|
|
|
|
@staticmethod
|
|
|
def l2(input, target):
|
|
|
return torch.nn.functional.mse_loss(input, target)
|
|
|
|
|
|
@staticmethod
|
|
|
def cross_entropy(input, target):
|
|
|
return torch.nn.functional.cross_entropy(input, target)
|
|
|
|
|
|
@staticmethod
|
|
|
def dice_loss(input: torch.Tensor, target: torch.Tensor, eps=1e-6):
|
|
|
input = torch.sigmoid(input)
|
|
|
input = input.view(-1)
|
|
|
target = target.view(-1)
|
|
|
|
|
|
intersection = (input * target).sum()
|
|
|
dice = (2. * intersection + eps) / (input.sum() + target.sum() + eps)
|
|
|
return 1 - dice
|
|
|
|
|
|
@staticmethod
|
|
|
def focal_loss(input: torch.Tensor, target: torch.Tensor, alpha=0.25, gamma=2.0, eps=1e-6):
|
|
|
input = torch.sigmoid(input)
|
|
|
input = input.view(-1)
|
|
|
target = target.view(-1)
|
|
|
|
|
|
bce_loss = torch.nn.functional.binary_cross_entropy(input, target, reduction='none')
|
|
|
p_t = input * target + (1 - input) * (1 - target)
|
|
|
focal_loss = alpha * (1 - p_t) ** gamma * bce_loss
|
|
|
return focal_loss.mean()
|
|
|
|
|
|
def __call__(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
|
if self.task == 'segmentation':
|
|
|
loss = 0.0
|
|
|
if self.loss == 'focal' or self.loss == 'default':
|
|
|
loss = ModelLoss.focal_loss(input, target, alpha=self.focal_alpha, gamma=self.focal_gamma)
|
|
|
elif self.loss == 'cross_entropy':
|
|
|
loss = ModelLoss.cross_entropy(input, target)
|
|
|
dice = ModelLoss.dice_loss(input, target)
|
|
|
return loss + dice
|
|
|
elif self.task == 'mae':
|
|
|
return ModelLoss.l2(input, target)
|
|
|
return None
|
|
|
|
|
|
|
|
|
|
|
|
class ModelMetrics():
|
|
|
def __init__(self, task='segmentation', device='cpu',threshold=0.5):
|
|
|
if task not in ['segmentation', 'mae']:
|
|
|
raise ValueError(f"Unsupported task: {task}")
|
|
|
|
|
|
self.task = task
|
|
|
self.device = device
|
|
|
self.threshold = threshold
|
|
|
|
|
|
@staticmethod
|
|
|
def iou_score(pred, target, eps=1e-6):
|
|
|
intersection = (pred * target).sum(dim=(1,2,3))
|
|
|
union = pred.sum(dim=(1,2,3)) + target.sum(dim=(1,2,3)) - intersection
|
|
|
return ((intersection + eps) / (union + eps)).mean()
|
|
|
|
|
|
@staticmethod
|
|
|
def dice_score(pred, target, eps=1e-6):
|
|
|
intersection = (pred * target).sum(dim=(1,2,3))
|
|
|
return ((2 * intersection + eps) /
|
|
|
(pred.sum(dim=(1,2,3)) + target.sum(dim=(1,2,3)) + eps)).mean()
|
|
|
|
|
|
@staticmethod
|
|
|
def pixel_accuracy(pred, target):
|
|
|
return (pred == target).float().mean()
|
|
|
|
|
|
@staticmethod
|
|
|
def pixel_precision(pred, target, eps=1e-6):
|
|
|
|
|
|
true_positive = (pred * target).sum()
|
|
|
predicted_positive = pred.sum()
|
|
|
return (true_positive.float() + eps) / (predicted_positive.float() + eps)
|
|
|
|
|
|
@staticmethod
|
|
|
def recall(pred, target, eps=1e-6):
|
|
|
true_positive = (pred * target).sum()
|
|
|
actual_positive = target.sum()
|
|
|
return (true_positive.float() + eps) / (actual_positive.float() + eps)
|
|
|
|
|
|
@staticmethod
|
|
|
def l1(input, target):
|
|
|
return torch.nn.functional.l1_loss(input, target)
|
|
|
|
|
|
@staticmethod
|
|
|
def l2(input, target):
|
|
|
return torch.nn.functional.mse_loss(input, target)
|
|
|
|
|
|
def getLabels(self):
|
|
|
if self.task == 'segmentation':
|
|
|
return ['iou', 'dice', 'accuracy', 'precision', 'recall']
|
|
|
elif self.task == 'mae':
|
|
|
return ['l1', 'l2']
|
|
|
return []
|
|
|
|
|
|
def __call__(self, pred: torch.Tensor, target: torch.Tensor):
|
|
|
"""
|
|
|
pred: (B, 1, H, W) => output of model BEFORE sigmoid
|
|
|
target: (B, H, W) or (B, 1, H, W)
|
|
|
"""
|
|
|
|
|
|
if self.task == 'segmentation':
|
|
|
|
|
|
pred = torch.sigmoid(pred)
|
|
|
pred = (pred > 0.5).float()
|
|
|
iou = self.iou_score(pred, target)
|
|
|
dice = self.dice_score(pred, target)
|
|
|
acc = self.pixel_accuracy(pred, target)
|
|
|
pres = self.pixel_precision(pred, target)
|
|
|
recall = self.recall(pred, target)
|
|
|
return iou.item(), dice.item(), acc.item(), pres.item(), recall.item()
|
|
|
|
|
|
if self.task == 'mae':
|
|
|
l1 = self.l1(pred, target)
|
|
|
l2 = self.l2(pred, target)
|
|
|
return l1.item(), l2.item()
|
|
|
|
|
|
def validate_mae(model, val_loader, metrics):
|
|
|
model.eval()
|
|
|
validation_accumulator = []
|
|
|
with torch.no_grad():
|
|
|
for sources, targets in val_loader:
|
|
|
sources = sources.cuda()
|
|
|
targets = targets.cuda()
|
|
|
|
|
|
pred = model(sources)
|
|
|
|
|
|
validation_accumulator.append(metrics(pred, targets))
|
|
|
|
|
|
return np.mean(validation_accumulator, axis=0) |
