Update app.py

app.py

@@ -8,6 +8,7 @@ import os
 import logging
 import requests
 from io import BytesIO
+import numpy as np
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -17,7 +18,10 @@ logger = logging.getLogger(__name__)
 num_classes = 3
 
 # Confidence threshold for reliable predictions
-CONFIDENCE_THRESHOLD = 0.5  # 50%
+CONFIDENCE_THRESHOLD = 0.8  # 80%
+
+# Entropy threshold for flat probability distribution (to detect non-maize/rice images)
+ENTROPY_THRESHOLD = 0.9  # Lower entropy means a more peaked distribution (more confident)
 
 # Download model from Hugging Face
 def download_model():
@@ -61,6 +65,16 @@ transform = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
 ])
 
+# Function to compute entropy of the probability distribution
+def compute_entropy(probabilities):
+    probs = probabilities.numpy()  # Convert to numpy array
+    # Avoid log(0) by adding a small epsilon
+    probs = np.clip(probs, 1e-10, 1.0)
+    entropy = -np.sum(probs * np.log(probs))
+    # Normalize entropy by the maximum possible entropy (log(num_classes))
+    max_entropy = np.log(num_classes)
+    return entropy / max_entropy
+
 # Prediction function for an uploaded image
 def predict_from_image_url(image_url):
     try:
@@ -82,25 +96,19 @@ def predict_from_image_url(image_url):
 
         # Define class information
         class_info = {
-            0: {"name": "Fall Army Worm", "problem_id": "126"},
-            1: {"name": "Phosphorus Deficiency", "problem_id": "142"},
-            2: {"name": "Bacterial Leaf Blight", "problem_id": "203"}
-        }
-
-        # Construct probabilities dictionary
-        probabilities_dict = {
-            "Fall Army Worm": f"{probabilities[0]*100:.2f}%",
-            "Phosphorus Deficiency": f"{probabilities[1]*100:.2f}%",
-            "Bacterial Leaf Blight": f"{probabilities[2]*100:.2f}%"
+            0: {"name": "Fall Army Worm", "problem_id": "126", "crop": "maize"},
+            1: {"name": "Phosphorus Deficiency", "problem_id": "142", "crop": "maize"},
+            2: {"name": "Bacterial Leaf Blight", "problem_id": "203", "crop": "rice"}
         }
 
         # Validate predicted class index
         if predicted_class not in class_info:
             logger.warning(f"Unexpected class prediction: {predicted_class} for image URL: {image_url}")
             return {
-                "status": "error",
-                "message": f"Unexpected class prediction (index {predicted_class}). Model may not be configured correctly.",
-                "probabilities": probabilities_dict
+                "status": "invalid",
+                "predicted_class": None,
+                "problem_id": None,
+                "confidence": None
             }
 
         # Get predicted class info
@@ -109,6 +117,24 @@ def predict_from_image_url(image_url):
         problem_id = predicted_info["problem_id"]
         confidence = probabilities[predicted_class].item()  # Confidence score for the predicted class
 
+        # Compute entropy of the probability distribution
+        entropy = compute_entropy(probabilities)
+        logger.info(f"Prediction entropy: {entropy:.4f}, confidence: {confidence:.4f} for image URL: {image_url}")
+
+        # Check if the image is likely maize or rice based on entropy and confidence
+        # High entropy (flat distribution) suggests the image may not be maize or rice
+        if entropy > ENTROPY_THRESHOLD:
+            logger.warning(
+                f"High entropy ({entropy:.4f} > {ENTROPY_THRESHOLD}) for image URL: {image_url}. "
+                "Image may not be of maize or rice."
+            )
+            return {
+                "status": "invalid",
+                "predicted_class": predicted_name,
+                "problem_id": problem_id,
+                "confidence": f"{confidence*100:.2f}%"
+            }
+
         # Check confidence threshold
         if confidence < CONFIDENCE_THRESHOLD:
             logger.warning(
@@ -116,29 +142,28 @@ def predict_from_image_url(image_url):
                 f"for image URL: {image_url}"
             )
             return {
-                "status": "uncertain",
-                "message": (
-                    f"Prediction confidence ({confidence*100:.2f}%) is below the threshold ({CONFIDENCE_THRESHOLD*100}%). "
-                    "This image may not belong to the trained classes (Fall Army Worm, Phosphorus Deficiency, Bacterial Leaf Blight)."
-                ),
+                "status": "invalid",
                 "predicted_class": predicted_name,
                 "problem_id": problem_id,
-                "confidence": f"{confidence*100:.2f}%",
-                "probabilities": probabilities_dict
+                "confidence": f"{confidence*100:.2f}%"
             }
 
         # Return successful prediction
         return {
-            "status": "success",
+            "status": "valid",
             "predicted_class": predicted_name,
             "problem_id": problem_id,
-            "confidence": f"{confidence*100:.2f}%",
-            "probabilities": probabilities_dict
+            "confidence": f"{confidence*100:.2f}%"
         }
 
     except Exception as e:
         logger.error(f"Error processing image URL {image_url}: {str(e)}")
-        return {"status": "error", "message": str(e)}
+        return {
+            "status": "invalid",
+            "predicted_class": None,
+            "problem_id": None,
+            "confidence": None
+        }
 
 # Gradio interface
 demo = gr.Interface(