Update main.py

main.py

@@ -1,12 +1,20 @@
 from fastapi import FastAPI, HTTPException
 from argopy import DataFetcher
 import matplotlib.pyplot as plt
-from fastapi.responses import StreamingResponse
-import io
-
+import uuid
+import os
 
 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/images"
+
 @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, 
@@ -18,17 +26,18 @@ def plot_trajectory(lon_min: float, lon_max: float, lat_min: float, lat_max: flo
         # Plotting
         plt.figure()
         f.plot('trajectory', add_legend=False)
-        buf = io.BytesIO()
-        plt.savefig(buf, format='png')
-        buf.seek(0)
+
+        # Save the plot to a file with a unique name
+        filename = str(uuid.uuid4()) + ".png"
+        file_path = os.path.join(image_directory, filename)
+        plt.savefig(file_path)
         plt.close()
-        StreamingResponse(buf, media_type="image/png")
-        
-        image_url = f"{base_url}plot_trajectory?lon_min={lon_min}&lon_max={lon_max}&lat_min={lat_min}&lat_max={lat_max}&depth_min={depth_min}&depth_max={depth_max}&date_start={date_start}&date_end={date_end}"
 
+        # Construct the URL for the saved image
+        image_url = f"{base_url}/{filename}"
 
         return {"image_url": image_url}
-        
+
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))