1

app.py

@@ -21,16 +21,13 @@ async def predict_ndvi_api(file: UploadFile = File(...)):
         contents = await file.read()
         img = Image.open(BytesIO(contents)).convert("RGB")
 
-        # Preprocess & predict
         norm_img = normalize_rgb(np.array(img))
         pred_ndvi = predict_ndvi(model, norm_img)
 
-        # Create visualization
-        vis_img = create_visualization(norm_img, pred_ndvi)
+        vis_img = create_visualization(norm_img, pred_ndvi)  # <-- should return PIL.Image
 
-        # Convert to bytes
         img_byte_arr = BytesIO()
-        vis_img.save(img_byte_arr, format='PNG')
+        vis_img.save(img_byte_arr, format='PNG')  # save image into buffer
         img_byte_arr.seek(0)
 
         return StreamingResponse(img_byte_arr, media_type="image/png")