Update app.py

app.py

@@ -3,8 +3,11 @@ import json
 import torch
 from torch import nn
 from torchvision import models, transforms
+from huggingface_hub import hf_hub_download
 from PIL import Image
+import requests
 import os
+from io import BytesIO
 
 # Define the number of classes
 num_classes = 2
@@ -26,7 +29,7 @@ def load_model(model_path):
 model_path = download_model()
 model = load_model(model_path)
 
-# Define transformation for image processing
+# Define the transformation for the input image
 transform = transforms.Compose([
     transforms.Resize(256),
     transforms.CenterCrop(224),
@@ -34,53 +37,61 @@ transform = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
 ])
 
-# Function to load and preprocess image
-def load_image_from_path(image_path):
-    if not os.path.exists(image_path):
-        raise FileNotFoundError(f"Image file not found at {image_path}")
-    image = Image.open(image_path)
-    image = transform(image).unsqueeze(0)  # Convert to tensor and add batch dimension
-    return image
+# Function to predict from image content
+def predict_from_image(image):
+    # Ensure the image is a PIL Image
+    if not isinstance(image, Image.Image):
+        raise ValueError("Invalid image format received. Please provide a valid image.")
 
-# Load the model (Example: ResNet50)
-def load_model():
-    model = models.resnet50(pretrained=True)
-    model.fc = nn.Linear(model.fc.in_features, num_classes)
-    model.load_state_dict(torch.load("model.pth"))
-    model.eval()
-    return model
+    # Apply transformations
+    image_tensor = transform(image).unsqueeze(0)
 
-# Predict from image tensor
-def predict(image_tensor):
+    # Predict
     with torch.no_grad():
         outputs = model(image_tensor)
         predicted_class = torch.argmax(outputs, dim=1).item()
-    return predicted_class
 
-# Initialize model
-model = load_model()
+    # Interpret the result
+    if predicted_class == 0:
+        return {"result": "The photo is of fall army worm with problem ID 126."}
+    elif predicted_class == 1:
+        return {"result": "The photo is of a healthy maize image."}
+    else:
+        return {"error": "Unexpected class prediction."}
 
-# Define the Gradio interface function
-def predict_from_file(file_path):
+# Function to handle image from URL or file path
+def predict_from_url_or_path(url=None, path=None):
     try:
-        # Load image from path
-        image_tensor = load_image_from_path(file_path)
-        # Get prediction
-        predicted_class = predict(image_tensor)
-        result = {"result": "Fall armyworm" if predicted_class == 0 else "Healthy maize"}
-        return result
+        # If URL is provided, fetch and process image
+        if url:
+            response = requests.get(url)
+            response.raise_for_status()  # Ensure the request was successful
+            image = Image.open(BytesIO(response.content))
+            return predict_from_image(image)
+        
+        # If path is provided, open the image from the local path
+        elif path:
+            if not os.path.exists(path):
+                return {"error": f"File not found at {path}"}
+            image = Image.open(path)
+            return predict_from_image(image)
+        else:
+            return {"error": "No valid input provided."}
     except Exception as e:
-        return {"error": str(e)}
+        return {"error": f"Failed to process the input: {str(e)}"}
 
-# Gradio Interface
+# Gradio interface
 iface = gr.Interface(
-    fn=predict_from_file,
-    inputs=gr.Textbox(label="Image Path (Local)"),
-    outputs=gr.JSON(),
+    fn=lambda url, path: predict_from_url_or_path(url=url, path=path),
+    inputs=[
+        gr.Textbox(label="Enter Image URL", placeholder="Provide a valid image URL (optional)", optional=True),
+        gr.Textbox(label="Or Enter Local Image Path", placeholder="Provide the local image path (optional)", optional=True),
+    ],
+    outputs=gr.JSON(label="Prediction Result"),
     live=True,
     title="Maize Anomaly Detection",
-    description="Send a local file path via POST request to trigger prediction.",
+    description="Provide either an image URL or a local file path to detect anomalies in maize crops.",
 )
 
-# Launch the Gradio app
-iface.launch(share=True, server_name="0.0.0.0", server_port=7860)
+# Launch the interface
+iface.launch(share=True, show_error=True)