Update app.py

app.py

@@ -1,51 +1,144 @@
 from fastapi.responses import StreamingResponse, JSONResponse
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, HTTPException
 from ndvi_predictor import load_model, normalize_rgb, predict_ndvi, create_visualization
+from yolo_predictor import load_yolo_model, predict_yolo, predict_pipeline
 from PIL import Image
 from io import BytesIO
 import numpy as np
 import zipfile
+import json
+import rasterio
+from rasterio.transform import from_bounds
+import tempfile
+import os
 
 app = FastAPI()
-model = load_model("ndvi_best_model.keras")
+
+# Load models at startup
+ndvi_model = load_model("ndvi_best_model.keras")
+yolo_model = load_yolo_model("4c_6c_regression.pt")
 
 @app.get("/")
 async def root():
-    return {"message": "Welcome to the NDVI prediction API!"}
+    return {"message": "Welcome to the NDVI and YOLO prediction API!"}
 
-@app.post("/predict/")
+@app.post("/predict_ndvi/")
 async def predict_ndvi_api(file: UploadFile = File(...)):
+    """Predict NDVI from RGB image"""
     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)
-
+        pred_ndvi = predict_ndvi(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)})
 
+@app.post("/predict_yolo/")
+async def predict_yolo_api(file: UploadFile = File(...)):
+    """Predict YOLO results from 4-channel TIFF image"""
+    try:
+        # Save uploaded file temporarily
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.tif') as tmp_file:
+            contents = await file.read()
+            tmp_file.write(contents)
+            tmp_file_path = tmp_file.name
+        
+        try:
+            # Predict using YOLO model
+            results = predict_yolo(yolo_model, tmp_file_path)
+            
+            # Convert results to JSON-serializable format
+            results_dict = {
+                "boxes": {
+                    "xyxyn": results.boxes.xyxyn.tolist() if results.boxes is not None else None,
+                    "conf": results.boxes.conf.tolist() if results.boxes is not None else None,
+                    "cls": results.boxes.cls.tolist() if results.boxes is not None else None
+                },
+                "classes": results.boxes.cls.tolist() if results.boxes is not None else None,
+                "names": results.names,
+                "growth_stages": getattr(results, 'growth_stages', None),
+                "orig_shape": results.orig_shape,
+                "speed": results.speed
+            }
+            
+            # Handle growth stages if present
+            if hasattr(results, 'boxes') and hasattr(results.boxes, 'data'):
+                # Extract growth stages from the results if available
+                if len(results.boxes.data[0]) > 6:  # Assuming growth stages are in the data
+                    growth_stages = results.boxes.data[:, 6].tolist()
+                    results_dict["growth_stages"] = growth_stages
+            
+            return JSONResponse(content=results_dict)
+            
+        finally:
+            # Clean up temporary file
+            os.unlink(tmp_file_path)
+            
     except Exception as e:
         return JSONResponse(status_code=500, content={"error": str(e)})
+
+@app.post("/predict_pipeline/")
+async def predict_pipeline_api(file: UploadFile = File(...)):
+    """Full pipeline: RGB -> NDVI -> 4-channel -> YOLO prediction"""
+    try:
+        # Save uploaded file temporarily
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.tif') as tmp_file:
+            contents = await file.read()
+            tmp_file.write(contents)
+            tmp_file_path = tmp_file.name
+        
+        try:
+            # Run the full pipeline
+            results = predict_pipeline(ndvi_model, yolo_model, tmp_file_path)
+            
+            # Convert results to JSON-serializable format
+            results_dict = {
+                "boxes": {
+                    "xyxyn": results.boxes.xyxyn.tolist() if results.boxes is not None else None,
+                    "conf": results.boxes.conf.tolist() if results.boxes is not None else None,
+                    "cls": results.boxes.cls.tolist() if results.boxes is not None else None
+                },
+                "classes": results.boxes.cls.tolist() if results.boxes is not None else None,
+                "names": results.names,
+                "growth_stages": getattr(results, 'growth_stages', None),
+                "orig_shape": results.orig_shape,
+                "speed": results.speed
+            }
+            
+            # Handle growth stages if present
+            if hasattr(results, 'boxes') and hasattr(results.boxes, 'data'):
+                if len(results.boxes.data[0]) > 6:
+                    growth_stages = results.boxes.data[:, 6].tolist()
+                    results_dict["growth_stages"] = growth_stages
+            
+            return JSONResponse(content=results_dict)
+            
+        finally:
+            # Clean up temporary file
+            os.unlink(tmp_file_path)
+            
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": str(e)})