Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from torch import nn
+import torch.nn as nn
 from torchvision import models, transforms
 from huggingface_hub import hf_hub_download
 from PIL import Image
@@ -9,6 +9,7 @@ import logging
 import requests
 from io import BytesIO
 import numpy as np
+from scipy.spatial.distance import mahalanobis
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -20,8 +21,8 @@ num_classes = 3
 # Confidence threshold for main model predictions
 CONFIDENCE_THRESHOLD = 0.8  # 80%
 
-# Threshold for OOD detection (maximum softmax probability)
-OPENMAX_THRESHOLD = 0.9  # If max softmax prob < 0.9, consider it OOD
+# Mahalanobis distance threshold for OOD detection
+MAHALANOBIS_THRESHOLD = 100.0  # Calibrate this using a validation set
 
 # Download model from Hugging Face
 def download_model():
@@ -53,6 +54,13 @@ def load_main_model(model_path):
     model.eval()
     return model
 
+# Load class statistics for Mahalanobis distance
+try:
+    class_statistics = torch.load("class_statistics.pth", map_location=torch.device("cpu"))
+except FileNotFoundError:
+    logger.error("class_statistics.pth not found. Please run the statistics computation script first.")
+    raise
+
 # Path to your model
 model_path = download_model()
 main_model = load_main_model(model_path)
@@ -65,12 +73,23 @@ transform = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
 ])
 
-# Simplified OpenMax-like OOD detection
-def openmax_ood_detection(probabilities):
-    max_prob = torch.max(probabilities).item()
-    logger.info(f"Max softmax probability: {max_prob:.4f}")
-    # If the maximum probability is below the threshold, consider it OOD
-    return max_prob >= OPENMAX_THRESHOLD
+# Compute Mahalanobis distance for OOD detection
+def compute_mahalanobis_distance(features, mean, cov):
+    # Compute the inverse covariance matrix
+    cov_inv = np.linalg.inv(cov + np.eye(cov.shape[0]) * 1e-6)  # Add small epsilon for numerical stability
+    return mahalanobis(features, mean, cov_inv)
+
+# OOD detection using Mahalanobis distance
+def is_in_distribution(features):
+    distances = []
+    for label in class_statistics:
+        mean = class_statistics[label]["mean"]
+        cov = class_statistics[label]["cov"]
+        distance = compute_mahalanobis_distance(features, mean, cov)
+        distances.append(distance)
+    min_distance = min(distances)
+    logger.info(f"Minimum Mahalanobis distance: {min_distance:.4f}")
+    return min_distance < MAHALANOBIS_THRESHOLD
 
 # Prediction function for an uploaded image
 def predict_from_image_url(image_url):
@@ -84,16 +103,17 @@ def predict_from_image_url(image_url):
         image_tensor = transform(image).unsqueeze(0)  # Shape: [1, 3, 224, 224]
         logger.info(f"Input image tensor shape: {image_tensor.shape}")
 
-        # Perform prediction
+        # Extract features from the penultimate layer
         with torch.no_grad():
-            outputs = main_model(image_tensor)  # Shape: [1, 3]
-            logger.info(f"Model output shape: {outputs.shape}")
-            logger.info(f"Raw logits: {outputs[0].numpy()}")
-            probabilities = torch.softmax(outputs, dim=1)[0]  # Convert to probabilities
-            logger.info(f"Softmax probabilities: {probabilities.numpy()}")
-
-        # OpenMax-like OOD detection
-        if not openmax_ood_detection(probabilities):
+            # Temporarily replace the final layer to get features
+            original_fc = main_model.fc
+            main_model.fc = nn.Identity()
+            features = main_model(image_tensor)  # Shape: [1, 2048]
+            main_model.fc = original_fc  # Restore the final layer
+            features = features[0].numpy()  # Convert to numpy
+
+        # Stage 1: OOD Detection using Mahalanobis distance
+        if not is_in_distribution(features):
             logger.warning(f"Image URL {image_url} detected as out-of-distribution.")
             return {
                 "status": "invalid",
@@ -102,8 +122,14 @@ def predict_from_image_url(image_url):
                 "confidence": None
             }
 
-        # Proceed with prediction
-        predicted_class = torch.argmax(outputs, dim=1).item()
+        # Stage 2: Main Model Prediction
+        with torch.no_grad():
+            outputs = main_model(image_tensor)  # Shape: [1, 3]
+            logger.info(f"Model output shape: {outputs.shape}")
+            logger.info(f"Raw logits: {outputs[0].numpy()}")
+            probabilities = torch.softmax(outputs, dim=1)[0]  # Convert to probabilities
+            logger.info(f"Softmax probabilities: {probabilities.numpy()}")
+            predicted_class = torch.argmax(outputs, dim=1).item()
 
         # Define class information
         class_info = {