mirror of
https://github.com/vladmandic/sdnext.git
synced 2026-01-27 15:02:48 +03:00
123 lines
3.6 KiB
Python
123 lines
3.6 KiB
Python
from collections import defaultdict
|
|
from typing import List, Union
|
|
|
|
from scipy.optimize import linear_sum_assignment
|
|
import PIL.Image as Image
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
__all__ = ['compute_iou', 'MeanEvaluator', 'load_mask', 'compute_ioa']
|
|
|
|
|
|
def compute_iou(a: torch.Tensor, b: torch.Tensor) -> float:
|
|
if a.shape[0] != b.shape[0]:
|
|
a = F.interpolate(a.unsqueeze(0).unsqueeze(0).float(), size=b.shape, mode='bicubic').squeeze()
|
|
a[a < 1] = 0
|
|
a[a >= 1] = 1
|
|
|
|
intersection = (a * b).sum()
|
|
union = a.sum() + b.sum() - intersection
|
|
|
|
return (intersection / (union + 1e-8)).item()
|
|
|
|
|
|
def compute_ioa(a: torch.Tensor, b: torch.Tensor) -> float:
|
|
if a.shape[0] != b.shape[0]:
|
|
a = F.interpolate(a.unsqueeze(0).unsqueeze(0).float(), size=b.shape, mode='bicubic').squeeze()
|
|
a[a < 1] = 0
|
|
a[a >= 1] = 1
|
|
|
|
intersection = (a * b).sum()
|
|
area = a.sum()
|
|
|
|
return (intersection / (area + 1e-8)).item()
|
|
|
|
|
|
def load_mask(path: str) -> torch.Tensor:
|
|
mask = np.array(Image.open(path))
|
|
mask = torch.from_numpy(mask).float()[:, :, 3] # use alpha channel
|
|
mask = (mask > 0).float()
|
|
|
|
return mask
|
|
|
|
|
|
class UnsupervisedEvaluator:
|
|
def __init__(self, name: str = 'UnsupervisedEvaluator'):
|
|
self.name = name
|
|
self.ious = defaultdict(list)
|
|
self.num_samples = 0
|
|
|
|
def log_iou(self, preds: Union[torch.Tensor, List[torch.Tensor]], truth: torch.Tensor, gt_idx: int = 0, pred_idx: int = 0):
|
|
if not isinstance(preds, list):
|
|
preds = [preds]
|
|
|
|
iou = max(compute_iou(pred, truth) for pred in preds)
|
|
self.ious[gt_idx].append((pred_idx, iou))
|
|
|
|
@property
|
|
def mean_iou(self) -> float:
|
|
n = max(max(self.ious), max([y[0] for x in self.ious.values() for y in x])) + 1
|
|
iou_matrix = np.zeros((n, n))
|
|
count_matrix = np.zeros((n, n))
|
|
|
|
for gt_idx, ious in self.ious.items():
|
|
for pred_idx, iou in ious:
|
|
iou_matrix[gt_idx, pred_idx] += iou
|
|
count_matrix[gt_idx, pred_idx] += 1
|
|
|
|
row_ind, col_ind = linear_sum_assignment(iou_matrix, maximize=True)
|
|
return iou_matrix[row_ind, col_ind].sum() / count_matrix[row_ind, col_ind].sum()
|
|
|
|
def increment(self):
|
|
self.num_samples += 1
|
|
|
|
def __len__(self) -> int:
|
|
return self.num_samples
|
|
|
|
def __str__(self):
|
|
return f'{self.name}<{self.mean_iou:.4f} (mIoU) {len(self)} samples>'
|
|
|
|
|
|
class MeanEvaluator:
|
|
def __init__(self, name: str = 'MeanEvaluator'):
|
|
self.ious: List[float] = []
|
|
self.intensities: List[float] = []
|
|
self.name = name
|
|
|
|
def log_iou(self, preds: Union[torch.Tensor, List[torch.Tensor]], truth: torch.Tensor):
|
|
if not isinstance(preds, list):
|
|
preds = [preds]
|
|
|
|
self.ious.append(max(compute_iou(pred, truth) for pred in preds))
|
|
return self
|
|
|
|
def log_intensity(self, pred: torch.Tensor):
|
|
self.intensities.append(pred.mean().item())
|
|
return self
|
|
|
|
@property
|
|
def mean_iou(self) -> float:
|
|
return np.mean(self.ious)
|
|
|
|
@property
|
|
def mean_intensity(self) -> float:
|
|
return np.mean(self.intensities)
|
|
|
|
@property
|
|
def ci95_miou(self) -> float:
|
|
return 1.96 * np.std(self.ious) / np.sqrt(len(self.ious))
|
|
|
|
def __len__(self) -> int:
|
|
return max(len(self.ious), len(self.intensities))
|
|
|
|
def __str__(self):
|
|
return f'{self.name}<{self.mean_iou:.4f} (±{self.ci95_miou:.3f} mIoU) {self.mean_intensity:.4f} (mInt) {len(self)} samples>'
|
|
|
|
|
|
if __name__ == '__main__':
|
|
mask = load_mask('truth/output/452/sink.gt.png')
|
|
|
|
print(MeanEvaluator().log_iou(mask, mask))
|