Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 import torch
 import torch.nn as nn
-import torch.serialization  # For add_safe_globals
 from torchvision import models, transforms
 from huggingface_hub import hf_hub_download
 from PIL import Image
@@ -25,14 +24,6 @@ CONFIDENCE_THRESHOLD = 0.8  # 80%
 # Mahalanobis distance threshold for OOD detection
 MAHALANOBIS_THRESHOLD = 100.0  # Calibrate this using a validation set
 
-# Allowlist NumPy globals for torch.load
-torch.serialization.add_safe_globals([
-    np.ndarray,
-    np._core.multiarray._reconstruct,
-    np.dtype,
-    _codecs.encode
-])
-
 # Download model from Hugging Face
 def download_model():
     model_path = hf_hub_download(repo_id="jays009/Resnet3", filename="model.pth")
@@ -84,9 +75,13 @@ transform = transforms.Compose([
 
 # 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.eye(cov.shape[0]) * 1e-6)  # Add small epsilon for numerical stability
-    return mahalanobis(features, mean, cov_inv)
+    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):
@@ -119,7 +114,7 @@ def predict_from_image_url(image_url):
             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
+            features = features[0].cpu().numpy()  # Convert to numpy for scipy
 
         # Stage 1: OOD Detection using Mahalanobis distance
         if not is_in_distribution(features):