initial commit

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+cloud_mask_visualization.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,6 @@
 ---
 license: mit
 ---
+
+
+Python 3.12

model.py

@@ -0,0 +1,199 @@
+"""
+Cloud Mask Prediction and Visualization Module
+
+This script processes Sentinel-2 satellite imagery bands to predict cloud masks
+using the omnicloudmask library. It reads blue, red, green, and near-infrared bands,
+resamples them as needed, creates a stacked array for prediction, and visualizes
+the cloud mask overlaid on the original RGB image.
+"""
+
+import rasterio
+import numpy as np
+from rasterio.enums import Resampling
+from omnicloudmask import predict_from_array
+import matplotlib.pyplot as plt
+from matplotlib.colors import ListedColormap
+import matplotlib.patches as mpatches
+
+def load_band(file_path, resample=False, target_height=None, target_width=None):
+    """
+    Load a single band from a raster file with optional resampling.
+    
+    Args:
+        file_path (str): Path to the raster file
+        resample (bool): Whether to resample the band
+        target_height (int, optional): Target height for resampling
+        target_width (int, optional): Target width for resampling
+        
+    Returns:
+        numpy.ndarray: Band data as float32 array
+    """
+    with rasterio.open(file_path) as src:
+        if resample and target_height is not None and target_width is not None:
+            band_data = src.read(
+                out_shape=(src.count, target_height, target_width),
+                resampling=Resampling.bilinear
+            )[0].astype(np.float32)
+        else:
+            band_data = src.read()[0].astype(np.float32)
+    
+    return band_data
+
+def prepare_input_array(base_path="jp2s/"):
+    """
+    Prepare a stacked array of satellite bands for cloud mask prediction.
+    
+    This function loads blue, red, green, and near-infrared bands from Sentinel-2 imagery,
+    resamples the NIR band if needed (from 20m to 10m resolution), and stacks the required
+    bands for cloud mask prediction in CHW (channel, height, width) format.
+    
+    Args:
+        base_path (str): Base directory containing the JP2 band files
+        
+    Returns:
+        tuple: (stacked_array, rgb_image)
+            - stacked_array: numpy.ndarray with bands stacked in CHW format for prediction
+            - rgb_image: numpy.ndarray with RGB bands for visualization
+    """
+    # Define paths to band files
+    band_paths = {
+        'blue': f"{base_path}B02.jp2",   # Blue band (10m)
+        'green': f"{base_path}B03.jp2",  # Green band (10m)
+        'red': f"{base_path}B04.jp2",    # Red band (10m)
+        'nir': f"{base_path}B8A.jp2"     # Near-infrared band (20m)
+    }
+
+    # Get dimensions from red band to use for resampling
+    with rasterio.open(band_paths['red']) as src:
+        target_height = src.height
+        target_width = src.width
+    
+    # Load bands (resample NIR band to match 10m resolution)
+    blue_data = load_band(band_paths['blue'])
+    green_data = load_band(band_paths['green'])
+    red_data = load_band(band_paths['red'])
+    nir_data = load_band(
+        band_paths['nir'], 
+        resample=True, 
+        target_height=target_height, 
+        target_width=target_width
+    )
+    
+    # Print band shapes for debugging
+    print(f"Band shapes - Blue: {blue_data.shape}, Green: {green_data.shape}, Red: {red_data.shape}, NIR: {nir_data.shape}")
+    
+    # Create RGB image for visualization (scale to 0-1 range)
+    # Adjust scaling factor based on your data's bit depth (e.g., 10000 for 16-bit Sentinel-2)
+    scale_factor = 10000.0  # Adjust based on your data
+    rgb_image = np.stack([
+        red_data / scale_factor, 
+        green_data / scale_factor, 
+        blue_data / scale_factor
+    ], axis=-1)
+    
+    # Clip values to 0-1 range
+    rgb_image = np.clip(rgb_image, 0, 1)
+    
+    # Stack bands in CHW format for cloud mask prediction (red, green, nir)
+    prediction_array = np.stack([red_data, green_data, nir_data], axis=0)
+    
+    return prediction_array, rgb_image
+
+def visualize_cloud_mask(rgb_image, cloud_mask, output_path="cloud_mask_visualization.png"):
+    """
+    Visualize the cloud mask overlaid on the original RGB image.
+    
+    Args:
+        rgb_image (numpy.ndarray): RGB image array (HWC format)
+        cloud_mask (numpy.ndarray): Predicted cloud mask
+        output_path (str): Path to save the visualization
+    """
+    # Fix the cloud mask shape if it has an extra dimension
+    if cloud_mask.ndim > 2:
+        # Check the shape and squeeze if needed
+        print(f"Original cloud mask shape: {cloud_mask.shape}")
+        cloud_mask = np.squeeze(cloud_mask)
+        print(f"Squeezed cloud mask shape: {cloud_mask.shape}")
+    
+    # Create figure with two subplots
+    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(18, 6))
+    
+    # Plot original RGB image
+    ax1.imshow(rgb_image)
+    ax1.set_title("Original RGB Image")
+    ax1.axis('off')
+    
+    # Define colormap for cloud mask
+    # 0=Clear, 1=Thick Cloud, 2=Thin Cloud, 3=Cloud Shadow
+    cloud_cmap = ListedColormap(['green', 'red', 'yellow', 'blue'])
+    
+    # Plot cloud mask
+    im = ax2.imshow(cloud_mask, cmap=cloud_cmap, vmin=0, vmax=3)
+    ax2.set_title("Cloud Mask")
+    ax2.axis('off')
+    
+    # Create legend patches
+    legend_patches = [
+        mpatches.Patch(color='green', label='Clear'),
+        mpatches.Patch(color='red', label='Thick Cloud'),
+        mpatches.Patch(color='yellow', label='Thin Cloud'),
+        mpatches.Patch(color='blue', label='Cloud Shadow')
+    ]
+    ax2.legend(handles=legend_patches, bbox_to_anchor=(1.05, 1), loc='upper left')
+    
+    # Plot RGB with semi-transparent cloud mask overlay
+    ax3.imshow(rgb_image)
+    
+    # Create a masked array with transparency
+    cloud_mask_rgba = np.zeros((*cloud_mask.shape, 4))
+    
+    # Set colors with alpha for each class
+    cloud_mask_rgba[cloud_mask == 0] = [0, 1, 0, 0.3]    # Clear - green with low opacity
+    cloud_mask_rgba[cloud_mask == 1] = [1, 0, 0, 0.5]    # Thick Cloud - red
+    cloud_mask_rgba[cloud_mask == 2] = [1, 1, 0, 0.5]    # Thin Cloud - yellow
+    cloud_mask_rgba[cloud_mask == 3] = [0, 0, 1, 0.5]    # Cloud Shadow - blue
+    
+    ax3.imshow(cloud_mask_rgba)
+    ax3.set_title("RGB with Cloud Mask Overlay")
+    ax3.axis('off')
+    
+    # Add legend to the overlay plot as well
+    ax3.legend(handles=legend_patches, bbox_to_anchor=(1.05, 1), loc='upper left')
+    
+    # Adjust layout and save
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=300, bbox_inches='tight')
+    plt.show()
+    
+    print(f"Visualization saved to {output_path}")
+
+def main():
+    """
+    Main function to run the cloud mask prediction and visualization workflow.
+    """
+    # Create input array from satellite bands and get RGB image for visualization
+    input_array, rgb_image = prepare_input_array()
+    
+    # Predict cloud mask using omnicloudmask
+    pred_mask = predict_from_array(input_array)
+    
+    # Print prediction results and shape
+    print("Cloud mask prediction results:")
+    print(f"Cloud mask shape: {pred_mask.shape}")
+    print(f"Unique classes in mask: {np.unique(pred_mask)}")
+    
+    # Calculate class distribution
+    if pred_mask.ndim > 2:
+        # Squeeze if needed for counting
+        flat_mask = np.squeeze(pred_mask)
+    else:
+        flat_mask = pred_mask
+        
+    print(f"Class distribution: Clear: {np.sum(flat_mask == 0)}, Thick Cloud: {np.sum(flat_mask == 1)}, "
+          f"Thin Cloud: {np.sum(flat_mask == 2)}, Cloud Shadow: {np.sum(flat_mask == 3)}")
+    
+    # Visualize the cloud mask on the original image
+    visualize_cloud_mask(rgb_image, pred_mask)
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,8 @@
+rasterio==1.3.11
+matplotlib==3.7.5
+fastai>=2.7
+timm>=0.9
+tqdm>=4.0
+rasterio>=1.3
+gdown>=5.1.0
+torch>=2.2