Update app.py

app.py

@@ -8,8 +8,6 @@ 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',
@@ -42,63 +40,29 @@ 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)
+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()
     
-#     # 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)
+    # Transform the target image and add a batch dimension
+    img = resnet50_transforms(img).unsqueeze(0)
     
-#     # 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))}
+    # 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)
     
-#     # Calculate the prediction time
-#     pred_time = round(timer() - start_time, 5)
+    # 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))}
     
-#     # Return the prediction dictionary and prediction time 
-#     return pred_labels_and_probs, pred_time
-
-def predict(img):
-    """Transforms and performs a prediction on img and returns prediction and time taken."""
-    start_time = timer()
+    # Calculate the prediction time
+    pred_time = round(timer() - start_time, 5)
     
-    try:
-        img = resnet50_transforms(img).unsqueeze(0)
-        resnet50.eval()
-        
-        with torch.inference_mode():
-            pred_probs = torch.softmax(resnet50(img), dim=1)
-            
-            # 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 base prediction dictionary
-            pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
-            
-            # OOD Detection - modify existing probabilities instead of adding new keys
-            if (max_prob > 0.95 and entropy < 0.05) or entropy > 0.1:
-                # Boost the probability of the first class and add a marker
-                pred_labels_and_probs[class_names[0]] = 0.99  # Use existing class
-                # You could also just print a warning or log it
-                print("May not be retina scan")
-        
-        pred_time = round(timer() - start_time, 5)
-        return pred_labels_and_probs, pred_time
-        
-    except Exception as e:
-        # Return dictionary with same structure as normal case
-        pred_labels_and_probs = {class_names[i]: 0.0 for i in range(len(class_names))}
-        pred_labels_and_probs[class_names[0]] = 1.0  # Show error in first class
-        return pred_labels_and_probs, 0.0
+    # Return the prediction dictionary and prediction time 
+    return pred_labels_and_probs, pred_time
 
 ### 4. Gradio app ###