2

app.py

@@ -24,10 +24,7 @@ async def predict_ndvi_api(file: UploadFile = File(...)):
         norm_img = normalize_rgb(np.array(img))
         pred_ndvi = predict_ndvi(model, norm_img)
 
-        vis_img = create_visualization(norm_img, pred_ndvi)  # <-- should return PIL.Image
-
-        img_byte_arr = BytesIO()
-        vis_img.save(img_byte_arr, format='PNG')  # save image into buffer
+        img_byte_arr = create_visualization(norm_img, pred_ndvi)  # already a BytesIO object
         img_byte_arr.seek(0)
 
         return StreamingResponse(img_byte_arr, media_type="image/png")