ParallelLLC/Segmentation
0
1"""2Segmentation Metrics3 4This module provides comprehensive metrics for evaluating segmentation performance5in few-shot and zero-shot learning scenarios.6"""7 8import torch9import torch.nn as nn10import numpy as np11from typing import Dict, List, Tuple, Optional12from sklearn.metrics import precision_recall_curve, average_precision_score13import cv214 15 16class SegmentationMetrics:17 """Comprehensive segmentation metrics calculator."""18 19 def __init__(self, threshold: float = 0.5):20 self.threshold = threshold21 22 def compute_metrics(23 self, 24 pred_mask: torch.Tensor, 25 gt_mask: torch.Tensor26 ) -> Dict[str, float]:27 """28 Compute comprehensive segmentation metrics.29 30 Args:31 pred_mask: Predicted mask tensor [H, W] or [1, H, W]32 gt_mask: Ground truth mask tensor [H, W] or [1, H, W]33 34 Returns:35 Dictionary containing various metrics36 """37 # Ensure masks are 2D38 if pred_mask.dim() == 3:39 pred_mask = pred_mask.squeeze(0)40 if gt_mask.dim() == 3:41 gt_mask = gt_mask.squeeze(0)42 43 # Convert to binary masks44 pred_binary = (pred_mask > self.threshold).float()45 gt_binary = (gt_mask > self.threshold).float()46 47 # Compute basic metrics48 metrics = {}49 50 # IoU (Intersection over Union)51 metrics['iou'] = self.compute_iou(pred_binary, gt_binary)52 53 # Dice coefficient54 metrics['dice'] = self.compute_dice(pred_binary, gt_binary)55 56 # Precision and Recall57 metrics['precision'] = self.compute_precision(pred_binary, gt_binary)58 metrics['recall'] = self.compute_recall(pred_binary, gt_binary)59 60 # F1 Score61 metrics['f1'] = self.compute_f1_score(pred_binary, gt_binary)62 63 # Accuracy64 metrics['accuracy'] = self.compute_accuracy(pred_binary, gt_binary)65 66 # Boundary metrics67 metrics['boundary_iou'] = self.compute_boundary_iou(pred_binary, gt_binary)68 metrics['hausdorff_distance'] = self.compute_hausdorff_distance(pred_binary, gt_binary)69 70 # Area metrics71 metrics['area_ratio'] = self.compute_area_ratio(pred_binary, gt_binary)72 73 return metrics74 75 def compute_iou(self, pred: torch.Tensor, gt: torch.Tensor) -> float:76 """Compute Intersection over Union."""77 intersection = (pred & gt).sum()78 union = (pred | gt).sum()79 return (intersection / union).item() if union > 0 else 0.080 81 def compute_dice(self, pred: torch.Tensor, gt: torch.Tensor) -> float:82 """Compute Dice coefficient."""83 intersection = (pred & gt).sum()84 total = pred.sum() + gt.sum()85 return (2 * intersection / total).item() if total > 0 else 0.086 87 def compute_precision(self, pred: torch.Tensor, gt: torch.Tensor) -> float:88 """Compute precision."""89 intersection = (pred & gt).sum()90 return (intersection / pred.sum()).item() if pred.sum() > 0 else 0.091 92 def compute_recall(self, pred: torch.Tensor, gt: torch.Tensor) -> float:93 """Compute recall."""94 intersection = (pred & gt).sum()95 return (intersection / gt.sum()).item() if gt.sum() > 0 else 0.096 97 def compute_f1_score(self, pred: torch.Tensor, gt: torch.Tensor) -> float:98 """Compute F1 score."""99 precision = self.compute_precision(pred, gt)100 recall = self.compute_recall(pred, gt)101 return (2 * precision * recall / (precision + recall)).item() if (precision + recall) > 0 else 0.0102 103 def compute_accuracy(self, pred: torch.Tensor, gt: torch.Tensor) -> float:104 """Compute pixel accuracy."""105 correct = (pred == gt).sum()106 total = pred.numel()107 return (correct / total).item()108 109 def compute_boundary_iou(self, pred: torch.Tensor, gt: torch.Tensor) -> float:110 """Compute boundary IoU."""111 # Extract boundaries112 pred_boundary = self.extract_boundary(pred)113 gt_boundary = self.extract_boundary(gt)114 115 # Compute IoU on boundaries116 return self.compute_iou(pred_boundary, gt_boundary)117 118 def extract_boundary(self, mask: torch.Tensor) -> torch.Tensor:119 """Extract boundary from binary mask."""120 mask_np = mask.cpu().numpy().astype(np.uint8)121 122 # Use morphological operations to extract boundary123 kernel = np.ones((3, 3), np.uint8)124 dilated = cv2.dilate(mask_np, kernel, iterations=1)125 eroded = cv2.erode(mask_np, kernel, iterations=1)126 boundary = dilated - eroded127 128 return torch.from_numpy(boundary).float()129 130 def compute_hausdorff_distance(self, pred: torch.Tensor, gt: torch.Tensor) -> float:131 """Compute Hausdorff distance between boundaries."""132 pred_boundary = self.extract_boundary(pred)133 gt_boundary = self.extract_boundary(gt)134 135 # Convert to numpy for distance computation136 pred_np = pred_boundary.cpu().numpy()137 gt_np = gt_boundary.cpu().numpy()138 139 # Find boundary points140 pred_points = np.column_stack(np.where(pred_np > 0))141 gt_points = np.column_stack(np.where(gt_np > 0))142 143 if len(pred_points) == 0 or len(gt_points) == 0:144 return float('inf')145 146 # Compute Hausdorff distance147 hausdorff_dist = self._hausdorff_distance(pred_points, gt_points)148 return hausdorff_dist149 150 def _hausdorff_distance(self, set1: np.ndarray, set2: np.ndarray) -> float:151 """Compute Hausdorff distance between two point sets."""152 def directed_hausdorff(set_a, set_b):153 min_distances = []154 for point_a in set_a:155 distances = np.linalg.norm(set_b - point_a, axis=1)156 min_distances.append(np.min(distances))157 return np.max(min_distances)158 159 d1 = directed_hausdorff(set1, set2)160 d2 = directed_hausdorff(set2, set1)161 return max(d1, d2)162 163 def compute_area_ratio(self, pred: torch.Tensor, gt: torch.Tensor) -> float:164 """Compute ratio of predicted area to ground truth area."""165 pred_area = pred.sum()166 gt_area = gt.sum()167 return (pred_area / gt_area).item() if gt_area > 0 else 0.0168 169 def compute_class_metrics(170 self, 171 predictions: Dict[str, torch.Tensor], 172 ground_truth: Dict[str, torch.Tensor]173 ) -> Dict[str, Dict[str, float]]:174 """Compute metrics for multiple classes."""175 class_metrics = {}176 177 for class_name in ground_truth.keys():178 if class_name in predictions:179 metrics = self.compute_metrics(predictions[class_name], ground_truth[class_name])180 class_metrics[class_name] = metrics181 else:182 # No prediction for this class183 class_metrics[class_name] = {184 'iou': 0.0,185 'dice': 0.0,186 'precision': 0.0,187 'recall': 0.0,188 'f1': 0.0,189 'accuracy': 0.0,190 'boundary_iou': 0.0,191 'hausdorff_distance': float('inf'),192 'area_ratio': 0.0193 }194 195 return class_metrics196 197 def compute_average_metrics(198 self, 199 class_metrics: Dict[str, Dict[str, float]]200 ) -> Dict[str, float]:201 """Compute average metrics across all classes."""202 if not class_metrics:203 return {}204 205 # Collect all metric names206 metric_names = list(class_metrics[list(class_metrics.keys())[0]].keys())207 208 # Compute averages209 averages = {}210 for metric_name in metric_names:211 values = [class_metrics[cls][metric_name] for cls in class_metrics.keys()]212 213 # Handle infinite values in Hausdorff distance214 if metric_name == 'hausdorff_distance':215 finite_values = [v for v in values if v != float('inf')]216 if finite_values:217 averages[metric_name] = np.mean(finite_values)218 else:219 averages[metric_name] = float('inf')220 else:221 averages[metric_name] = np.mean(values)222 223 return averages224 225 226class FewShotMetrics:227 """Specialized metrics for few-shot learning evaluation."""228 229 def __init__(self):230 self.segmentation_metrics = SegmentationMetrics()231 232 def compute_episode_metrics(233 self, 234 episode_results: List[Dict]235 ) -> Dict[str, float]:236 """Compute metrics across multiple episodes."""237 all_metrics = []238 239 for episode in episode_results:240 if 'metrics' in episode:241 all_metrics.append(episode['metrics'])242 243 if not all_metrics:244 return {}245 246 # Compute episode-level statistics247 episode_stats = {}248 metric_names = all_metrics[0].keys()249 250 for metric_name in metric_names:251 values = [ep[metric_name] for ep in all_metrics if metric_name in ep]252 if values:253 episode_stats[f'mean_{metric_name}'] = np.mean(values)254 episode_stats[f'std_{metric_name}'] = np.std(values)255 episode_stats[f'min_{metric_name}'] = np.min(values)256 episode_stats[f'max_{metric_name}'] = np.max(values)257 258 return episode_stats259 260 def compute_shot_analysis(261 self, 262 results_by_shots: Dict[int, List[Dict]]263 ) -> Dict[str, Dict[str, float]]:264 """Analyze performance across different numbers of shots."""265 shot_analysis = {}266 267 for num_shots, results in results_by_shots.items():268 episode_metrics = self.compute_episode_metrics(results)269 shot_analysis[f'{num_shots}_shots'] = episode_metrics270 271 return shot_analysis272 273 274class ZeroShotMetrics:275 """Specialized metrics for zero-shot learning evaluation."""276 277 def __init__(self):278 self.segmentation_metrics = SegmentationMetrics()279 280 def compute_prompt_strategy_comparison(281 self, 282 strategy_results: Dict[str, List[Dict]]283 ) -> Dict[str, Dict[str, float]]:284 """Compare different prompt strategies."""285 strategy_comparison = {}286 287 for strategy_name, results in strategy_results.items():288 # Compute average metrics for this strategy289 avg_metrics = {}290 if results:291 metric_names = results[0].keys()292 for metric_name in metric_names:293 values = [r[metric_name] for r in results if metric_name in r]294 if values:295 avg_metrics[f'mean_{metric_name}'] = np.mean(values)296 avg_metrics[f'std_{metric_name}'] = np.std(values)297 298 strategy_comparison[strategy_name] = avg_metrics299 300 return strategy_comparison301 302 def compute_attention_analysis(303 self, 304 with_attention: List[Dict], 305 without_attention: List[Dict]306 ) -> Dict[str, float]:307 """Analyze the impact of attention mechanisms."""308 if not with_attention or not without_attention:309 return {}310 311 # Compute average metrics312 with_attention_avg = {}313 without_attention_avg = {}314 315 metric_names = with_attention[0].keys()316 for metric_name in metric_names:317 with_values = [r[metric_name] for r in with_attention if metric_name in r]318 without_values = [r[metric_name] for r in without_attention if metric_name in r]319 320 if with_values:321 with_attention_avg[metric_name] = np.mean(with_values)322 if without_values:323 without_attention_avg[metric_name] = np.mean(without_values)324 325 # Compute improvements326 improvements = {}327 for metric_name in with_attention_avg.keys():328 if metric_name in without_attention_avg:329 improvement = with_attention_avg[metric_name] - without_attention_avg[metric_name]330 improvements[f'{metric_name}_improvement'] = improvement331 332 return {333 'with_attention': with_attention_avg,334 'without_attention': without_attention_avg,335 'improvements': improvements336 }