Update app.py

app.py

@@ -17,6 +17,10 @@ CONFIDENCE_THRESHOLD = 0.7
 
 def predict_image(image):
     try:
+        # Validate input
+        if image.mode != "RGB":
+            return {"Error": 1.0}, "⚠️ Please upload a valid cervical cell image (RGB image required)."
+
         # Preprocess
         image = image.resize((224, 224)).convert("RGB")
         img_array = np.array(image, dtype=np.float32) / 255.0
@@ -27,22 +31,25 @@ def predict_image(image):
         interpreter.invoke()
         output = interpreter.get_tensor(output_details[0]['index'])[0]  # shape (num_classes,)
 
-        # Normalize if needed
+        # Normalize
         probs = tf.nn.softmax(output).numpy()
 
+        # Check if prediction is below confidence threshold
+        if np.max(probs) < CONFIDENCE_THRESHOLD:
+            return {"Error": 1.0}, "⚠️ The model is unsure. Please upload a clearer/correct image of cervical cells."
+
         # Convert to dict for Gradio Label
         probs_dict = {class_names[i]: float(probs[i]) for i in range(len(class_names))}
-
-        return probs_dict
+        return probs_dict, f"✅ Prediction: {class_names[np.argmax(probs)]} ({np.max(probs)*100:.2f}%)"
 
     except Exception as e:
-        return {"Error": 1.0}  # dummy error output
+        return {"Error": 1.0}, f"⚠️ Something went wrong. Please upload a correct cervical cell image. ({str(e)})"
 
 # Gradio UI
 gr.Interface(
     fn=predict_image,
     inputs=gr.Image(type="pil"),
-    outputs=gr.Label(num_top_classes=len(class_names)),  # shows all classes with bars
+    outputs=[gr.Label(num_top_classes=len(class_names)), gr.Textbox()],
     title="Cervical Cancer Classification",
-    description="Upload an image. The model shows probabilities for each class."
+    description="Upload an image. The model shows probabilities for each class and warns if the image is incorrect."
 ).launch()