Update main.py

main.py

@@ -1,20 +1,12 @@
 from fastapi import FastAPI, HTTPException
+from fastapi.responses import JSONResponse
 from argopy import DataFetcher
 import matplotlib.pyplot as plt
-import uuid
-import os
+import io
+import base64
 
 app = FastAPI()
 
-# Directory where images will be saved
-# This directory must be publicly accessible through your web server
-#image_directory = "/images"
-#os.makedirs(image_directory, exist_ok=True)
-
-# Base URL for the images
-# Replace with your server's base URL
-base_url = "https://zliang-argoapi.hf.space"
-
 @app.get("/plot_trajectory")
 def plot_trajectory(lon_min: float, lon_max: float, lat_min: float, lat_max: float, 
                     depth_min: float, depth_max: float, 
@@ -27,21 +19,22 @@ def plot_trajectory(lon_min: float, lon_max: float, lat_min: float, lat_max: flo
         plt.figure()
         f.plot('trajectory', add_legend=False)
 
-        # Save the plot to a file with a unique name
-        filename = str(uuid.uuid4()) + ".png"
-        plt.savefig(filename)
+        # Save the plot to a BytesIO buffer
+        buf = io.BytesIO()
+        plt.savefig(buf, format='png')
+        buf.seek(0)
         plt.close()
 
-        # Construct the URL for the saved image
-        image_url = f"{base_url}/{filename}"
+        # Encode the image in base64 and prepare it for JSON response
+        encoded_img = base64.b64encode(buf.read()).decode('utf-8')
+        response = {"image_base64": encoded_img}
 
-        return {"image_url": image_url}
+        return JSONResponse(content=response)
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
 
-
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)