Update app.py

app.py

@@ -3,53 +3,79 @@ import tensorflow as tf
 import numpy as np
 from PIL import Image
 
-# Load the TFLite model
+# 🧠 Load your TFLite model
 interpreter = tf.lite.Interpreter(model_path="stool_model.tflite")
 interpreter.allocate_tensors()
 
 input_details = interpreter.get_input_details()
 output_details = interpreter.get_output_details()
 
-# Define label names (adjust if your model uses different ones)
+# 🏷️ Define your classes
 labels = ["bloody", "hard stool", "normal", "parasite", "watery"]
 
+# 🧩 Image preprocessing
 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
 
+# 🚫 Detect if the uploaded image is NOT a stool image
+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))
+
+    # 🧠 Basic heuristic checks
+    # These values are adjustable based on your dataset
+    if brightness > 220 or brightness < 25:
+        return True  # too bright or dark
+    if contrast < 25:
+        return True  # too flat / low texture
+    if avg_color[0] > 180 and avg_color[1] < 80 and avg_color[2] < 80:
+        return True  # too red
+    if avg_color[0] < 50 and avg_color[1] > 180:
+        return True  # too greenish
+    if avg_color[2] > 200:
+        return True  # too blueish
+
+    return False
+
+# 🧠 Classification function
 def classify_image(image):
     try:
-        img = image.convert("RGB")
-        arr = np.asarray(img).astype(np.float32)
-        brightness = arr.mean()
-        contrast = arr.std()
-
-        # 🧠 Simple sanity check for stool-like features (brownish tone + moderate contrast)
-        # You can adjust these thresholds depending on your dataset
-        if brightness > 220 or brightness < 20 or contrast < 25:
+        # 🚫 Check if this is not stool
+        if is_not_stool_image(image):
             return {"Not stool image": 1.0}
 
-        # Normal prediction process
+        # ✅ Proceed with model prediction
         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]
+
+        # 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))
+
+        # Extra sanity rule: if top score < 0.4, label as uncertain
+        top_label, top_score = list(sorted_results.items())[0]
+        if top_score < 0.4:
+            return {"Uncertain / unclear stool image": top_score}
+
         return sorted_results
 
     except Exception as e:
-        return {"error": str(e)}
-
+        return {"Error": str(e)}
 
+# 🎨 Gradio UI
 demo = gr.Interface(
     fn=classify_image,
-    inputs=gr.Image(type="pil", label="Upload stool image"),
-    outputs=gr.Label(num_top_classes=3, label="Predictions"),
-    title="Stool Diagnosis Model",
-    description="Upload a stool image for classification."
+    inputs=gr.Image(type="pil", label="📸 Upload stool image"),
+    outputs=gr.Label(num_top_classes=3, label="Predicted diagnosis"),
+    title="🐾 Stool Diagnosis AI",
+    description="Upload a stool image for analysis. The model predicts stool type or rejects unrelated photos."
 )
 
 if __name__ == "__main__":