Update app.py

app.py

@@ -11,12 +11,28 @@ import rasterio
 from rasterio.transform import from_bounds
 import tempfile
 import os
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
 # Load models at startup
-ndvi_model = load_model("ndvi_best_model.keras")
-yolo_model = load_yolo_model("4c_6c_regression.pt")
+try:
+    ndvi_model = load_model("ndvi_best_model.keras")
+    logger.info("NDVI model loaded successfully")
+except Exception as e:
+    logger.error(f"Failed to load NDVI model: {e}")
+    ndvi_model = None
+
+try:
+    yolo_model = load_yolo_model("4c_6c_regression.pt")
+    logger.info("YOLO model loaded successfully")
+except Exception as e:
+    logger.error(f"Failed to load YOLO model: {e}")
+    yolo_model = None
 
 @app.get("/")
 async def root():
@@ -25,6 +41,9 @@ async def root():
 @app.post("/predict_ndvi/")
 async def predict_ndvi_api(file: UploadFile = File(...)):
     """Predict NDVI from RGB image"""
+    if ndvi_model is None:
+        return JSONResponse(status_code=500, content={"error": "NDVI model not loaded"})
+    
     try:
         contents = await file.read()
         img = Image.open(BytesIO(contents)).convert("RGB")
@@ -54,14 +73,18 @@ async def predict_ndvi_api(file: UploadFile = File(...)):
             headers={"Content-Disposition": "attachment; filename=ndvi_output.zip"}
         )
     except Exception as e:
+        logger.error(f"Error in predict_ndvi_api: {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"""
+    if yolo_model is None:
+        return JSONResponse(status_code=500, content={"error": "YOLO model not loaded"})
+    
     try:
         # Save uploaded file temporarily with proper extension
-        file_extension = '.tiff' if file.filename.lower().endswith(('.tif', '.tiff')) else '.tiff'
+        file_extension = '.tiff' if file.filename and file.filename.lower().endswith(('.tif', '.tiff')) else '.tiff'
         
         with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as tmp_file:
             contents = await file.read()
@@ -74,6 +97,26 @@ async def predict_yolo_api(file: UploadFile = File(...)):
             if not os.path.exists(tmp_file_path) or os.path.getsize(tmp_file_path) == 0:
                 raise ValueError("Failed to create temporary file")
             
+            logger.info(f"Processing YOLO prediction for file: {file.filename}, temp path: {tmp_file_path}")
+            
+            # Additional validation: check if file has 4 channels
+            try:
+                import tifffile
+                test_array = tifffile.imread(tmp_file_path)
+                if len(test_array.shape) == 3:
+                    if test_array.shape[0] == 4 or test_array.shape[2] == 4:
+                        channels = 4
+                    else:
+                        channels = test_array.shape[0] if test_array.shape[0] <= 4 else test_array.shape[2]
+                else:
+                    channels = 1
+                
+                if channels != 4:
+                    raise ValueError(f"YOLO model expects 4-channel images, but uploaded file has {channels} channels")
+                    
+            except Exception as validation_error:
+                logger.warning(f"Could not validate channels: {validation_error}")
+            
             # Predict using YOLO model
             results = predict_yolo(yolo_model, tmp_file_path)
             
@@ -98,6 +141,7 @@ async def predict_yolo_api(file: UploadFile = File(...)):
                         growth_stages = results.boxes.data[:, 6:].tolist()
                         results_dict["growth_stages"] = growth_stages
             
+            logger.info(f"YOLO prediction completed successfully")
             return JSONResponse(content=results_dict)
             
         finally:
@@ -106,14 +150,25 @@ async def predict_yolo_api(file: UploadFile = File(...)):
                 os.unlink(tmp_file_path)
             
     except Exception as e:
+        logger.error(f"Error in predict_yolo_api: {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"""
+    if ndvi_model is None or yolo_model is None:
+        return JSONResponse(status_code=500, content={"error": "Models not loaded properly"})
+    
     try:
         # Save uploaded file temporarily with proper extension
-        file_extension = '.tiff' if file.filename.lower().endswith(('.tif', '.tiff')) else '.jpg'
+        file_extension = '.jpg'
+        if file.filename:
+            if file.filename.lower().endswith(('.tif', '.tiff')):
+                file_extension = '.tiff'
+            elif file.filename.lower().endswith(('.png', '.PNG')):
+                file_extension = '.png'
+            elif file.filename.lower().endswith(('.jpeg', '.jpg', '.JPG', '.JPEG')):
+                file_extension = '.jpg'
         
         with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as tmp_file:
             contents = await file.read()
@@ -126,6 +181,8 @@ async def predict_pipeline_api(file: UploadFile = File(...)):
             if not os.path.exists(tmp_file_path) or os.path.getsize(tmp_file_path) == 0:
                 raise ValueError("Failed to create temporary file")
             
+            logger.info(f"Processing pipeline for file: {file.filename}, temp path: {tmp_file_path}")
+            
             # Run the full pipeline
             results = predict_pipeline(ndvi_model, yolo_model, tmp_file_path)
             
@@ -150,6 +207,7 @@ async def predict_pipeline_api(file: UploadFile = File(...)):
                         growth_stages = results.boxes.data[:, 6:].tolist()
                         results_dict["growth_stages"] = growth_stages
             
+            logger.info(f"Pipeline prediction completed successfully")
             return JSONResponse(content=results_dict)
             
         finally:
@@ -158,4 +216,5 @@ async def predict_pipeline_api(file: UploadFile = File(...)):
                 os.unlink(tmp_file_path)
             
     except Exception as e:
+        logger.error(f"Error in predict_pipeline_api: {e}")
         return JSONResponse(status_code=500, content={"error": str(e)})