Update app.py

app.py

@@ -3,52 +3,53 @@ import torch.nn.functional as F
 from transformers import AutoModelForImageClassification
 from torchvision import transforms
 from PIL import Image, ImageStat
-from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi import FastAPI, File, UploadFile
 from fastapi.responses import HTMLResponse
 import io
-import os
 import gc
 
 # ==========================================
-# 1. CONFIGURATION (Rural India Context)
+# 1. CONFIGURATION (With STRICT Guardrails)
 # ==========================================
 MODELS = {
     "lungs": {
         "id": "nickmuchi/vit-finetuned-chest-xray-pneumonia",
         "desc": "Tuberculosis & Pneumonia (Chest X-Ray)",
         "safe": ["NORMAL", "normal", "No Pneumonia"],
-        "check_color": True # X-Rays must be black & white
+        # Rule: Saturation must be LOW (Grayscale)
+        "guardrails": {"max_sat": 35} 
     },
     "blood": { 
         "id": "mrm8488/vit-base-patch16-224-finetuned-malaria-detection",
         "desc": "Malaria Screening (Microscopic Slide)",
         "safe": ["Uninfected", "uninfected"],
-        "check_color": False
+        # Rule: Must be bright (Backlit slide)
+        "guardrails": {"min_bright": 60}
     },
     "eye": {
         "id": "AventIQ-AI/resnet18-cataract-detection-system", 
         "desc": "Cataract Detection (Smartphone Eye Photo)",
         "safe": ["Normal", "normal", "healthy"],
-        "check_color": False
+        # Rule: Saturation must be HIGH (Color photo) - Blocks X-Rays
+        "guardrails": {"min_sat": 20}
     },
     "skin": {
         "id": "Anwarkh1/Skin_Cancer-Image_Classification",
         "desc": "Dermatology & Lesion Analysis",
         "safe": ["Benign", "benign", "nv", "bkl"],
-        "check_color": False
+        # Rule: Saturation must be HIGH (Color photo) - Blocks X-Rays
+        "guardrails": {"min_sat": 20}
     }
 }
 
 # ==========================================
-# 2. AI ENGINE (The Brain)
+# 2. AI ENGINE
 # ==========================================
 class MedicalEngine:
     def __init__(self):
-        # Force CPU to avoid memory crashes on Free Tier
         self.device = "cpu"
-        print("✅ System Initialized: Medical Engine Ready (Lazy Loading)")
+        print("✅ System Initialized: Medical Engine Ready")
         
-        # Standard Image Transformation
         self.transform = transforms.Compose([
             transforms.Resize((224, 224)),
             transforms.ToTensor(),
@@ -56,59 +57,84 @@ class MedicalEngine:
         ])
 
     def validate_image(self, image, task):
-        """Guardrail: Rejects obvious bad images (e.g. Selfies for X-Ray)"""
-        if MODELS[task]["check_color"]:
-            # Check saturation for X-Rays
-            stat = ImageStat.Stat(image.convert('HSV'))
-            saturation = stat.mean[1] # 0 = Gray, 255 = Color
-            if saturation > 35: # Threshold
-                return False, "⚠️ Invalid Image: This looks like a color photo. Please upload a Black & White X-Ray."
+        """
+        Universal Guardrails:
+        - Prevents X-Rays in Skin/Eye tabs (Checks Min Saturation)
+        - Prevents Selfies in X-Ray tab (Checks Max Saturation)
+        """
+        rules = MODELS[task].get("guardrails", {})
+        
+        # Convert to HSV (Hue, Saturation, Value)
+        # Saturation (Index 1): 0 = Gray, 255 = Color
+        # Value (Index 2): 0 = Dark, 255 = Bright
+        stat = ImageStat.Stat(image.convert('HSV'))
+        avg_sat = stat.mean[1]
+        avg_bright = stat.mean[2]
+
+        # 1. Check Max Saturation (Block Colorful images)
+        if "max_sat" in rules and avg_sat > rules["max_sat"]:
+            return False, f"⚠️ Invalid Image: Too colorful ({int(avg_sat)}). This looks like a photo, not an X-Ray."
+
+        # 2. Check Min Saturation (Block Grayscale images)
+        if "min_sat" in rules and avg_sat < rules["min_sat"]:
+            return False, f"⚠️ Invalid Image: Too gray ({int(avg_sat)}). This looks like an X-Ray/Doc. Please upload a color photo."
+
+        # 3. Check Min Brightness (Block Dark images)
+        if "min_bright" in rules and avg_bright < rules["min_bright"]:
+            return False, "⚠️ Invalid Image: Too dark. Microscope slides must be backlit."
+
         return True, ""
 
     def predict(self, image_bytes, task):
-        # A. Validation
+        # A. Load Image
         try:
             image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
         except:
             return {"error": "File is not a valid image."}
 
+        # B. Run Validation
         is_valid, msg = self.validate_image(image, task)
         if not is_valid:
-            return {"error": msg, "risk": "INVALID"}
+            # Return a special "INVALID" state
+            return {
+                "task": task,
+                "desc": MODELS[task]["desc"],
+                "prediction": {"label": "Invalid Image", "score": 0.0},
+                "risk": "INVALID",
+                "error": msg
+            }
 
-        # B. Load Model (On Demand)
-        print(f"⏳ Downloading/Loading Model for: {task}...")
+        # C. Load Model
+        print(f"⏳ Loading Model for: {task}...")
         try:
             model_id = MODELS[task]["id"]
             model = AutoModelForImageClassification.from_pretrained(model_id)
             model.to(self.device)
             model.eval()
         except Exception as e:
-            print(f"❌ Model Load Error: {e}")
             return {"error": "Failed to load AI model. Try again."}
 
-        # C. Inference
+        # D. Inference
         try:
             inputs = self.transform(image).unsqueeze(0).to(self.device)
             with torch.no_grad():
                 outputs = model(inputs)
                 probs = F.softmax(outputs.logits, dim=-1)
             
-            # Extract Results
             results = []
             for i, score in enumerate(probs[0]):
                 label = model.config.id2label[i]
                 results.append({"label": label, "score": float(score)})
             results.sort(key=lambda x: x['score'], reverse=True)
-            
             top = results[0]
             
-            # D. Risk Logic
+            # E. Risk Logic
             safe_words = MODELS[task]["safe"]
             is_safe = any(s.lower() in top["label"].lower() for s in safe_words)
             
             if top["score"] < 0.5:
-                risk = "UNCERTAIN" # Low confidence
+                risk = "UNCERTAIN"
+                top["label"] = "Inconclusive / Unknown"
             elif is_safe:
                 risk = "LOW"
             else:
@@ -118,7 +144,6 @@ class MedicalEngine:
             return {"error": f"Prediction Error: {str(e)}"}
         
         finally:
-            # E. Cleanup RAM (Critical for Free Tier)
             del model
             gc.collect()
 
@@ -130,22 +155,17 @@ class MedicalEngine:
         }
 
 # ==========================================
-# 3. API SERVER (FastAPI)
+# 3. API & UI
 # ==========================================
 app = FastAPI()
 engine = MedicalEngine()
 
 @app.post("/predict/{task}")
 async def predict_route(task: str, file: UploadFile = File(...)):
-    if task not in MODELS:
-        return {"error": "Invalid Task"}
-    
+    if task not in MODELS: return {"error": "Invalid Task"}
     content = await file.read()
     return engine.predict(content, task)
 
-# ==========================================
-# 4. FRONTEND UI (HTML embedded)
-# ==========================================
 @app.get("/", response_class=HTMLResponse)
 def home():
     return """
@@ -207,7 +227,7 @@ def home():
                         <i class="fas fa-cloud-upload-alt text-3xl text-gray-400 mb-2"></i>
                         <p class="text-gray-500 text-sm">Tap to upload image</p>
                     </div>
-                    <img id="preview" class="hidden mx-auto max-h-48 rounded shadow">
+                    <img id="preview" class="hidden mx-auto max-h-48 rounded shadow object-contain">
                 </div>
             </div>
 
@@ -218,7 +238,6 @@ def home():
             <div id="loader" class="hidden text-center py-6">
                 <div class="inline-block animate-spin rounded-full h-8 w-8 border-4 border-blue-500 border-t-transparent"></div>
                 <p class="text-sm text-gray-500 mt-2 font-semibold">Downloading AI Model & Analyzing...</p>
-                <p class="text-xs text-gray-400">(This takes ~20s on first run)</p>
             </div>
 
             <div id="result-box" class="hidden mt-6 border-t pt-6">
@@ -244,28 +263,21 @@ def home():
         </div>
     </div>
 
-    <footer class="text-center p-4 text-xs text-gray-400">
-        ⚠️ Disclaimer: AI Research Tool. Verify with a Doctor.
-    </footer>
-
     <script>
         let currTask = null;
         let currFile = null;
 
         function setTask(task) {
             currTask = task;
-            
-            // Highlight Buttons
             document.querySelectorAll('button[id^="btn-"]').forEach(b => b.classList.remove('ring-2', 'ring-blue-400', 'border-blue-500'));
             document.getElementById('btn-'+task).classList.add('ring-2', 'ring-blue-400', 'border-blue-500');
-            
-            // Reset UI
             document.getElementById('header-text').innerHTML = `Upload <span class="uppercase text-blue-600">${task}</span> Image`;
             document.getElementById('inputs').classList.remove('opacity-50', 'pointer-events-none');
             document.getElementById('result-box').classList.add('hidden');
             document.getElementById('run-btn').classList.add('hidden');
             document.getElementById('placeholder').classList.remove('hidden');
             document.getElementById('preview').classList.add('hidden');
+            document.getElementById('preview').src = "";
             currFile = null;
         }
 
@@ -286,13 +298,8 @@ def home():
 
         async function analyze() {
             if (!currTask || !currFile) return;
-            
-            if (!document.getElementById('p-name').value) {
-                alert("Please enter Patient Name.");
-                return;
-            }
+            if (!document.getElementById('p-name').value) { alert("Please enter Patient Name."); return; }
 
-            // Show Loader
             document.getElementById('run-btn').classList.add('hidden');
             document.getElementById('loader').classList.remove('hidden');
             document.getElementById('result-box').classList.add('hidden');
@@ -301,21 +308,27 @@ def home():
             formData.append("file", currFile);
 
             try {
-                // Determine API URL (Handles Hugging Face URL structure)
-                let url = "/predict/" + currTask;
-                
-                let res = await fetch(url, { method: "POST", body: formData });
+                let res = await fetch("/predict/" + currTask, { method: "POST", body: formData });
                 let data = await res.json();
 
-                if (data.error) {
-                    alert("Error: " + data.error);
-                    resetLoading();
+                document.getElementById('loader').classList.add('hidden');
+                document.getElementById('result-box').classList.remove('hidden');
+
+                if (data.risk === "INVALID") {
+                    document.getElementById('res-label').innerText = "Image Rejected";
+                    document.getElementById('res-conf').innerText = "--";
+                    let badge = document.getElementById('res-badge');
+                    badge.className = "px-3 py-1 rounded text-sm font-bold uppercase bg-gray-200 text-gray-700";
+                    badge.innerText = "INVALID";
+                    
+                    let alertBox = document.getElementById('alert-box');
+                    alertBox.className = "mt-4 p-3 bg-gray-100 text-gray-800 rounded border border-gray-300 text-sm";
+                    alertBox.classList.remove('hidden');
+                    document.getElementById('alert-text').innerText = data.error;
                     return;
                 }
 
-                // Show Results
-                document.getElementById('loader').classList.add('hidden');
-                document.getElementById('result-box').classList.remove('hidden');
+                if (data.error) { alert("Error: " + data.error); resetLoading(); return; }
 
                 document.getElementById('res-label').innerText = data.prediction.label;
                 document.getElementById('res-conf').innerText = (data.prediction.score * 100).toFixed(1) + "%";
@@ -325,24 +338,25 @@ def home():
 
                 if (data.risk === "HIGH") {
                     badge.className = "px-3 py-1 rounded text-sm font-bold uppercase bg-red-100 text-red-700";
+                    alertBox.className = "mt-4 p-3 bg-red-50 text-red-800 rounded border border-red-200 text-sm";
                     alertBox.classList.remove('hidden');
                     document.getElementById('alert-text').innerText = "High Risk. Immediate Referral Recommended.";
                 } else if (data.risk === "MODERATE") {
                     badge.className = "px-3 py-1 rounded text-sm font-bold uppercase bg-yellow-100 text-yellow-700";
+                    alertBox.className = "mt-4 p-3 bg-yellow-50 text-yellow-800 rounded border border-yellow-200 text-sm";
                     alertBox.classList.remove('hidden');
                     document.getElementById('alert-text').innerText = "Moderate Risk. Consult Doctor.";
-                } else if (data.risk === "INVALID" || data.risk === "UNCERTAIN") {
+                } else if (data.risk === "UNCERTAIN") {
                     badge.className = "px-3 py-1 rounded text-sm font-bold uppercase bg-gray-200 text-gray-700";
+                    alertBox.className = "mt-4 p-3 bg-gray-100 text-gray-800 rounded border border-gray-200 text-sm";
                     alertBox.classList.remove('hidden');
-                    document.getElementById('alert-text').innerText = "Image Unclear / Invalid. Retake Photo.";
+                    document.getElementById('alert-text').innerText = "Image Unclear. Retake Photo.";
                 } else {
                     badge.className = "px-3 py-1 rounded text-sm font-bold uppercase bg-green-100 text-green-700";
                     alertBox.classList.add('hidden');
                 }
-                
                 badge.innerText = data.risk + " RISK";
 
-                // Sync Animation
                 setTimeout(() => {
                     document.getElementById('sync-msg').innerHTML = "<i class='fas fa-check-circle'></i> Synced!";
                     document.getElementById('sync-msg').className = "text-green-600 font-bold";