Upload folder using huggingface_hub

image_caption/get_caption.py

@@ -12,6 +12,7 @@ import json
 import requests
 import base64
 from io import BytesIO
+import PIL
 
 def get_image_caption(img_obj, prompt=None):
     """
@@ -27,7 +28,17 @@ def get_image_caption(img_obj, prompt=None):
 
 
     buffered = BytesIO()
-    img_obj.save(buffered, format="JPEG")
+    if img_obj is None:
+        return "No image data provided!"
+    try:
+        img_obj.save(buffered, format="JPEG")
+    except PIL.UnidentifiedImageError as e:
+        print("Error in saving the image", e)
+        return "Invalid image format!"
+    except Exception as e:
+        print("Some error occurred while saving the image", e)
+        return "Some error occurred while loading the image!"
+
     img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
     url = os.environ.get("IMAGE_CAPTION_ENDPOINT")
     auth_token = os.environ.get("IMAGE_CAPTION_AUTH_TOKEN")