update app.py

app.py

@@ -29,6 +29,11 @@ NUM_CLASSES = 2
 CLASS_NAMES = ["Elliptical", "Spiral"]
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# 🔴 Calibration + OOD thresholds
+TEMPERATURE = 2.5          # softens overconfidence
+CONF_THRESHOLD = 0.60      # below this → uncertain
+ENTROPY_THRESHOLD = 0.85   # high entropy → uncertain
+
 # Image preprocessing
 preprocess = transforms.Compose([
     transforms.Resize((224, 224)),
@@ -47,7 +52,7 @@ def load_model():
     if os.path.exists(MODEL_PATH):
         try:
             state_dict = torch.load(MODEL_PATH, map_location=DEVICE)
-            model.load_state_dict(state_dict)
+            model.load_state_dict((torch.load(MODEL_PATH, map_location=DEVICE))
             print(f"Model loaded successfully from {MODEL_PATH}")
         except Exception as e:
             print(f"Error loading model: {e}")
@@ -140,10 +145,24 @@ def predict_galaxy(image):
     img_tensor = preprocess(image).unsqueeze(0).to(DEVICE)
     img_tensor.requires_grad = True
     
-    outputs = model(img_tensor)
-    probs = F.softmax(outputs, dim=1)
-    pred_class = outputs.argmax(dim=1).item()
-    confidence = probs[0][pred_class].item()
+    # 🔴 Temperature scaling
+    scaled_logits = logits / TEMPERATURE
+    probs = F.softmax(scaled_logits, dim=1)[0]
+
+    confidence, pred_class = torch.max(probs, dim=0)
+
+    # 🔴 Entropy-based uncertainty
+    entropy = -(probs * torch.log(probs + 1e-8)).sum().item()
+
+    if confidence.item() < CONF_THRESHOLD or entropy > ENTROPY_THRESHOLD:
+        result_text = (
+            "**Prediction:** Uncertain / Not a Galaxy\n"
+            "**Confidence:** Low"
+        )
+        overlay_img = image.resize((224, 224))
+        return overlay_img, result_text
+
+
     
     gradcam = GradCAM(model, model.layer4)
     cam = gradcam.generate_cam(img_tensor, pred_class)
@@ -153,7 +172,11 @@ def predict_galaxy(image):
     overlay_rgb = cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB)
     overlay_pil = Image.fromarray(overlay_rgb)
     
-    result_text = f"Predicted Class: {CLASS_NAMES[pred_class]}\nConfidence: {confidence:.2%}"
+    # 🔴 Separate lines (as requested)
+    result_text = (
+        f"**Prediction:** {CLASS_NAMES[pred_class.item()]}\n"
+        f"**Confidence:** {confidence.item() * 100:.2f}%"
+    )
     
     return overlay_pil, result_text