app.py

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 zipfile

app = FastAPI()
model = load_model("ndvi_best_model.keras")

@app.get("/")
async def root():
    return {"message": "Welcome to the NDVI prediction API!"}

@app.post("/predict/")
async def predict_ndvi_api(file: UploadFile = File(...)):
    try:
        contents = await file.read()
        img = Image.open(BytesIO(contents)).convert("RGB")

        norm_img = normalize_rgb(np.array(img))
        pred_ndvi = predict_ndvi(model, norm_img)

        # 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)})