.npy and png in zip output

app.py

@@ -1,13 +1,10 @@
-# app.py
-from fastapi import FastAPI, File, UploadFile
 from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi import FastAPI, File, UploadFile
 from ndvi_predictor import load_model, normalize_rgb, predict_ndvi, create_visualization
 from PIL import Image
 from io import BytesIO
 import numpy as np
-import io
-import os
-import base64
+import zipfile
 
 app = FastAPI()
 model = load_model("ndvi_best_model.keras")
@@ -25,22 +22,30 @@ async def predict_ndvi_api(file: UploadFile = File(...)):
         norm_img = normalize_rgb(np.array(img))
         pred_ndvi = predict_ndvi(model, norm_img)
 
-        # Create preview (small visualization)
-        img_io = create_visualization(norm_img, pred_ndvi)
-        img_io.seek(0)
-        preview_img = Image.open(img_io)
-        preview_img.thumbnail((256, 256))  # Resize for fast rendering
-        preview_bytes = BytesIO()
-        preview_img.save(preview_bytes, format="PNG")
-        preview_bytes.seek(0)
-
-        img_base64 = base64.b64encode(preview_bytes.read()).decode("utf-8")
-        ndvi_array = pred_ndvi.tolist()  # Convert NumPy array to list
-
-        return {
-            "preview_base64": img_base64,
-            "ndvi_array": ndvi_array  # This can be big depending on image size
-        }
+        # Visualization image as PNG
+        vis_img_bytes = create_visualization(norm_img, pred_ndvi)
+        vis_img_bytes.seek(0)
+
+        # NDVI band as .npy
+        ndvi_bytes = BytesIO()
+        np.save(ndvi_bytes, pred_ndvi)
+        ndvi_bytes.seek(0)
+
+        # Create a ZIP containing both files
+
+        zip_buf = BytesIO()
+        with zipfile.ZipFile(zip_buf, "w") as zip_file:
+            zip_file.writestr("ndvi_image.png", vis_img_bytes.read())
+            ndvi_bytes.seek(0)
+            zip_file.writestr("ndvi_band.npy", ndvi_bytes.read())
+
+        zip_buf.seek(0)
+
+        return StreamingResponse(
+            zip_buf,
+            media_type="application/x-zip-compressed",
+            headers={"Content-Disposition": "attachment; filename=ndvi_output.zip"}
+        )
 
     except Exception as e:
         return JSONResponse(status_code=500, content={"error": str(e)})