Update app.py

app.py

@@ -8,6 +8,8 @@ from model import create_resnet50_model
 from timeit import default_timer as timer
 from typing import Tuple, Dict
 
+import torch.nn.functional as F
+
 # Setup class names
 class_names = ['CRVO',
   'Choroidal Nevus',
@@ -40,28 +42,56 @@ resnet50.load_state_dict(
 ### 3. Predict function ###
 
 # Create predict function
+# def predict(img) -> Tuple[Dict, float]:
+#     """Transforms and performs a prediction on img and returns prediction and time taken.
+#     """
+#     # Start the timer
+#     start_time = timer()
+    
+#     # Transform the target image and add a batch dimension
+#     img = resnet50_transforms(img).unsqueeze(0)
+    
+#     # Put model into evaluation mode and turn on inference mode
+#     resnet50.eval()
+#     with torch.inference_mode():
+#         # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
+#         pred_probs = torch.softmax(resnet50(img), dim=1)
+    
+#     # Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
+#     pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
+    
+#     # Calculate the prediction time
+#     pred_time = round(timer() - start_time, 5)
+    
+#     # Return the prediction dictionary and prediction time 
+#     return pred_labels_and_probs, pred_time
+
 def predict(img) -> Tuple[Dict, float]:
-    """Transforms and performs a prediction on img and returns prediction and time taken.
-    """
+    """Transforms and performs a prediction on img and returns prediction and time taken."""
     # Start the timer
     start_time = timer()
-    
     # Transform the target image and add a batch dimension
     img = resnet50_transforms(img).unsqueeze(0)
-    
     # Put model into evaluation mode and turn on inference mode
     resnet50.eval()
     with torch.inference_mode():
         # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
         pred_probs = torch.softmax(resnet50(img), dim=1)
-    
-    # Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
-    pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
+        
+        # Calculate entropy for OOD detection
+        entropy = -torch.sum(pred_probs * torch.log(pred_probs + 1e-8)).item()
+        max_prob = torch.max(pred_probs).item()
+        
+        # Create a prediction label and prediction probability dictionary for each prediction class
+        pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
+        
+        # OOD Detection: Flag suspicious predictions
+        if (max_prob > 0.95 and entropy < 0.2) or entropy > 2.0:
+            pred_labels_and_probs["May not be a retina scan"] = 0.99
     
     # Calculate the prediction time
     pred_time = round(timer() - start_time, 5)
-    
-    # Return the prediction dictionary and prediction time 
+    # Return the prediction dictionary and prediction time
     return pred_labels_and_probs, pred_time
 
 ### 4. Gradio app ###