Update app.py

app.py

@@ -4,15 +4,12 @@ import torch.nn as nn
 from torchvision import models, transforms
 from huggingface_hub import hf_hub_download
 from PIL import Image
-import os
 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)
+logging.basicConfig(level=logging.WARNING)
 logger = logging.getLogger(__name__)
 
 # Define the number of classes
@@ -21,8 +18,8 @@ num_classes = 3
 # Confidence threshold for main model predictions
 CONFIDENCE_THRESHOLD = 0.8  # 80%
 
-# Mahalanobis distance threshold for OOD detection
-MAHALANOBIS_THRESHOLD = 400.0  # Calibrate this using a validation set
+# Energy threshold for OOD detection (to be calibrated)
+ENERGY_THRESHOLD = -5.0  # Placeholder, will calibrate
 
 # Download model from Hugging Face
 def download_model():
@@ -54,13 +51,6 @@ 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 ensure the file is in the same directory as app.py.")
-    raise
-
 # Path to your model
 model_path = download_model()
 main_model = load_main_model(model_path)
@@ -73,27 +63,15 @@ transform = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
 ])
 
-# Compute Mahalanobis distance for OOD detection
-def compute_mahalanobis_distance(features, mean, cov):
-    # Convert PyTorch tensors to NumPy arrays for scipy
-    features_np = features
-    mean_np = mean.cpu().numpy()
-    cov_np = cov.cpu().numpy()
-    # Compute the inverse covariance matrix
-    cov_inv = np.linalg.inv(cov_np + np.eye(cov_np.shape[0]) * 1e-6)  # Add small epsilon for numerical stability
-    return mahalanobis(features_np, mean_np, 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
+# Compute energy score for OOD detection
+def compute_energy_score(logits, temperature=1.0):
+    return -temperature * torch.logsumexp(logits / temperature, dim=1).item()
+
+# OOD detection using energy score
+def is_in_distribution(logits):
+    energy = compute_energy_score(logits)
+    logger.info(f"Energy score: {energy:.4f}")  # Log for calibration
+    return energy < ENERGY_THRESHOLD  # Lower (more negative) energy means ID
 
 # Prediction function for an uploaded image
 def predict_from_image_url(image_url):
@@ -105,35 +83,23 @@ def predict_from_image_url(image_url):
 
         # Apply transformations
         image_tensor = transform(image).unsqueeze(0)  # Shape: [1, 3, 224, 224]
-        logger.info(f"Input image tensor shape: {image_tensor.shape}")
 
-        # Extract features from the penultimate layer
+        # Stage 1: OOD Detection using energy score
         with torch.no_grad():
-            # 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].cpu().numpy()  # Convert to numpy for scipy
-
-        # 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",
-                "predicted_class": None,
-                "problem_id": None,
-                "confidence": None
-            }
+            logits = main_model(image_tensor)  # Shape: [1, 3]
+            if not is_in_distribution(logits):
+                logger.warning(f"Image URL {image_url} detected as out-of-distribution.")
+                return {
+                    "status": "invalid",
+                    "predicted_class": None,
+                    "problem_id": None,
+                    "confidence": None
+                }
 
         # 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()
+            probabilities = torch.softmax(logits, dim=1)[0]  # Convert to probabilities
+            predicted_class = torch.argmax(logits, dim=1).item()
 
         # Define class information
         class_info = {