Update app.py

app.py

@@ -1,7 +1,8 @@
 import gradio as gr
 import tensorflow as tf
 import numpy as np
-from PIL import Image
+from PIL import Image, ImageStat
+import cv2
 
 # 🧠 Load your TFLite model
 interpreter = tf.lite.Interpreter(model_path="stool_model.tflite")
@@ -13,7 +14,7 @@ output_details = interpreter.get_output_details()
 # 🏷️ Labels of your stool model
 labels = ["bloody", "hard stool", "normal", "parasite", "watery"]
 
-# 💬 Diagnosis summaries for each label
+# 💬 Diagnosis summaries
 diagnosis_advice = {
     "bloody": "Possible hemorrhagic gastroenteritis. ⚠️ Visit a vet immediately as it may indicate internal bleeding or infection.",
     "hard stool": "Your pet may be constipated. 💧 Encourage hydration and increase dietary fiber intake.",
@@ -24,78 +25,95 @@ diagnosis_advice = {
     "Uncertain / unclear stool image": "⚠️ The image is unclear or confidence is low. Try retaking the photo in better lighting."
 }
 
-# 🧩 Preprocess image for TFLite
+# 🧩 Image preprocessing for the TFLite model
 def preprocess_image(img: Image.Image):
     img = img.convert("RGB").resize((128, 128))
     arr = np.asarray(img).astype(np.float32) / 255.0
     arr = np.expand_dims(arr, axis=0)
     return arr
 
-# 🚫 Detection filter for non-stool images
+# 🚫 Strict non-stool detection
 def is_not_stool_image(image):
-    arr = np.asarray(image.convert("RGB")).astype(np.float32)
-    brightness = arr.mean()
-    contrast = arr.std()
-    avg_color = arr.mean(axis=(0, 1))
-
-    if brightness > 220 or brightness < 25:
-        return True  # too bright or too dark
+    """Multi-layer detection for non-stool photos"""
+    np_img = np.array(image.convert("RGB"))
+    stat = ImageStat.Stat(image)
+
+    # 🧠 1️⃣ Brightness & contrast checks
+    brightness = stat.mean[0]
+    contrast = stat.stddev[0]
+    if brightness > 210 or brightness < 30:
+        return True  # too bright or dark
     if contrast < 25:
         return True  # too flat / no texture
-    if avg_color[0] > 180 and avg_color[1] < 80 and avg_color[2] < 80:
-        return True  # too reddish
-    if avg_color[1] > 200 and avg_color[0] < 100:
-        return True  # too greenish
-    if avg_color[2] > 200:
-        return True  # too bluish
-    return False
-
-# 🧠 Main classification logic
+
+    # 🧠 2️⃣ Dominant color detection (stool tends to be brownish)
+    avg_color = np.mean(np_img.reshape(-1, 3), axis=0)
+    r, g, b = avg_color
+
+    # Detect unnatural colors (too blue, green, or red)
+    if (r > 200 and g < 100 and b < 100) or (g > 180 and r < 100) or (b > 180):
+        return True
+
+    # 🧠 3️⃣ Color tone validation — real stool is usually brown (RGB ratios)
+    brownish_score = (r * 0.5 + g * 0.3 + b * 0.2)
+    grayish_score = abs(r - g) + abs(g - b)
+    if brownish_score < 60 or grayish_score > 80:
+        return True
+
+    # 🧠 4️⃣ Texture detection with OpenCV
+    gray = cv2.cvtColor(np_img, cv2.COLOR_RGB2GRAY)
+    lap_var = cv2.Laplacian(gray, cv2.CV_64F).var()  # blur measure
+    if lap_var < 50:
+        return True  # image too smooth or blurry (likely not stool)
+
+    return False  # passes all checks, likely stool
+
+# 🧠 Main classifier
 def classify_image(image):
     try:
-        # Step 1: Validate stool image
+        # 🧩 Step 1: Reject non-stool images first
         if is_not_stool_image(image):
             return {"Not stool image": 1.0}, diagnosis_advice["Not stool image"]
 
-        # Step 2: Run TFLite model
+        # 🧩 Step 2: Run TFLite model
         input_data = preprocess_image(image)
         interpreter.set_tensor(input_details[0]['index'], input_data)
         interpreter.invoke()
         output_data = interpreter.get_tensor(output_details[0]['index'])[0]
 
-        # Step 3: Sort predictions by confidence
+        # 🧩 Step 3: Sort predictions
         results = {labels[i]: float(output_data[i]) for i in range(len(labels))}
         sorted_results = dict(sorted(results.items(), key=lambda x: x[1], reverse=True))
         top_label, top_score = list(sorted_results.items())[0]
 
-        # Step 4: Handle low confidence
-        if top_score < 0.4:
+        # 🧩 Step 4: Handle low-confidence predictions
+        if top_score < 0.45:
             return {"Uncertain / unclear stool image": top_score}, diagnosis_advice["Uncertain / unclear stool image"]
 
-        # Step 5: Add readable confidence percentages
+        # 🧩 Step 5: Add readable confidence %
         formatted_results = {
             f"{label} ({score * 100:.2f}%)": score for label, score in sorted_results.items()
         }
 
-        # Step 6: Return label + human diagnosis
+        # 🩺 Step 6: Generate human-readable advice
         advice = diagnosis_advice.get(top_label, "No advice available for this diagnosis.")
         return formatted_results, advice
 
     except Exception as e:
-        return {"Error": str(e)}, f"⚠️ Error encountered: {str(e)}"
+        return {"Error": 0.0}, f"⚠️ Error: {str(e)}"
 
 # 🎨 Gradio Interface
 demo = gr.Interface(
     fn=classify_image,
-    inputs=gr.Image(type="pil", label="📸 Upload stool image"),
+    inputs=gr.Image(type="pil", label="📸 Upload Stool Image"),
     outputs=[
         gr.Label(num_top_classes=3, label="Predicted Diagnosis & Confidence"),
         gr.Textbox(label="💬 Diagnosis Summary", lines=3),
     ],
-    title="🐾 Pet Stool Diagnosis AI (with Confidence & Advice)",
+    title="🐾 Pet Stool Diagnosis AI (Strict Validation)",
     description=(
-        "Upload a stool image to get an AI-based diagnosis with confidence scores and health advice. "
-        "This tool checks image validity, identifies stool condition, and provides next-step guidance."
+        "Upload a stool image for AI-based diagnosis. The system will first verify if it's a real stool photo, "
+        "then predict the stool type and give a simple health summary."
     ),
 )