Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py (Use this code for Hugging Face)
+# app.py
 import torch
 import gradio as gr
 from PIL import Image
@@ -9,11 +9,9 @@ MODEL_NAME = "microsoft/swin-tiny-patch4-window7-224"
 MODEL_PATH = "best_model_swin.pth"
 NUM_CLASSES = 3
 CLASS_NAMES = ['COVID19', 'NORMAL', 'PNEUMONIA']
-device = torch.device("cpu") # Use CPU for free-tier hosting
+device = torch.device("cpu")
 
-# --- ADDED ---
 # We will reject any prediction where the model's top guess is below 90% confidence.
-# You can adjust this value (e.g., to 0.95 or 0.85)
 CONFIDENCE_THRESHOLD = 0.90  
 
 processor = ViTImageProcessor.from_pretrained(MODEL_NAME)
@@ -41,7 +39,6 @@ def classify_image(input_image: Image.Image):
         
     probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
     
-    # --- START OF MODIFICATION ---
     
     # Get the top class and its confidence score
     top_confidence, top_idx = torch.max(probabilities, dim=1)
@@ -53,7 +50,6 @@ def classify_image(input_image: Image.Image):
         # Return a custom label for low-confidence predictions
         return {f"Invalid Image or Low Confidence ({top_class_name})": top_confidence_score}
 
-    # --- END OF MODIFICATION ---
 
     # If confidence is high enough, return the normal dictionary
     confidences = {CLASS_NAMES[i]: prob.item() for i, prob in enumerate(probabilities[0])}