Update app.py

app.py

@@ -1,161 +1,169 @@
 import torch
 import torch.nn.functional as F
-from transformers import AutoModelForImageClassification
+from transformers import AutoModelForImageClassification, pipeline
 from torchvision import transforms
 from PIL import Image, ImageStat
 from fastapi import FastAPI, File, UploadFile
 from fastapi.responses import HTMLResponse
+import numpy as np
 import io
 import gc
+import librosa
+import soundfile as sf
 
 # ==========================================
-# 1. CONFIGURATION (With STRICT Guardrails)
+# 1. CONFIGURATION
 # ==========================================
 MODELS = {
     "lungs": {
+        "type": "image",
         "id": "nickmuchi/vit-finetuned-chest-xray-pneumonia",
         "desc": "Tuberculosis & Pneumonia (Chest X-Ray)",
         "safe": ["NORMAL", "normal", "No Pneumonia"],
-        # Rule: Saturation must be LOW (Grayscale)
-        "guardrails": {"max_sat": 35} 
+        "rules": {"max_saturation": 30, "reject_msg": "Invalid: Too colorful. Upload B&W X-Ray."}
     },
-    "blood": { 
-        "id": "mrm8488/vit-base-patch16-224-finetuned-malaria-detection",
-        "desc": "Malaria Screening (Microscopic Slide)",
-        "safe": ["Uninfected", "uninfected"],
-        # Rule: Must be bright (Backlit slide)
-        "guardrails": {"min_bright": 60}
+    "cough": { 
+        "type": "audio",
+        "id": "MIT/ast-finetuned-audioset-10-10-0.4593",
+        "desc": "COPD & Respiratory Screening (Cough Audio)",
+        "target_labels": ["Cough", "Throat clearing", "Respiratory sounds", "Wheeze"],
+        "rules": {"min_duration": 1.0, "reject_msg": "Invalid: Audio too short or silent."}
     },
     "eye": {
+        "type": "image",
         "id": "AventIQ-AI/resnet18-cataract-detection-system", 
-        "desc": "Cataract Detection (Smartphone Eye Photo)",
+        "desc": "Cataract Detection (Eye Photo)",
         "safe": ["Normal", "normal", "healthy"],
-        # Rule: Saturation must be HIGH (Color photo) - Blocks X-Rays
-        "guardrails": {"min_sat": 20}
+        "rules": {"min_saturation": 20, "min_white_ratio": 0.05, "reject_msg": "Invalid: No eye detected."}
     },
     "skin": {
+        "type": "image",
         "id": "Anwarkh1/Skin_Cancer-Image_Classification",
-        "desc": "Dermatology & Lesion Analysis",
+        "desc": "Dermatology Analysis (Lesion Photo)",
         "safe": ["Benign", "benign", "nv", "bkl"],
-        # Rule: Saturation must be HIGH (Color photo) - Blocks X-Rays
-        "guardrails": {"min_sat": 20}
+        "rules": {"min_saturation": 20, "max_white_ratio": 0.25, "reject_msg": "Invalid: Not a skin close-up."}
     }
 }
 
 # ==========================================
-# 2. AI ENGINE
+# 2. MEDICAL ENGINE
 # ==========================================
 class MedicalEngine:
     def __init__(self):
         self.device = "cpu"
         print("✅ System Initialized: Medical Engine Ready")
         
-        self.transform = transforms.Compose([
+        self.img_transform = transforms.Compose([
             transforms.Resize((224, 224)),
             transforms.ToTensor(),
             transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
         ])
 
+    # --- IMAGE VALIDATION ---
     def validate_image(self, image, task):
-        """
-        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'))
+        rules = MODELS[task].get("rules", {})
+        img_hsv = image.convert('HSV')
+        stat = ImageStat.Stat(img_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."
-
+        
+        # Forensics
+        img_np = np.array(img_hsv)
+        white_pixels = np.logical_and(img_np[:,:,1] < 40, img_np[:,:,2] > 180)
+        white_ratio = np.sum(white_pixels) / white_pixels.size
+
+        if "max_saturation" in rules and avg_sat > rules["max_saturation"]: return False, rules["reject_msg"]
+        if "min_saturation" in rules and avg_sat < rules["min_saturation"]: return False, rules["reject_msg"]
+        if "min_white_ratio" in rules and white_ratio < rules["min_white_ratio"]: return False, rules["reject_msg"]
+        if "max_white_ratio" in rules and white_ratio > rules["max_white_ratio"]: return False, rules["reject_msg"]
+        
         return True, ""
 
-    def predict(self, image_bytes, task):
-        # 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 a special "INVALID" state
-            return {
-                "task": task,
-                "desc": MODELS[task]["desc"],
-                "prediction": {"label": "Invalid Image", "score": 0.0},
-                "risk": "INVALID",
-                "error": msg
-            }
+    # --- AUDIO VALIDATION ---
+    def validate_audio(self, audio_array, sr):
+        duration = len(audio_array) / sr
+        if duration < 0.5: return False, "Audio too short (< 0.5s)."
+        if np.max(np.abs(audio_array)) < 0.01: return False, "Audio is silent/empty."
+        return True, ""
 
-        # 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:
-            return {"error": "Failed to load AI model. Try again."}
-
-        # 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)
-            
-            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]
-            
-            # 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"
-                top["label"] = "Inconclusive / Unknown"
-            elif is_safe:
-                risk = "LOW"
-            else:
-                risk = "HIGH" if top["score"] > 0.75 else "MODERATE"
-
-        except Exception as e:
-            return {"error": f"Prediction Error: {str(e)}"}
+    # --- PREDICTION LOGIC ---
+    def predict(self, file_bytes, task):
+        model_cfg = MODELS[task]
         
-        finally:
-            del model
-            gc.collect()
-
-        return {
-            "task": task,
-            "desc": MODELS[task]["desc"],
-            "prediction": top,
-            "risk": risk
-        }
+        # === AUDIO PIPELINE ===
+        if model_cfg["type"] == "audio":
+            print(f"⏳ Processing Audio for {task}...")
+            try:
+                # Load Audio
+                audio, sr = librosa.load(io.BytesIO(file_bytes), sr=16000)
+                
+                # Guardrail
+                is_valid, msg = self.validate_audio(audio, sr)
+                if not is_valid: return {"error": msg, "risk": "INVALID"}
+
+                # Inference
+                classifier = pipeline("audio-classification", model=model_cfg["id"])
+                # We save to temp file because pipeline prefers paths or structured input
+                sf.write("temp.wav", audio, sr)
+                outputs = classifier("temp.wav")
+                
+                # Logic: Check if top prediction is a "Cough"
+                top = outputs[0]
+                target_labels = model_cfg["target_labels"]
+                
+                # Check if ANY of the top 3 results match "Cough"
+                is_cough = any(target in res['label'] for res in outputs[:3] for target in target_labels)
+                
+                risk = "HIGH" if is_cough and top['score'] > 0.4 else "LOW"
+                label = f"Detected: {top['label']}" if is_cough else "Normal / Background Noise"
+                
+                return {
+                    "task": task,
+                    "desc": model_cfg["desc"],
+                    "prediction": {"label": label, "score": top['score']},
+                    "risk": risk
+                }
+            except Exception as e:
+                return {"error": f"Audio Error: {str(e)}"}
+
+        # === IMAGE PIPELINE ===
+        else:
+            print(f"⏳ Processing Image for {task}...")
+            try:
+                image = Image.open(io.BytesIO(file_bytes)).convert("RGB")
+                is_valid, msg = self.validate_image(image, task)
+                if not is_valid: 
+                    return {"task": task, "risk": "INVALID", "error": msg, "prediction": {"label": "Rejected", "score": 0.0}}
+
+                model = AutoModelForImageClassification.from_pretrained(model_cfg["id"])
+                model.to(self.device)
+                model.eval()
+
+                inputs = self.transform(image).unsqueeze(0).to(self.device)
+                with torch.no_grad():
+                    outputs = model(inputs)
+                    probs = F.softmax(outputs.logits, dim=-1)
+                
+                results = [{"label": model.config.id2label[i], "score": float(score)} for i, score in enumerate(probs[0])]
+                results.sort(key=lambda x: x['score'], reverse=True)
+                top = results[0]
+                
+                safe_words = model_cfg["safe"]
+                is_safe = any(s.lower() in top["label"].lower() for s in safe_words)
+                
+                if top["score"] < 0.5: risk = "UNCERTAIN"
+                elif is_safe: risk = "LOW"
+                else: risk = "HIGH" if top["score"] > 0.75 else "MODERATE"
+                
+                del model
+                gc.collect()
+
+                return {"task": task, "desc": model_cfg["desc"], "prediction": top, "risk": risk}
+
+            except Exception as e:
+                return {"error": f"Image Error: {str(e)}"}
 
 # ==========================================
-# 3. API & UI
+# 3. API & FRONTEND
 # ==========================================
 app = FastAPI()
 engine = MedicalEngine()
@@ -198,16 +206,16 @@ def home():
     <div class="container mx-auto mt-8 p-4 max-w-3xl flex-grow">
         
         <div class="grid grid-cols-4 gap-2 mb-6">
-            <button onclick="setTask('lungs')" id="btn-lungs" class="p-3 bg-white rounded-lg shadow hover:bg-blue-50 border-2 border-transparent transition text-sm font-bold text-gray-600">
+            <button onclick="setTask('lungs', 'image')" id="btn-lungs" class="p-3 bg-white rounded-lg shadow hover:bg-blue-50 border-2 border-transparent transition text-sm font-bold text-gray-600">
                 <i class="fas fa-lungs text-blue-500 block text-2xl mb-1"></i> Lungs
             </button>
-            <button onclick="setTask('blood')" id="btn-blood" class="p-3 bg-white rounded-lg shadow hover:bg-red-50 border-2 border-transparent transition text-sm font-bold text-gray-600">
-                <i class="fas fa-burn text-red-500 block text-2xl mb-1"></i> Malaria
+            <button onclick="setTask('cough', 'audio')" id="btn-cough" class="p-3 bg-white rounded-lg shadow hover:bg-teal-50 border-2 border-transparent transition text-sm font-bold text-gray-600">
+                <i class="fas fa-head-side-cough text-teal-500 block text-2xl mb-1"></i> COPD
             </button>
-            <button onclick="setTask('eye')" id="btn-eye" class="p-3 bg-white rounded-lg shadow hover:bg-indigo-50 border-2 border-transparent transition text-sm font-bold text-gray-600">
+            <button onclick="setTask('eye', 'image')" id="btn-eye" class="p-3 bg-white rounded-lg shadow hover:bg-indigo-50 border-2 border-transparent transition text-sm font-bold text-gray-600">
                 <i class="fas fa-eye text-indigo-500 block text-2xl mb-1"></i> Eye
             </button>
-            <button onclick="setTask('skin')" id="btn-skin" class="p-3 bg-white rounded-lg shadow hover:bg-orange-50 border-2 border-transparent transition text-sm font-bold text-gray-600">
+            <button onclick="setTask('skin', 'image')" id="btn-skin" class="p-3 bg-white rounded-lg shadow hover:bg-orange-50 border-2 border-transparent transition text-sm font-bold text-gray-600">
                 <i class="fas fa-hand-dots text-orange-500 block text-2xl mb-1"></i> Skin
             </button>
         </div>
@@ -222,12 +230,15 @@ def home():
                 </div>
                 
                 <div onclick="document.getElementById('file-input').click()" class="border-2 border-dashed border-gray-300 rounded-xl p-8 text-center cursor-pointer hover:bg-gray-50 transition">
-                    <input type="file" id="file-input" class="hidden" accept="image/*" onchange="showPreview(event)" onclick="this.value=null">
+                    <input type="file" id="file-input" class="hidden" onchange="showPreview(event)" onclick="this.value=null">
+                    
                     <div id="placeholder">
-                        <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>
+                        <i id="upload-icon" class="fas fa-cloud-upload-alt text-3xl text-gray-400 mb-2"></i>
+                        <p id="upload-text" class="text-gray-500 text-sm">Tap to upload</p>
                     </div>
-                    <img id="preview" class="hidden mx-auto max-h-48 rounded shadow object-contain">
+                    
+                    <img id="img-preview" class="hidden mx-auto max-h-48 rounded shadow object-contain">
+                    <audio id="audio-preview" controls class="hidden w-full mt-4"></audio>
                 </div>
             </div>
 
@@ -237,14 +248,14 @@ 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-sm text-gray-500 mt-2 font-semibold">Analyzing Data...</p>
             </div>
 
             <div id="result-box" class="hidden mt-6 border-t pt-6">
                 <div class="flex justify-between items-start">
                     <div>
-                        <p class="text-xs font-bold text-gray-400 uppercase">Prediction</p>
-                        <h1 id="res-label" class="text-3xl font-extrabold text-gray-800">--</h1>
+                        <p class="text-xs font-bold text-gray-400 uppercase">Analysis</p>
+                        <h1 id="res-label" class="text-2xl font-extrabold text-gray-800">--</h1>
                         <p class="text-sm text-gray-600 mt-1">Confidence: <span id="res-conf" class="font-mono font-bold">--</span></p>
                     </div>
                     <span id="res-badge" class="px-3 py-1 rounded text-sm font-bold uppercase">--</span>
@@ -265,34 +276,66 @@ def home():
 
     <script>
         let currTask = null;
+        let currType = 'image';
         let currFile = null;
 
-        function setTask(task) {
+        function setTask(task, type) {
             currTask = task;
+            currType = type;
+            
+            // 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');
-            document.getElementById('header-text').innerHTML = `Upload <span class="uppercase text-blue-600">${task}</span> Image`;
+            
+            // UI Text
+            let taskName = task === 'cough' ? 'COUGH Audio' : task.toUpperCase() + ' Image';
+            document.getElementById('header-text').innerHTML = `Upload <span class="uppercase text-blue-600">${taskName}</span>`;
+            
+            // Input Type Handling
+            let input = document.getElementById('file-input');
+            let icon = document.getElementById('upload-icon');
+            let text = document.getElementById('upload-text');
+            
+            if (type === 'audio') {
+                input.accept = "audio/*";
+                icon.className = "fas fa-microphone-alt text-3xl text-teal-500 mb-2";
+                text.innerText = "Tap to upload Cough Audio (.wav/.mp3)";
+            } else {
+                input.accept = "image/*";
+                icon.className = "fas fa-cloud-upload-alt text-3xl text-gray-400 mb-2";
+                text.innerText = "Tap to upload Image";
+            }
+
+            // Reset
             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 = "";
+            document.getElementById('img-preview').classList.add('hidden');
+            document.getElementById('audio-preview').classList.add('hidden');
             currFile = null;
         }
 
         function showPreview(event) {
             if (event.target.files && event.target.files[0]) {
                 currFile = event.target.files[0];
-                let reader = new FileReader();
-                reader.onload = function(e) {
-                    document.getElementById('preview').src = e.target.result;
-                    document.getElementById('preview').classList.remove('hidden');
-                    document.getElementById('placeholder').classList.add('hidden');
-                    document.getElementById('run-btn').classList.remove('hidden');
-                    document.getElementById('result-box').classList.add('hidden');
-                };
-                reader.readAsDataURL(currFile);
+                let url = URL.createObjectURL(currFile);
+                
+                if (currType === 'audio') {
+                    let aud = document.getElementById('audio-preview');
+                    aud.src = url;
+                    aud.classList.remove('hidden');
+                    document.getElementById('img-preview').classList.add('hidden');
+                } else {
+                    let img = document.getElementById('img-preview');
+                    img.src = url;
+                    img.classList.remove('hidden');
+                    document.getElementById('audio-preview').classList.add('hidden');
+                }
+                
+                document.getElementById('placeholder').classList.add('hidden');
+                document.getElementById('run-btn').classList.remove('hidden');
+                document.getElementById('result-box').classList.add('hidden');
             }
         }
 
@@ -315,47 +358,28 @@ def home():
                 document.getElementById('result-box').classList.remove('hidden');
 
                 if (data.risk === "INVALID") {
-                    document.getElementById('res-label').innerText = "Image Rejected";
+                    document.getElementById('res-label').innerText = "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;
+                    updateBadge("INVALID", "bg-gray-200", "text-gray-700");
+                    showAlert(data.error || "Invalid File", "bg-gray-100");
                     return;
                 }
 
-                if (data.error) { alert("Error: " + data.error); resetLoading(); return; }
+                if (data.error) { alert("Error: " + data.error); document.getElementById('run-btn').classList.remove('hidden'); return; }
 
                 document.getElementById('res-label').innerText = data.prediction.label;
                 document.getElementById('res-conf').innerText = (data.prediction.score * 100).toFixed(1) + "%";
 
-                let badge = document.getElementById('res-badge');
-                let alertBox = document.getElementById('alert-box');
-
                 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.";
+                    updateBadge("HIGH RISK", "bg-red-100", "text-red-700");
+                    showAlert("Respiratory Symptom Detected. Refer to PHC.", "bg-red-50", "text-red-800");
                 } 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 === "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. Retake Photo.";
+                    updateBadge("MODERATE", "bg-yellow-100", "text-yellow-700");
+                    showAlert("Moderate Risk. Monitor symptoms.", "bg-yellow-50", "text-yellow-800");
                 } else {
-                    badge.className = "px-3 py-1 rounded text-sm font-bold uppercase bg-green-100 text-green-700";
-                    alertBox.classList.add('hidden');
+                    updateBadge("LOW RISK", "bg-green-100", "text-green-700");
+                    document.getElementById('alert-box').classList.add('hidden');
                 }
-                badge.innerText = data.risk + " RISK";
 
                 setTimeout(() => {
                     document.getElementById('sync-msg').innerHTML = "<i class='fas fa-check-circle'></i> Synced!";
@@ -363,15 +387,23 @@ def home():
                 }, 2000);
 
             } catch (e) {
-                alert("Connection Failed. Refresh and try again.");
+                alert("Connection Failed.");
                 console.error(e);
-                resetLoading();
+                document.getElementById('loader').classList.add('hidden');
+                document.getElementById('run-btn').classList.remove('hidden');
             }
         }
 
-        function resetLoading() {
-            document.getElementById('loader').classList.add('hidden');
-            document.getElementById('run-btn').classList.remove('hidden');
+        function updateBadge(text, bg, color) {
+            let b = document.getElementById('res-badge');
+            b.innerText = text;
+            b.className = `px-3 py-1 rounded text-sm font-bold uppercase ${bg} ${color}`;
+        }
+        function showAlert(msg, bg, color) {
+            let a = document.getElementById('alert-box');
+            a.className = `mt-4 p-3 rounded border border-gray-200 text-sm ${bg} ${color || ''}`;
+            a.classList.remove('hidden');
+            document.getElementById('alert-text').innerText = msg;
         }
     </script>
 </body>