Upload handler.py

handler.py

@@ -6,6 +6,8 @@ import base64
 from ultralytics import YOLO
 import os
 import gdown
+from PIL import Image
+import io
 
 class EndpointHandler:
     def __init__(self, path='.'):  # pass api key to model
@@ -72,6 +74,15 @@ class EndpointHandler:
             cv2.imwrite("result.jpg", clothes)  
         # create base 64 object
         jpg_as_text = base64.b64encode(buffer).decode("utf-8") # Decode bytes to string")
+        # base64_encoded = base64.b64encode(image_bytes).decode("utf-8")
+        
+        # Get the image format
+        image_format = Image.open(io.BytesIO(buffer)).format.lower()
+        
+        # Construct the data URI
+        data_uri = f"data:image/{image_format};base64,{jpg_as_text}"
+        
+        return data_uri
         print("checkpoint 4")
 ###########################################################################
         return jpg_as_text