Update main.py

main.py

@@ -27,11 +27,15 @@ async def upload_file(file: UploadFile = File(...)):
     try:
         result = client.predict(temp_file_path, api_name="/predict")
         
-        # Convert the image data to base64 string
-        with open(result[0], "rb") as image_file:
-            image_data = base64.b64encode(image_file.read()).decode("utf-8")
-
-        return {"sketch_image_base64": image_data}
+        # Check if the result is valid
+        if result and len(result) > 0:
+            # Convert the image data to base64 string
+            with open(result[0], "rb") as image_file:
+                image_data = base64.b64encode(image_file.read()).decode("utf-8")
+
+            return {"sketch_image_base64": f"data:image/png;base64,{image_data}"}
+        else:
+            return {"error": "No sketch image generated."}
     finally:
         os.unlink(temp_file_path)