Update app.py

app.py

@@ -17,24 +17,39 @@ CONFIDENCE_THRESHOLD = 0.25
 
 def predict_image(image):
     try:
+        # Preprocess
         image = image.resize((224, 224)).convert("RGB")
         img_array = np.array(image, dtype=np.float32) / 255.0
         img_array = np.expand_dims(img_array, axis=0)
 
+        # Run inference
         interpreter.set_tensor(input_details[0]['index'], img_array)
         interpreter.invoke()
-        output = interpreter.get_tensor(output_details[0]['index'])
+        output = interpreter.get_tensor(output_details[0]['index'])[0]  # shape (num_classes,)
 
-        class_idx = int(np.argmax(output))
-        confidence = float(np.max(output))
+        # Normalize if needed (sometimes TFLite outputs logits)
+        probs = tf.nn.softmax(output).numpy()
+
+        # Get predicted class
+        class_idx = int(np.argmax(probs))
+        confidence = float(np.max(probs))
+
+        # Format output (show every class probability)
+        results = []
+        for i, prob in enumerate(probs):
+            results.append(f"{class_names[i]}: {prob*100:.2f}%")
+
+        results_text = "\n".join(results)
 
         if confidence < CONFIDENCE_THRESHOLD:
-            return f"⚠️ Low confidence ({confidence:.2f}). The model is unsure. Please try a clearer image."
+            return f"⚠️ Low confidence ({confidence:.2f}). The model is unsure.\n\nProbabilities:\n{results_text}"
         else:
-            return f"✅ Prediction: {class_names[class_idx]}"
+            return f"✅ Prediction: {class_names[class_idx]} ({confidence*100:.2f}%)\n\nProbabilities:\n{results_text}"
+
     except Exception as e:
         return f"Error: {str(e)}"
 
+# Gradio UI
 gr.Interface(
     fn=predict_image,
     inputs=gr.Image(type="pil"),