Update app.py

app.py

@@ -4,6 +4,9 @@ from torch import nn
 from torchvision import models, transforms
 from huggingface_hub import hf_hub_download
 from PIL import Image
+import requests
+import base64
+from io import BytesIO
 
 # Define the number of classes
 num_classes = 2  # Update with the actual number of classes in your dataset (e.g., 2 for healthy and anomalous)
@@ -33,8 +36,30 @@ transform = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),  # Normalize the image (ImageNet mean and std)
 ])
 
+# Function to convert image from URL to PIL image
+def url_to_image(image_url):
+    response = requests.get(image_url)
+    img = Image.open(BytesIO(response.content))
+    return img
+
+# Function to convert base64 string to PIL image
+def base64_to_pil(base64_string):
+    img_data = base64.b64decode(base64_string)
+    return Image.open(BytesIO(img_data))
+
 # Define the prediction function
-def predict(image):
+def predict(image_input):
+    # If input is a string (URL or base64 encoded), handle accordingly
+    if isinstance(image_input, str):
+        if image_input.startswith("http"):  # If URL
+            image = url_to_image(image_input)
+        elif image_input.startswith("data:image"):  # If base64 string
+            image = base64_to_pil(image_input)
+        else:  # Local image path
+            image = Image.open(image_input)
+    else:
+        image = image_input  # If the input is already a PIL image
+
     # Apply the necessary transformations to the image
     image = transform(image).unsqueeze(0)  # Add batch dimension
     image = image.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))  # Move to GPU if available
@@ -57,9 +82,9 @@ iface = gr.Interface(
     inputs=gr.Image(type="pil"),  # Image input
     outputs=gr.Textbox(),  # Output: Predicted class
     live=True,  # Updates as the user uploads an image
-    title="maize Anomaly Detection",
-    description="Upload an image of maize to detect anomalies like disease or pest infestation."
+    title="Maize Anomaly Detection",
+    description="Upload an image of maize to detect anomalies like disease or pest infestation. You can provide local paths, URLs, or base64-encoded images."
 )
 
 # Launch the Gradio interface
-iface.launch(share=True)
+iface.launch(share=True)  # Share=True to make the app accessible online