Update yolo_predictor.py

yolo_predictor.py

@@ -1,9 +1,14 @@
 # yolo_predictor.py
 import os
+import logging
 import rasterio
 from ultralytics import YOLO
 import tifffile
 
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
 def load_yolo_model(model_path):
     """Load YOLO model from .pt file"""
     return YOLO(model_path)
@@ -22,8 +27,11 @@ def validate_4channel_tiff(image_path):
         ValueError: If validation fails
     """
     if not os.path.exists(image_path):
+        logger.error(f"Image file does not exist: {image_path}")
         raise ValueError(f"Image file does not exist: {image_path}")
     
+    logger.info(f"Validating TIFF file: {image_path}")
+    
     try:
         # Primary validation with tifffile
         img_array = tifffile.imread(image_path)
@@ -42,25 +50,30 @@ def validate_4channel_tiff(image_path):
                 channels = min(img_array.shape[0], img_array.shape[2])
                 height, width = img_array.shape[0], img_array.shape[1]
         else:
+            logger.error(f"Invalid image shape: {img_array.shape}. Expected 3D array with 4 channels.")
             raise ValueError(f"Invalid image shape: {img_array.shape}. Expected 3D array with 4 channels.")
             
         if channels != 4:
+            logger.error(f"YOLO model expects 4-channel images, but got {channels} channels")
             raise ValueError(f"YOLO model expects 4-channel images, but got {channels} channels")
             
-        print(f"Validation successful: {channels} channels, {height}x{width}, dtype: {img_array.dtype}")
+        logger.info(f"Validation successful: {channels} channels, {height}x{width}, dtype: {img_array.dtype}")
         return True
         
     except Exception as e:
+        logger.warning(f"Tifffile validation failed: {str(e)}, trying rasterio fallback")
         # Fallback validation with rasterio
         try:
             with rasterio.open(image_path) as src:
                 if src.count != 4:
+                    logger.error(f"YOLO model expects 4-channel images, but got {src.count} channels")
                     raise ValueError(f"YOLO model expects 4-channel images, but got {src.count} channels")
                     
-                print(f"Validation successful (rasterio): {src.count} channels, {src.width}x{src.height}, dtype: {src.dtypes[0]}")
+                logger.info(f"Validation successful (rasterio): {src.count} channels, {src.width}x{src.height}, dtype: {src.dtypes[0]}")
                 return True
                 
         except Exception as e2:
+            logger.error(f"Could not validate TIFF file. Tifffile error: {str(e)}, Rasterio error: {str(e2)}")
             raise ValueError(f"Could not validate TIFF file. Errors: tifffile={str(e)}, rasterio={str(e2)}")
 
 def predict_yolo(yolo_model, image_path, conf=0.001):
@@ -75,12 +88,16 @@ def predict_yolo(yolo_model, image_path, conf=0.001):
     Returns:
         results: YOLO results object
     """
+    logger.info(f"Starting YOLO prediction on: {image_path} with confidence: {conf}")
+    
     # Validate input file
     validate_4channel_tiff(image_path)
     
+    logger.info("Running YOLO model inference...")
     # Run YOLO prediction directly on the input file
     results = yolo_model([image_path], conf=conf)
     
+    logger.info(f"YOLO prediction completed. Results type: {type(results[0])}")
     return results[0]  # Return first result
 
 def predict_pipeline(ndvi_model, yolo_model, image_path, conf=0.001):
@@ -96,5 +113,8 @@ def predict_pipeline(ndvi_model, yolo_model, image_path, conf=0.001):
     Returns:
         results: YOLO results object
     """
+    logger.info(f"Starting prediction pipeline for: {image_path}")
     # Simply validate and run prediction on the uploaded file
-    return predict_yolo(yolo_model, image_path, conf=conf)
+    result = predict_yolo(yolo_model, image_path, conf=conf)
+    logger.info("Prediction pipeline completed successfully")
+    return result