Update app.py

app.py

@@ -11,15 +11,16 @@ from pytorch_grad_cam.utils.image import show_cam_on_image
 import os
 import datetime
 
+# -----------------------
 # Setup
+# -----------------------
 device = torch.device("cpu")
 save_dir = "/home/user/app/saved_predictions"
-if not os.path.exists(save_dir):
-    os.makedirs(save_dir)
-    print("📁 Folder created:", save_dir)
 os.makedirs(save_dir, exist_ok=True)
 
+# -----------------------
 # Load model
+# -----------------------
 model = models.resnet50(weights=None)
 model.fc = torch.nn.Linear(model.fc.in_features, 2)
 model.load_state_dict(torch.load("resnet50_dr_classifier.pth", map_location=device))
@@ -38,8 +39,48 @@ transform = transforms.Compose([
                          [0.229, 0.224, 0.225])
 ])
 
+# -----------------------
+# Helper: basic fundus check
+# -----------------------
+def looks_like_fundus(image):
+    """
+    Heuristic to check if an image is likely a retinal fundus scan.
+    Fundus images typically have a brighter central region (retina)
+    and a darker outer border (background).
+    """
+    img = np.array(image.convert("L").resize((224, 224)))
+
+    # Central crop
+    center = img[40:184, 40:184]
+    center_mean = center.mean()
+
+    # Border = everything outside the central crop
+    border = img.copy()
+    border[40:184, 40:184] = 0
+    border_pixels = border[border > 0]
+
+    # If no real border pixels, don't block it
+    if border_pixels.size == 0:
+        return True
+
+    border_mean = border_pixels.mean()
+
+    # Fundus: center clearly brighter than border
+    return center_mean > border_mean + 8
+
+
+# -----------------------
 # Predict and save
+# -----------------------
 def predict_retinopathy(image):
+    # 1. Block non-retina images BEFORE model / Grad-CAM
+    if not looks_like_fundus(image):
+        raise gr.Error(
+            "The uploaded image does not appear to be a retinal fundus scan. "
+            "Please upload a valid ophthalmic retinal image for Diabetic Retinopathy assessment."
+        )
+
+    # 2. Normal pipeline for valid retinal images
     timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
     img = image.convert("RGB").resize((224, 224))
     img_tensor = transform(img).unsqueeze(0).to(device)
@@ -55,7 +96,10 @@ def predict_retinopathy(image):
     # Grad-CAM
     rgb_img_np = np.array(img).astype(np.float32) / 255.0
     rgb_img_np = np.ascontiguousarray(rgb_img_np)
-    grayscale_cam = cam(input_tensor=img_tensor, targets=[ClassifierOutputTarget(pred)])[0]
+    grayscale_cam = cam(
+        input_tensor=img_tensor,
+        targets=[ClassifierOutputTarget(pred)]
+    )[0]
     cam_image = show_cam_on_image(rgb_img_np, grayscale_cam, use_rgb=True)
     cam_pil = Image.fromarray(cam_image)
 
@@ -65,8 +109,11 @@ def predict_retinopathy(image):
 
     return cam_pil, f"{label} (Confidence: {confidence:.2f})"
 
+
+# -----------------------
 # Gradio app
-gr.Interface(
+# -----------------------
+demo = gr.Interface(
     fn=predict_retinopathy,
     inputs=gr.Image(type="pil", label="Upload Retinal Image"),
     outputs=[
@@ -82,4 +129,7 @@ gr.Interface(
         "⚕️ **OpthaDetect** is an AI-powered ophthalmic decision-support tool. "
         "It highlights retinal risk regions using Grad-CAM for better clinical interpretability."
     )
-).launch()
+)
+
+if __name__ == "__main__":
+    demo.launch()