Update app.py

app.py

@@ -2,10 +2,10 @@ import torch
 import torch.nn.functional as F
 from transformers import AutoModelForImageClassification, pipeline
 from torchvision import transforms
-from PIL import Image
-import numpy as np
+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
@@ -13,7 +13,7 @@ import soundfile as sf
 from datetime import datetime
 
 # ==========================================
-# 1. CONFIGURATION & GUARDRAIL RULES
+# 1. CONFIGURATION
 # ==========================================
 MODELS = {
     "lungs": {
@@ -21,14 +21,14 @@ MODELS = {
         "id": "nickmuchi/vit-finetuned-chest-xray-pneumonia",
         "desc": "Chest X-Ray Analysis",
         "safe": ["NORMAL", "normal", "No Pneumonia"],
-        # Rule: Must be Grayscale (Low Saturation)
         "rules": {"max_sat": 30, "reject_msg": "Invalid: Too colorful. Please upload a B&W X-Ray."}
     },
     "cough": { 
         "type": "audio",
         "id": "MIT/ast-finetuned-audioset-10-10-0.4593",
         "desc": "Respiratory Audio Analysis",
-        "target_labels": ["Cough", "Throat clearing", "Respiratory sounds", "Wheeze"],
+        # The AI looks for these specific sound tags
+        "target_labels": ["Cough", "Throat clearing", "Respiratory sounds", "Wheeze", "Gasping"],
         "rules": {"min_duration": 0.5, "reject_msg": "Invalid: Audio too short or silent."}
     },
     "fracture": {
@@ -36,7 +36,6 @@ MODELS = {
         "id": "dima806/bone_fracture_detection",
         "desc": "Bone Trauma X-Ray",
         "safe": ["normal", "healed"],
-        # Rule: Must be Grayscale
         "rules": {"max_sat": 30, "reject_msg": "Invalid: Too colorful. Please upload a B&W X-Ray."}
     },
     "brain": {
@@ -44,7 +43,6 @@ MODELS = {
         "id": "Hemgg/brain-tumor-classification",
         "desc": "Brain MRI Scan Analysis",
         "safe": ["no_tumor"],
-        # Rule: Must be Grayscale
         "rules": {"max_sat": 30, "reject_msg": "Invalid: Too colorful. Please upload a B&W MRI Scan."}
     },
     "eye": {
@@ -52,29 +50,19 @@ MODELS = {
         "id": "AventIQ-AI/resnet18-cataract-detection-system", 
         "desc": "Ophthalmology Scan",
         "safe": ["Normal", "normal", "healthy"],
-        # Rule: Must be Color AND have White Sclera
-        "rules": {
-            "min_sat": 20, 
-            "min_white": 0.05, # Eyes have >5% white pixels
-            "reject_msg": "Invalid: No eye detected (Missing white sclera)."
-        }
+        "rules": {"min_sat": 20, "min_white": 0.05, "reject_msg": "Invalid: No eye detected (Missing white sclera)."}
     },
     "skin": {
         "type": "image",
         "id": "Anwarkh1/Skin_Cancer-Image_Classification",
         "desc": "Dermatology Lesion Scan",
         "safe": ["Benign", "benign", "nv", "bkl"],
-        # Rule: Must be Color BUT NOT have huge white patches (like an eye)
-        "rules": {
-            "min_sat": 20, 
-            "max_white": 0.15, # Skin shouldn't have >15% pure white pixels
-            "reject_msg": "Invalid: Image looks like an Eye or Document (Too much white area)."
-        }
+        "rules": {"min_sat": 20, "max_white": 0.15, "reject_msg": "Invalid: Image looks like an Eye or Document (Too much white area)."}
     }
 }
 
 # ==========================================
-# 2. MEDICAL ENGINE (With Forensic Logic)
+# 2. MEDICAL ENGINE
 # ==========================================
 class MedicalEngine:
     def __init__(self):
@@ -88,38 +76,24 @@ class MedicalEngine:
         ])
 
     def validate_image(self, image, task):
-        """
-        Forensic Guardrails using NumPy
-        """
         rules = MODELS[task].get("rules", {})
-        
-        # Convert to HSV (Hue, Saturation, Value)
         img_hsv = image.convert('HSV')
         img_np = np.array(img_hsv)
         
-        # Calculate Averages
+        # Calculate Stats
         s_channel = img_np[:, :, 1]
         v_channel = img_np[:, :, 2]
         avg_sat = np.mean(s_channel)
         
-        # Calculate White Ratio (Low Saturation + High Brightness)
-        # Thresholds: Saturation < 40/255 AND Value > 180/255
+        # White Pixel Ratio
         white_pixels = np.logical_and(s_channel < 40, v_channel > 180)
         white_ratio = np.sum(white_pixels) / white_pixels.size
 
-        print(f"🔍 Analysis [{task}]: Sat={int(avg_sat)}, WhiteRatio={white_ratio:.3f}")
-
-        # 1. Saturation Check (Color vs Gray)
-        if "max_sat" in rules and avg_sat > rules["max_sat"]:
-            return False, rules["reject_msg"]
-        if "min_sat" in rules and avg_sat < rules["min_sat"]:
-            return False, "Invalid: Image is Black & White. Color photo required."
-
-        # 2. White Ratio Check (Eye vs Skin)
-        if "min_white" in rules and white_ratio < rules["min_white"]:
-            return False, rules["reject_msg"]
-        if "max_white" in rules and white_ratio > rules["max_white"]:
-            return False, rules["reject_msg"]
+        # Guardrails
+        if "max_sat" in rules and avg_sat > rules["max_sat"]: return False, rules["reject_msg"]
+        if "min_sat" in rules and avg_sat < rules["min_sat"]: return False, "Invalid: Image is Black & White. Color photo required."
+        if "min_white" in rules and white_ratio < rules["min_white"]: return False, rules["reject_msg"]
+        if "max_white" in rules and white_ratio > rules["max_white"]: return False, rules["reject_msg"]
         
         return True, ""
 
@@ -135,20 +109,30 @@ class MedicalEngine:
         # === AUDIO PIPELINE ===
         if model_cfg["type"] == "audio":
             try:
+                # Load Audio (Librosa handles mp3/wav if ffmpeg is installed)
                 audio, sr = librosa.load(io.BytesIO(file_bytes), sr=16000)
                 is_valid, msg = self.validate_audio(audio, sr)
                 if not is_valid: return {"error": msg, "risk": "INVALID"}
 
+                # Run AST Model
                 classifier = pipeline("audio-classification", model=model_cfg["id"])
+                # Pipeline requires file path, so we save temp
                 sf.write("temp.wav", audio, sr)
                 outputs = classifier("temp.wav")
                 
+                # Logic: Is "Cough" or "Wheeze" in the top predictions?
                 top = outputs[0]
                 target_labels = model_cfg["target_labels"]
-                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"
+                # Check top 3 predictions for any respiratory issue
+                is_respiratory = any(t in res['label'] for res in outputs[:3] for t in target_labels)
+                
+                if is_respiratory:
+                    risk = "HIGH"
+                    label = f"Detected: {top['label']}"
+                else:
+                    risk = "LOW"
+                    label = "Normal / Background Noise"
                 
                 return {"task": task, "desc": model_cfg["desc"], "prediction": {"label": label, "score": top['score']}, "risk": risk}
             except Exception as e:
@@ -436,13 +420,14 @@ def home():
             let txt = document.getElementById('upload-text');
             
             if (type === 'audio') {
-                input.accept = "audio/*";
+                input.accept = ".wav, .mp3, audio/*";
                 icon.className = "fas fa-microphone-alt text-4xl text-teal-500 mb-2";
+                txt.innerHTML = "Tap to upload Audio<br><span class='text-xs'>(.wav, .mp3)</span>";
             } else {
                 input.accept = "image/*";
                 icon.className = "fas fa-cloud-upload-alt text-4xl text-gray-400 mb-2";
+                txt.innerText = t['txt_upload'];
             }
-            txt.innerText = t['txt_upload'];
 
             document.getElementById('inputs').classList.remove('opacity-50', 'pointer-events-none');
             document.getElementById('result-box').classList.add('hidden');