Update app.py

app.py

@@ -13,31 +13,23 @@ model = load_model("ndvi_best_model.keras")
 
 @app.post("/predict/")
 async def predict_ndvi_api(file: UploadFile = File(...)):
-    contents = await file.read()
-    img = Image.open(BytesIO(contents)).convert("RGB")
-    rgb_np = np.array(img)
-    rgb_norm = normalize_rgb(rgb_np)
-    ndvi = predict_ndvi(model, rgb_norm)
+    try:
+        contents = await file.read()
+        img = Image.open(BytesIO(contents)).convert("RGB")
 
-    # Generate visual output
-    vis_image = create_visualization(rgb_np, ndvi)
-    vis_filename = "ndvi_visualization.png"
+        # Preprocess & predict
+        norm_img = normalize_rgb(np.array(img))
+        pred_ndvi = predict_ndvi(model, norm_img)
 
-    # Generate GeoTIFF in memory
-    geotiff_buffer = io.BytesIO()
-    import rasterio
-    profile = {
-        'driver': 'GTiff',
-        'height': ndvi.shape[0],
-        'width': ndvi.shape[1],
-        'count': 1,
-        'dtype': 'float32'
-    }
-    with rasterio.open(geotiff_buffer, 'w', **profile) as dst:
-        dst.write(ndvi, 1)
-    geotiff_buffer.seek(0)
+        # Create visualization
+        vis_img = create_visualization(norm_img, pred_ndvi)
+
+        # Convert to bytes
+        img_byte_arr = BytesIO()
+        vis_img.save(img_byte_arr, format='PNG')
+        img_byte_arr.seek(0)
+
+        return StreamingResponse(img_byte_arr, media_type="image/png")
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": str(e)})
 
-    return {
-        "ndvi_tiff": StreamingResponse(geotiff_buffer, media_type="image/tiff", headers={"Content-Disposition": "attachment; filename=predicted_ndvi.tif"}),
-        "visualization": StreamingResponse(vis_image, media_type="image/png", headers={"Content-Disposition": "attachment; filename=ndvi_visualization.png"})
-    }