Update app.py

app.py

@@ -11,10 +11,14 @@ from io import BytesIO
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Define the number of classes
 num_classes = 3
 
+# Confidence threshold for reliable predictions
+CONFIDENCE_THRESHOLD = 0.5  # 50%
+
 # Download model from Hugging Face
 def download_model():
     model_path = hf_hub_download(repo_id="jays009/Resnet3", filename="model.pth")
@@ -26,7 +30,7 @@ def load_model(model_path):
     num_features = model.fc.in_features
     model.fc = nn.Sequential(
         nn.Dropout(0.5),
-        nn.Linear(num_features, 3)  # 3 classes
+        nn.Linear(num_features, num_classes)  # 3 classes
     )
 
     # Load the checkpoint
@@ -67,48 +71,74 @@ def predict_from_image_url(image_url):
 
         # Apply transformations
         image_tensor = transform(image).unsqueeze(0)  # Shape: [1, 3, 224, 224]
-        print(f"Input image tensor shape: {image_tensor.shape}")  # Debug: Should be [1, 3, 224, 224]
+        logger.info(f"Input image tensor shape: {image_tensor.shape}")
 
         # Perform prediction
         with torch.no_grad():
             outputs = model(image_tensor)  # Shape: [1, 3]
-            print(f"Model output shape: {outputs.shape}")  # Debug: Should be [1, 3]
+            logger.info(f"Model output shape: {outputs.shape}")
             probabilities = torch.softmax(outputs, dim=1)[0]  # Convert to probabilities
             predicted_class = torch.argmax(outputs, dim=1).item()
 
-        # Interpret the result
-        if predicted_class == 0:
-            return {
-                "result": "The photo is of Fall Army Worm with problem ID 126.",
-                "probabilities": {
-                    "Fall Army Worm": f"{probabilities[0]*100:.2f}%",
-                    "Phosphorus Deficiency": f"{probabilities[1]*100:.2f}%",
-                    "Bacterial Leaf Blight": f"{probabilities[2]*100:.2f}%"
-                }
-            }
-        elif predicted_class == 1:
+        # 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}%"
+        }
+
+        # 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 {
-                "result": "The photo shows symptoms of Phosphorus Deficiency with Problem ID 142.",
-                "probabilities": {
-                    "Fall Army Worm": f"{probabilities[0]*100:.2f}%",
-                    "Phosphorus Deficiency": f"{probabilities[1]*100:.2f}%",
-                    "Bacterial Leaf Blight": f"{probabilities[2]*100:.2f}%"
-                }
+                "status": "error",
+                "message": f"Unexpected class prediction (index {predicted_class}). Model may not be configured correctly.",
+                "probabilities": probabilities_dict
             }
-        elif predicted_class == 2:
+
+        # Get predicted class info
+        predicted_info = class_info[predicted_class]
+        predicted_name = predicted_info["name"]
+        problem_id = predicted_info["problem_id"]
+        confidence = probabilities[predicted_class].item()  # Confidence score for the predicted class
+
+        # Check confidence threshold
+        if confidence < CONFIDENCE_THRESHOLD:
+            logger.warning(
+                f"Low confidence prediction: {predicted_name} with confidence {confidence*100:.2f}% "
+                f"for image URL: {image_url}"
+            )
             return {
-                "result": "The photo shows symptoms of Bacterial Leaf Blight with Problem ID 203.",
-                "probabilities": {
-                    "Fall Army Worm": f"{probabilities[0]*100:.2f}%",
-                    "Phosphorus Deficiency": f"{probabilities[1]*100:.2f}%",
-                    "Bacterial Leaf Blight": f"{probabilities[2]*100:.2f}%"
-                }
+                "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)."
+                ),
+                "predicted_class": predicted_name,
+                "problem_id": problem_id,
+                "confidence": f"{confidence*100:.2f}%",
+                "probabilities": probabilities_dict
             }
-        else:
-            return {"error": "Unexpected class prediction."}
+
+        # Return successful prediction
+        return {
+            "status": "success",
+            "predicted_class": predicted_name,
+            "problem_id": problem_id,
+            "confidence": f"{confidence*100:.2f}%",
+            "probabilities": probabilities_dict
+        }
 
     except Exception as e:
-        return {"error": str(e)}
+        logger.error(f"Error processing image URL {image_url}: {str(e)}")
+        return {"status": "error", "message": str(e)}
 
 # Gradio interface
 demo = gr.Interface(