Update app.py

app.py

@@ -20,6 +20,9 @@ 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
+
 # Download model from Hugging Face
 def download_model():
     model_path = hf_hub_download(repo_id="jays009/Resnet3", filename="model.pth")
@@ -62,6 +65,13 @@ 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
+
 # Prediction function for an uploaded image
 def predict_from_image_url(image_url):
     try:
@@ -78,12 +88,22 @@ def predict_from_image_url(image_url):
         with torch.no_grad():
             outputs = main_model(image_tensor)  # Shape: [1, 3]
             logger.info(f"Model output shape: {outputs.shape}")
-            # Log raw logits
             logger.info(f"Raw logits: {outputs[0].numpy()}")
             probabilities = torch.softmax(outputs, dim=1)[0]  # Convert to probabilities
-            # Log softmax probabilities
             logger.info(f"Softmax probabilities: {probabilities.numpy()}")
-            predicted_class = torch.argmax(outputs, dim=1).item()
+
+        # OpenMax-like OOD detection
+        if not openmax_ood_detection(probabilities):
+            logger.warning(f"Image URL {image_url} detected as out-of-distribution.")
+            return {
+                "status": "invalid",
+                "predicted_class": None,
+                "problem_id": None,
+                "confidence": None
+            }
+
+        # Proceed with prediction
+        predicted_class = torch.argmax(outputs, dim=1).item()
 
         # Define class information
         class_info = {