Performance enhancement of image statistics. Not calculating in RGB mode, as not needed. Batching for cal_on_everything() agains OOM.

Author: ChristopherKotthoff

detectree2/preprocessing/tiling.py

@@ -673,22 +673,79 @@ def calculate_image_statistics(file_path,
         def calc_on_everything():
             logger.info("Processing entire image...")
             band_stats = []
+            
+            # Define chunk size for reading (e.g. 2048 rows)
+            chunk_height = 2048
+            
             for band_idx in range(1, src.count + 1):
                 if band_idx - 1 in ignore_bands_indices:
                     continue
-                band = src.read(band_idx).astype(float)
-                # Mask out bad values
-                mask = (np.isnan(band) | np.isin(band, values_to_ignore))
-                valid_data = band[~mask]
-
-                if valid_data.size > 0:
-                    min_val, max_val = np.percentile(valid_data, [1, 99])
+                
+                # Accumulators for exact stats
+                total_count = 0
+                total_sum = 0.0
+                total_sum_sq = 0.0
+                global_min = float('inf')
+                global_max = float('-inf')
+                
+                # Buffer for percentiles
+                percentile_buffer = []
+                buffer_size = 0
+                MAX_BUFFER = 5_000_000 # 5 million pixels ~ 40MB
+                
+                for row_off in tqdm(range(0, height, chunk_height), desc=f"Calculating stats for band {band_idx}", leave=False):
+                    h = min(chunk_height, height - row_off)
+                    window = rasterio.windows.Window(0, row_off, width, h)
+                    
+                    band_chunk = src.read(band_idx, window=window).astype(float)
+                    
+                    # Mask out bad values
+                    mask = (np.isnan(band_chunk) | np.isin(band_chunk, values_to_ignore))
+                    valid_chunk = band_chunk[~mask]
+                    
+                    if valid_chunk.size > 0:
+                        # Update exact stats
+                        c_min = np.min(valid_chunk)
+                        c_max = np.max(valid_chunk)
+                        c_sum = np.sum(valid_chunk)
+                        c_sum_sq = np.sum(valid_chunk ** 2)
+                        c_count = valid_chunk.size
+                        
+                        if c_min < global_min: global_min = c_min
+                        if c_max > global_max: global_max = c_max
+                        total_sum += c_sum
+                        total_sum_sq += c_sum_sq
+                        total_count += c_count
+                        
+                        # Update percentile buffer
+                        percentile_buffer.append(valid_chunk)
+                        buffer_size += c_count
+                        
+                        if buffer_size > MAX_BUFFER:
+                            merged = np.concatenate(percentile_buffer)
+                            # Downsample to keep memory usage low
+                            merged = merged[::2] 
+                            percentile_buffer = [merged]
+                            buffer_size = merged.size
+
+                if total_count > 0:
+                    # Finalize percentiles
+                    if percentile_buffer:
+                        final_buffer = np.concatenate(percentile_buffer)
+                        min_val, max_val = np.percentile(final_buffer, [1, 99])
+                    else:
+                        min_val, max_val = global_min, global_max
+
+                    mean_val = total_sum / total_count
+                    # Variance = (SumSq / N) - Mean^2
+                    variance = (total_sum_sq / total_count) - (mean_val ** 2)
+                    std_dev = np.sqrt(max(0, variance))
 
                     stats = {
-                        "mean": float(np.mean(valid_data)),
+                        "mean": float(mean_val),
                         "min": float(min_val),
                         "max": float(max_val),
-                        "std_dev": float(np.std(valid_data)),
+                        "std_dev": float(std_dev),
                     }
                 else:
                     stats = {
@@ -835,11 +892,12 @@ def tile_data(
 
     tile_coordinates = _calculate_tile_placements(img_path, buffer, tile_width, tile_height, crowns, tile_placement,
                                                   overlapping_tiles)
+    
     image_statistics = calculate_image_statistics(img_path,
                                                   values_to_ignore=additional_nodata,
                                                   mode=mode,
-                                                  ignore_bands_indices=ignore_bands_indices)
-
+                                                  ignore_bands_indices=ignore_bands_indices) if mode == "ms" else None # Only needed for multispectral data
+    
     tile_args = [
         (img_path, out_dir, buffer, tile_width, tile_height, dtype_bool, minx, miny, crs, tilename, crowns, threshold,
          nan_threshold, mode, class_column, mask_gdf, additional_nodata, image_statistics, ignore_bands_indices,

docs/source/tutorials/02_data_preparation.rst

@@ -88,7 +88,7 @@ The ``tile_data`` function exposes many parameters to control how tiles are crea
 
 - ``nan_threshold``: The maximum proportion of a tile that can be NaN (or other no-data values) before it is discarded.
 
-- ``use_convex_mask``: When ``True``, this creates a tight "wrapper" polygon (a convex hull) around all the training crowns within a tile. Any pixels outside this wrapper are masked out. This is a way to reduce noise by forcing the model to ignore parts of the tile that are far from any labeled object.
+- ``use_convex_mask``: When ``True``, this creates a tight "wrapper" polygon (a convex hull) around all the training crowns within a tile. Any pixels outside this wrapper are masked out. This is a way to reduce noise by forcing the model to ignore parts of the tile that are far from any labeled object. Note that the masked out area counts towards the ``nan_threshold``, so you may need to increase ``nan_threshold`` when using this option.
 
 - ``enhance_rgb_contrast``: When ``True`` (for RGB images only), this applies a percentile contrast stretch. It calculates the 0.2 and 99.8 percentile pixel values and rescales the image to a 1-255 range. This is effective for normalizing hazy, dark, or washed-out imagery. It allows the model to more easily differentiate between tree crowns. 0 is reserved for masked-out areas.