Merge branch 'pr/12' into pr/13

app.py

@@ -13,6 +13,7 @@ from rasterio.transform import from_bounds
 import tempfile
 import os
 import logging
+from resize_image import resize_image_optimized, resize_image_simple
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -39,6 +40,7 @@ except Exception as e:
 async def root():
     return {"message": "Welcome to the NDVI and YOLO prediction API!"}
 
+# Example usage in your predict_ndvi endpoint:
 @app.post("/predict_ndvi/")
 async def predict_ndvi_api(file: UploadFile = File(...)):
     """Predict NDVI from RGB image"""
@@ -46,11 +48,25 @@ async def predict_ndvi_api(file: UploadFile = File(...)):
         return JSONResponse(status_code=500, content={"error": "NDVI model not loaded"})
     
     try:
+        # Define target size (height, width)
+        target_size = (640, 640)
+        
         contents = await file.read()
         img = Image.open(BytesIO(contents)).convert("RGB")
-        norm_img = normalize_rgb(np.array(img))
+        
+        # Convert to numpy array
+        rgb_array = np.array(img)
+        
+        # Resize image to target size
+        rgb_resized = resize_image_optimized(rgb_array, target_size)
+        
+        # Normalize the resized image
+        norm_img = normalize_rgb(rgb_resized)
+        
+        # Predict NDVI
         pred_ndvi = predict_ndvi(ndvi_model, norm_img)
         
+        # Rest of the endpoint remains the same...
         # Visualization image as PNG
         vis_img_bytes = create_visualization(norm_img, pred_ndvi)
         vis_img_bytes.seek(0)
@@ -77,6 +93,7 @@ async def predict_ndvi_api(file: UploadFile = File(...)):
         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"""

ndvi_predictor.py

@@ -1,6 +1,6 @@
 # ndvi_predictor.py
 import os
-os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
+# os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
 os.environ["SM_FRAMEWORK"] = "tf.keras"
 import segmentation_models as sm
 import tensorflow as tf
@@ -11,6 +11,9 @@ import rasterio
 import matplotlib.pyplot as plt
 from PIL import Image
 import io
+from tensorflow.keras.models import model_from_json
+import traceback
+import gc
 
 # Custom loss functions and activation functions
 def balanced_mse_loss(y_true, y_pred):
@@ -27,31 +30,36 @@ def custom_mae(y_true, y_pred):
 
 def load_model(models_dir):
     """Load NDVI prediction model with custom objects"""
+    
     # Define custom objects dictionary
     custom_objects = {
         'balanced_mse_loss': balanced_mse_loss,
         'custom_mae': custom_mae
     }
     
-    # Load model architecture
-    with open(os.path.join(models_dir, "model_architecture.json"), "r") as json_file:
-        model_json = json_file.read()
-    
-    model = model_from_json(model_json, custom_objects=custom_objects)
-    
-    # Load weights
-    model.load_weights(os.path.join(models_dir, "best_model_weights.weights.h5"))
-    
-    # Compile model with custom functions
-    optimizer = tf.keras.optimizers.AdamW(learning_rate=0.0005, weight_decay=1e-4)
-    
-    model.compile(
-        optimizer=optimizer,
-        loss=balanced_mse_loss,
-        metrics=[custom_mae, 'mse']
-    )
-    
-    return model
+    try:
+        # Load model architecture
+        with open(os.path.join(models_dir, "model_architecture.json"), "r") as json_file:
+            model_json = json_file.read()
+        
+        model = model_from_json(model_json, custom_objects=custom_objects)
+        
+        # Load weights
+        model.load_weights(os.path.join(models_dir, "best_model_weights.weights.h5"))
+        
+        # Compile model with custom functions
+        optimizer = tf.keras.optimizers.AdamW(learning_rate=0.0005, weight_decay=1e-4)
+        
+        model.compile(
+            optimizer=optimizer,
+            loss=balanced_mse_loss,
+            metrics=[custom_mae, 'mse']
+        )
+        
+        return model
+    except Exception as e:
+        traceback.print_exc()
+        return None
 
 def normalize_rgb(rgb):
     """Normalize RGB image to [0, 1] range using percentile normalization"""
@@ -71,7 +79,7 @@ def normalize_rgb(rgb):
 
 def predict_ndvi(model, rgb_np):
     """
-    Predict NDVI from RGB image using tiled approach for large images
+    Faster NDVI prediction with larger tiles and more efficient processing
     
     Args:
         model: Loaded NDVI prediction model
@@ -81,46 +89,81 @@ def predict_ndvi(model, rgb_np):
         ndvi_pred: Predicted NDVI as numpy array (H, W) in range [-1, 1]
     """
     height, width = rgb_np.shape[:2]
+    
+    # Larger tiles for faster processing
     tile_size = 512
-    stride = int(tile_size * 0.7)
+    stride = int(tile_size * 0.75)  # 25% overlap
+    
+    # For smaller images, process whole image at once
+    if height <= tile_size and width <= tile_size:
+        # Pad to tile size if needed
+        pad_height = max(0, tile_size - height)
+        pad_width = max(0, tile_size - width)
+        if pad_height > 0 or pad_width > 0:
+            rgb_padded = np.pad(rgb_np, ((0, pad_height), (0, pad_width), (0, 0)), mode='reflect')
+        else:
+            rgb_padded = rgb_np
+        
+        # Single prediction
+        pred = model.predict(np.expand_dims(rgb_padded, axis=0), verbose=0, batch_size=1)[0, :, :, 0]
+        return pred[:height, :width]
     
     # Initialize output arrays
     ndvi_pred = np.zeros((height, width), dtype=np.float32)
     weight_map = np.zeros((height, width), dtype=np.float32)
     
-    # Handle small images by padding
-    if height < tile_size or width < tile_size:
-        pad_height = max(0, tile_size - height)
-        pad_width = max(0, tile_size - width)
-        rgb_padded = np.pad(rgb_np, ((0, pad_height), (0, pad_width), (0, 0)), mode='reflect')
-        height_padded, width_padded = rgb_padded.shape[0], rgb_padded.shape[1]
-    else:
-        rgb_padded = rgb_np
-        height_padded, width_padded = height, width
-    
-    # Process image tiles
-    for i in range(0, height_padded - tile_size + 1, stride):
-        for j in range(0, width_padded - tile_size + 1, stride):
+    # Pre-compute weights for efficiency
+    y, x = np.mgrid[0:tile_size, 0:tile_size]
+    base_weights = np.minimum(np.minimum(x, tile_size - x - 1), np.minimum(y, tile_size - y - 1))
+    base_weights = np.clip(base_weights, 0, 64) / 64
+    
+    # Collect all tiles first
+    tiles = []
+    positions = []
+    
+    for i in range(0, height, stride):
+        for j in range(0, width, stride):
+            # Calculate actual tile bounds
+            end_i = min(i + tile_size, height)
+            end_j = min(j + tile_size, width)
+            actual_height = end_i - i
+            actual_width = end_j - j
+            
             # Extract tile
-            tile = rgb_padded[i:i+tile_size, j:j+tile_size, :]
+            tile = rgb_np[i:end_i, j:end_j, :]
             
-            # Create distance-based weights for blending
-            y, x = np.mgrid[0:tile_size, 0:tile_size]
-            weights = np.minimum(np.minimum(x, tile_size - x - 1), np.minimum(y, tile_size - y - 1))
-            weights = np.clip(weights, 0, 50) / 50
+            # Pad if necessary
+            if actual_height < tile_size or actual_width < tile_size:
+                pad_height = tile_size - actual_height
+                pad_width = tile_size - actual_width
+                tile = np.pad(tile, ((0, pad_height), (0, pad_width), (0, 0)), mode='reflect')
             
-            # Predict NDVI for tile
-            tile_pred = model.predict(np.expand_dims(tile, axis=0), verbose=0)[0, :, :, 0]
+            tiles.append(tile)
+            positions.append((i, j, actual_height, actual_width))
+    
+    # Process all tiles in larger batches
+    batch_size = 8  # Process 8 tiles at once
+    for batch_start in range(0, len(tiles), batch_size):
+        batch_end = min(batch_start + batch_size, len(tiles))
+        batch_tiles = np.array(tiles[batch_start:batch_end])
+        
+        # Predict batch
+        batch_preds = model.predict(batch_tiles, verbose=0, batch_size=batch_size)
+        
+        # Apply predictions
+        for k in range(batch_end - batch_start):
+            pred = batch_preds[k, :, :, 0]
+            i, j, actual_height, actual_width = positions[batch_start + k]
             
-            # Determine valid region (handle edge cases)
-            valid_height = min(tile_size, height - i)
-            valid_width = min(tile_size, width - j)
+            # Use appropriate weights
+            weights = base_weights[:actual_height, :actual_width]
             
-            # Accumulate weighted predictions
-            ndvi_pred[i:i+valid_height, j:j+valid_width] += (
-                tile_pred[:valid_height, :valid_width] * weights[:valid_height, :valid_width]
-            )
-            weight_map[i:i+valid_height, j:j+valid_width] += weights[:valid_height, :valid_width]
+            # Add to output
+            ndvi_pred[i:i+actual_height, j:j+actual_width] += pred[:actual_height, :actual_width] * weights
+            weight_map[i:i+actual_height, j:j+actual_width] += weights
+        
+        # Clean up batch
+        del batch_tiles, batch_preds
     
     # Normalize by weights
     mask = weight_map > 0

resize_image.py

@@ -0,0 +1,62 @@
+from PIL import Image
+import numpy as np
+
+# Alternative: Simple resize function using PIL directly
+def resize_image_simple(image_array, target_size):
+    """
+    Simple resize function using PIL
+    
+    Args:
+        image_array: Input image array (H, W, C)
+        target_size: Tuple (height, width)
+    
+    Returns:
+        Resized image array
+    """
+    # Ensure image is in correct format
+    if image_array.max() <= 1:
+        image_array = (image_array * 255).astype(np.uint8)
+    
+    # Convert to PIL Image
+    image_pil = Image.fromarray(image_array)
+    
+    # Resize (PIL uses width, height format)
+    resized_pil = image_pil.resize((target_size[1], target_size[0]), Image.LANCZOS)
+    
+    # Convert back to numpy array and normalize back to [0, 1]
+    resized_array = np.array(resized_pil).astype(np.float32) / 255.0
+    
+    return resized_array
+
+def resize_image_optimized(image_array, target_size):
+    """
+    Resize image to target size with memory optimization
+    
+    Args:
+        image_array: Input image array (H, W, C)
+        target_size: Tuple (height, width) representing target dimensions
+    
+    Returns:
+        Resized image array
+    """
+    # Convert numpy array to PIL Image
+    if image_array.dtype != np.uint8:
+        # Convert to uint8 if not already
+        if image_array.max() <= 1:
+            image_array = (image_array * 255).astype(np.uint8)
+        else:
+            image_array = image_array.astype(np.uint8)
+    
+    image_pil = Image.fromarray(image_array)
+    
+    # Resize image (PIL uses (width, height) format)
+    resized_pil = image_pil.resize((target_size[1], target_size[0]), Image.LANCZOS)
+    
+    # Convert back to numpy array
+    result = np.array(resized_pil)
+    
+    # Clean up
+    image_pil.close()
+    resized_pil.close()
+    
+    return result