Update app.py

app.py

@@ -44,21 +44,31 @@ transform = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
 ])
 
-def predict_from_image(image):
+def predict(input_data):
     try:
-        # Check if the input is a PIL Image type (Gradio automatically provides a PIL image)
-        if not isinstance(image, Image.Image):
+        # Check if the input is a URL or image (Gradio will pass either based on user input)
+        if isinstance(input_data, str):  # If it's a string, assume it's a URL
+            response = requests.get(input_data)
+            if response.status_code == 200:
+                img = Image.open(BytesIO(response.content))
+            else:
+                return json.dumps({"error": "Unable to fetch image from the URL."})
+        else:  # If it's not a string, assume it's an image file
+            img = input_data
+
+        # Check if the input is a PIL Image type
+        if not isinstance(img, Image.Image):
             return json.dumps({"error": "Invalid image format received. Please provide a valid image."})
 
         # Apply transformations to the image
-        image = transform(image).unsqueeze(0)
+        img = transform(img).unsqueeze(0)
 
         # Move the image to the correct device
-        image = image.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+        img = img.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
 
         # Make predictions
         with torch.no_grad():
-            outputs = model(image)
+            outputs = model(img)
             predicted_class = torch.argmax(outputs, dim=1).item()
 
         # Return the result based on the predicted class
@@ -73,23 +83,9 @@ def predict_from_image(image):
         return json.dumps({"error": f"Error processing image: {e}"})
 
 
-def predict_from_url(url):
-    try:
-        # Check if the URL is valid and try fetching the image
-        response = requests.get(url)
-        if response.status_code == 200:
-            img = Image.open(BytesIO(response.content))
-            # Call the predict function for the image
-            return predict_from_image(img)
-        else:
-            return json.dumps({"error": "Unable to fetch image from the URL."})
-    except Exception as e:
-        return json.dumps({"error": f"Error fetching image from URL: {e}"})
-
-
 # Create the Gradio interface with both local file upload and URL input
 iface = gr.Interface(
-    fn=predict_from_image,
+    fn=predict,
     inputs=[gr.Image(type="pil", label="Upload an image or provide a local path"), 
             gr.Textbox(label="Or enter image URL (if available)", placeholder="Enter a URL for the image")],
     outputs=gr.Textbox(label="Prediction Result"),