Update app.py

app.py

@@ -20,13 +20,28 @@ pipe = pipeline(task="image-classification", model=model_repo)
 # Define a custom prediction function
 import numpy as np
 
-# Define a custom prediction function
+from PIL import Image
+from io import BytesIO
+import numpy as np
+import base64
+
 def predict(image):
-    # Convert the image to numpy array if it's not already in that format
-    if not isinstance(image, np.ndarray):
-        raise ValueError("Input image must be a numpy array.")
+    # Check if the input image is a base64 encoded string
+    if isinstance(image, str):
+        # Decode the base64 encoded image string and convert it to a PIL image object
+        image_data = BytesIO(base64.b64decode(image))
+        pil_image = Image.open(image_data)
+        # Convert the PIL image to a numpy array
+        image_np = np.array(pil_image)
+    elif isinstance(image, np.ndarray):
+        # If the input image is already a numpy array, use it directly
+        image_np = image
+    else:
+        # If the input is neither a base64 encoded string nor a numpy array, raise an error
+        raise ValueError("Input image must be either a base64 encoded string or a numpy array.")
+
     # Get the predictions from the pipeline
-    predictions = pipe(image)
+    predictions = pipe(image_np)
     # Get the predicted label index
     predicted_index = predictions[0]['label']
     # Map the index to the corresponding disease name using id2label
@@ -36,6 +51,7 @@ def predict(image):
     return f"{label_name} ({confidence_score:.2f})"
 
 
+
 # Create Gradio interface
 iface = gr.Interface(fn=predict,
                      inputs=gr.Image(type="numpy"),