Update metrics.py

Author: Oppie2333

src/patchcore/metrics.py

@@ -1,9 +1,13 @@
+"""Anomaly metrics."""
 import numpy as np
-from sklearn.metrics import average_precision_score
+from sklearn import metrics
 
-def compute_image_ap(anomaly_prediction_weights, anomaly_ground_truth_labels):
+
+def compute_imagewise_retrieval_metrics(
+    anomaly_prediction_weights, anomaly_ground_truth_labels
+):
     """
-    Computes the average precision (AP) for image-wise anomaly detection.
+    Computes retrieval statistics (AUROC, FPR, TPR).
 
     Args:
         anomaly_prediction_weights: [np.array or list] [N] Assignment weights
@@ -12,11 +16,19 @@ def compute_image_ap(anomaly_prediction_weights, anomaly_ground_truth_labels):
         anomaly_ground_truth_labels: [np.array or list] [N] Binary labels - 1
                                     if image is an anomaly, 0 if not.
     """
-    return average_precision_score(anomaly_ground_truth_labels, anomaly_prediction_weights)
+    fpr, tpr, thresholds = metrics.roc_curve(
+        anomaly_ground_truth_labels, anomaly_prediction_weights
+    )
+    auroc = metrics.roc_auc_score(
+        anomaly_ground_truth_labels, anomaly_prediction_weights
+    )
+    return {"auroc": auroc, "fpr": fpr, "tpr": tpr, "threshold": thresholds}
+
 
-def compute_pixel_ap(anomaly_segmentations, ground_truth_masks):
+def compute_pixelwise_retrieval_metrics(anomaly_segmentations, ground_truth_masks):
     """
-    Computes the average precision (AP) for pixel-wise anomaly detection.
+    Computes pixel-wise statistics (AUROC, FPR, TPR) for anomaly segmentations
+    and ground truth segmentation masks.
 
     Args:
         anomaly_segmentations: [list of np.arrays or np.array] [NxHxW] Contains
@@ -32,17 +44,33 @@ def compute_pixel_ap(anomaly_segmentations, ground_truth_masks):
     flat_anomaly_segmentations = anomaly_segmentations.ravel()
     flat_ground_truth_masks = ground_truth_masks.ravel()
 
-    return average_precision_score(flat_ground_truth_masks.astype(int), flat_anomaly_segmentations)
+    fpr, tpr, thresholds = metrics.roc_curve(
+        flat_ground_truth_masks.astype(int), flat_anomaly_segmentations
+    )
+    auroc = metrics.roc_auc_score(
+        flat_ground_truth_masks.astype(int), flat_anomaly_segmentations
+    )
 
-def compute_pro(anomaly_segmentations, ground_truth_masks):
-    """
-    Computes the PRO score. This is a placeholder function and needs to be implemented.
+    precision, recall, thresholds = metrics.precision_recall_curve(
+        flat_ground_truth_masks.astype(int), flat_anomaly_segmentations
+    )
+    F1_scores = np.divide(
+        2 * precision * recall,
+        precision + recall,
+        out=np.zeros_like(precision),
+        where=(precision + recall) != 0,
+    )
 
-    Args:
-        anomaly_segmentations: [list of np.arrays or np.array] [NxHxW] Contains
-                                generated segmentation masks.
-        ground_truth_masks: [list of np.arrays or np.array] [NxHxW] Contains
-                            predefined ground truth segmentation masks
-    """
-    # Placeholder implementation, replace with actual PRO calculation
-    return np.random.rand()
+    optimal_threshold = thresholds[np.argmax(F1_scores)]
+    predictions = (flat_anomaly_segmentations >= optimal_threshold).astype(int)
+    fpr_optim = np.mean(predictions > flat_ground_truth_masks)
+    fnr_optim = np.mean(predictions < flat_ground_truth_masks)
+
+    return {
+        "auroc": auroc,
+        "fpr": fpr,
+        "tpr": tpr,
+        "optimal_threshold": optimal_threshold,
+        "optimal_fpr": fpr_optim,
+        "optimal_fnr": fnr_optim,
+    }