Update app.py

app.py

@@ -7,6 +7,7 @@ import torchvision
 import numpy as np
 import scipy.ndimage as ndimage
 import gradio as gr
+import nibabel as nib # Added for direct raw loading
 from monai.transforms import (
     Compose, LoadImaged, EnsureChannelFirstd, 
     Resized, ScaleIntensityd, ToTensord, Lambdad,
@@ -47,19 +48,14 @@ class GatekeeperService:
         else:
             print(f"⚠️ Model file '{model_path}' not found.")
 
-        # --- AI PREPROCESSING (MATCHING TRAINING EXACTLY) ---
-        # We revert to ScaleIntensityd (Min/Max) because that is what the model learned.
+        # --- AI PIPELINE (PREPROCESSED) ---
         self.transforms = Compose([
             LoadImaged(keys=["image"]),
             EnsureChannelFirstd(keys=["image"]),
             Orientationd(keys=["image"], axcodes="RAS"), 
-            # We enforce the same spacing used in training (likely 1.5mm iso)
             Spacingd(keys=["image"], pixdim=(1.5, 1.5, 1.5), mode=("bilinear")), 
             Lambdad(keys=["image"], func=lambda x: x[0:1]), 
-            
-            # FIX: Reverted to Simple Scaling to match Training Data distribution
             ScaleIntensityd(keys=["image"]), 
-            
             Resized(keys=["image"], spatial_size=(96, 96, 96)), 
             ToTensord(keys=["image"])
         ])
@@ -69,7 +65,6 @@ class GatekeeperService:
         swi_path = None
         if not file_paths: return None, None
 
-        print(f"Scanning {len(file_paths)} uploaded files...")
         for path in file_paths:
             filename = os.path.basename(path).lower()
             if "flair" in filename or "t2_flair" in filename:
@@ -80,32 +75,51 @@ class GatekeeperService:
 
     def physics_check_hemorrhage(self, swi_path):
         """
-        Independent Logic: Uses its own robust normalization logic 
-        separate from the AI pipeline.
+        FIXED: Uses RAW nibabel loading to avoid interpolation artifacts.
         """
         if not swi_path: return False, 0.0
         try:
-            loader = LoadImaged(keys=["swi"])
-            data = loader({"swi": swi_path})
-            vol = data["swi"].numpy()
+            # 1. Load RAW Data (No Monai interpolation)
+            img_obj = nib.load(swi_path)
+            vol = img_obj.get_fdata()
             
-            # ROBUST NORMALIZATION (For Physics Only)
-            # Use percentiles to handle outliers in raw hospital data
+            # 2. Robust Normalization
+            # Map 0-99th percentile to 0-1
             p99 = np.percentile(vol, 99)
             if p99 > 0: vol = vol / p99
-            vol = np.clip(vol, 0, 1)
             
-            # Physics Rule: Signal Void < 0.05
-            dark_mask = vol < 0.05
+            # 3. Physics Rule
+            # Since we made a "Black Hole", pixels are exactly 0.0 or close to it.
+            # We assume anything < 2% intensity is a void (Blood or Air).
+            dark_mask = vol < 0.02
+            
+            # 4. Cluster Analysis
             labeled_array, num_features = ndimage.label(dark_mask)
-            voxel_vol = 1.5**3 
+            
+            # Use raw voxel size from header for accuracy
+            zooms = img_obj.header.get_zooms()
+            voxel_vol = zooms[0] * zooms[1] * zooms[2]
+            
+            largest_bleed_size = 0
             
             for i in range(1, num_features+1):
-                size = np.sum(labeled_array == i) * voxel_vol
-                # Size Check: 2mL - 200mL
-                if 2000 < size < 200000: 
-                    return True, size
+                # Count voxels in this cluster
+                voxel_count = np.sum(labeled_array == i)
+                size_mm3 = voxel_count * voxel_vol
+                
+                # Logic: 
+                # - Background Air is huge (> 500,000 mm3)
+                # - Noise is small (< 500 mm3)
+                # - Bleed is in between (500 mm3 - 300,000 mm3)
+                if 500 < size_mm3 < 300000:
+                    if size_mm3 > largest_bleed_size:
+                        largest_bleed_size = size_mm3
+            
+            if largest_bleed_size > 0:
+                return True, largest_bleed_size
+                
             return False, 0.0
+            
         except Exception as e:
             print(f"Physics Check Failed: {e}")
             return False, 0.0
@@ -125,15 +139,17 @@ class GatekeeperService:
             response["reasoning"].append("CRITICAL: No FLAIR sequence found.")
             return response
 
-        # 1. Physics Check
+        # 1. Physics Check (Priority 1)
         if swi_path:
             is_bleed, size = self.physics_check_hemorrhage(swi_path)
             if is_bleed:
                 response["status"] = "COMPLETED"
                 response["triage_category"] = "HEMORRHAGE"
                 response["next_agent"] = "REPORT_AGENT"
-                response["reasoning"].append(f"Physics Rule: Signal void ({size:.0f}mm3) detected on SWI.")
+                response["reasoning"].append(f"Physics Rule: Signal void ({size:.0f}mm3) detected on SWI. Thrombolysis Contraindicated.")
                 return response
+            else:
+                response["reasoning"].append("Physics Rule: SWI scan clear (No active bleed detected).")
 
         # 2. AI Prediction
         try:
@@ -144,21 +160,20 @@ class GatekeeperService:
                 output = self.model(tensor_data)
                 probs = torch.softmax(output, dim=1)
                 
-                # CONFIDENCE THRESHOLDING LOGIC
-                # If "Ischemic" is the 2nd choice but very close to "Tumor", 
-                # we should prioritize Ischemic for safety (Stroke Alert).
-                ischemic_idx = self.classes.index("ischemic")
-                ischemic_prob = probs[0, ischemic_idx].item()
-                
+                # BIAS CORRECTION: 
+                # If model is unsure between Tumor/Ischemic, prioritize Ischemic for safety.
                 pred_idx = torch.argmax(probs, dim=1).item()
                 confidence = probs[0, pred_idx].item()
                 
-                # Override Logic: If Model is unsure (Tumor vs Ischemic) but Ischemic > 20%, alert Stroke.
-                # It is safer to send a Tumor to the Stroke Vision Agent (which will find 0 volume) 
-                # than to send a Stroke home.
-                if self.classes[pred_idx] == "tumor" and ischemic_prob > 0.20:
+                # Get probabilities
+                p_ischemic = probs[0, self.classes.index("ischemic")].item()
+                p_tumor = probs[0, self.classes.index("tumor")].item()
+                
+                # Override: If Tumor is predicted but Ischemic is > 20%, force Ischemic Alert.
+                # Reason: Better to send a Tumor to Vision Agent (Safety) than miss a Stroke.
+                if self.classes[pred_idx] == "tumor" and p_ischemic > 0.20:
                     diagnosis = "ischemic"
-                    response["reasoning"].append(f"AI Alert: Possible Stroke pattern ({ischemic_prob:.2%}). Prioritizing over Tumor signal.")
+                    response["reasoning"].append(f"AI Alert: Complex pattern. Prioritizing Stroke (Prob: {p_ischemic:.2%}) over Tumor to ensure safety.")
                 else:
                     diagnosis = self.classes[pred_idx]
 
@@ -184,7 +199,7 @@ class GatekeeperService:
         return response
 
 # ==========================================
-# 3. INITIALIZATION
+# 3. INITIALIZATION & UI
 # ==========================================
 MODEL_FILE = "gatekeeper_hybrid_trained.pth"
 agent = GatekeeperService(MODEL_FILE)
@@ -192,9 +207,6 @@ agent = GatekeeperService(MODEL_FILE)
 def predict_mri_protocol(file_list):
     return agent.process_batch(file_list)
 
-# ==========================================
-# 4. GRADIO UI
-# ==========================================
 with gr.Blocks() as demo:
     gr.Markdown("# 🧠 Neuro-Triage: Gatekeeper Agent")