Update app.py

app.py

@@ -6,9 +6,9 @@ from torchvision import models, transforms
 from huggingface_hub import hf_hub_download
 from PIL import Image
 import requests
+import os
 import base64
 from io import BytesIO
-import os
 
 # Define the number of classes
 num_classes = 2
@@ -46,46 +46,22 @@ transform = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
 ])
 
-def predict(data):
+def predict(image):
     try:
-        # Check if the data is a list and not empty
-        if not isinstance(data, list) or len(data) == 0:
-            return json.dumps({"error": "Input data should be a non-empty list."})
-        
-        # Extract the image path
-        image_input = data[0].get('path', None)
-        if not image_input:
-            return json.dumps({"error": "No image path provided."})
+        print(f"Received image input: {image}")
 
-        print(f"Received image input: {image_input}")
-
-        # Handle URLs
-        if isinstance(image_input, str) and (image_input.startswith("http://") or image_input.startswith("https://")):
-            try:
-                response = requests.get(image_input)
-                response.raise_for_status()  # Check for HTTP errors
-                image = Image.open(BytesIO(response.content))
-                print(f"Fetched image from URL: {image}")
-            except Exception as e:
-                print(f"Error fetching image from URL: {e}")
-                return json.dumps({"error": f"Error fetching image from URL: {e}"})
-        
-        # Check if the image path is a valid local path
-        elif isinstance(image_input, str) and os.path.exists(image_input):
-            try:
-                image = Image.open(image_input)
-                print(f"Loaded image from local path: {image}")
-            except Exception as e:
-                return json.dumps({"error": f"Error loading image from local path: {e}"})
-        
-        else:
-            return json.dumps({"error": "Invalid image path. Ensure it's a valid URL or local path."})
+        # Check if the input is a PIL Image type (Gradio already provides a PIL image)
+        if not isinstance(image, Image.Image):
+            return json.dumps({"error": "Invalid image format received. Please provide a valid image."})
 
-        # Apply the transformations and make prediction
+        # Apply transformations to the image
         image = transform(image).unsqueeze(0)
         print(f"Transformed image tensor: {image.shape}")
+
+        # Move the image to the correct device
         image = image.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
 
+        # Make predictions
         with torch.no_grad():
             outputs = model(image)
             predicted_class = torch.argmax(outputs, dim=1).item()
@@ -98,7 +74,7 @@ def predict(data):
             return json.dumps({"result": "The photo you've sent is of a healthy maize image."})
         else:
             return json.dumps({"error": "Unexpected class prediction."})
-    
+
     except Exception as e:
         print(f"Error processing image: {e}")
         return json.dumps({"error": f"Error processing image: {e}"})
@@ -106,7 +82,7 @@ def predict(data):
 # Create the Gradio interface
 iface = gr.Interface(
     fn=predict,
-    inputs=gr.JSON(label="Input JSON"),
+    inputs=gr.Image(type="pil", label="Upload an image or provide a URL or local path"),
     outputs=gr.Textbox(label="Prediction Result"),
     live=True,
     title="Maize Anomaly Detection",
@@ -115,4 +91,3 @@ iface = gr.Interface(
 
 # Launch the Gradio interface 
 iface.launch(share=True, show_error=True)
-