app.py

import psutil
import gradio as gr
import numpy as np
import os
import cv2
import matplotlib.pyplot as plt
from huggingface_hub import snapshot_download
import rasterio
from rasterio.enums import Resampling
from rasterio.plot import reshape_as_image
import sys

# Download the entire repository to a subdirectory
repo_id = "truthdotphd/cloud-detection"
repo_subdir = "."
repo_dir = snapshot_download(repo_id=repo_id, local_dir=repo_subdir)

# Add the repository directory to the Python path
sys.path.append(repo_dir)

# Import the necessary functions from the downloaded modules
try:
    from omnicloudmask import predict_from_array
except ImportError:
    omnicloudmask_dir = os.path.join(repo_dir, "omnicloudmask")
    if os.path.exists(omnicloudmask_dir):
        sys.path.append(omnicloudmask_dir)
        from omnicloudmask import predict_from_array
    else:
        raise ImportError("Could not find the omnicloudmask module in the downloaded repository")

def visualize_rgb(red_file, green_file, blue_file, nir_file):
    """
    Create and display an RGB visualization immediately after images are uploaded.
    """
    if not all([red_file, green_file, blue_file, nir_file]):
        return None
    
    try:
        # Get dimensions from red band to use for resampling
        with rasterio.open(red_file) as src:
            target_height = src.height
            target_width = src.width
        
        # Load bands
        blue_data = load_band(blue_file)
        green_data = load_band(green_file)
        red_data = load_band(red_file)
        
        # Compute max values for each channel for dynamic normalization
        red_max = np.max(red_data)
        green_max = np.max(green_data)
        blue_max = np.max(blue_data)
        
        # Create RGB image for visualization with dynamic normalization
        rgb_image = np.zeros((red_data.shape[0], red_data.shape[1], 3), dtype=np.float32)
        
        # Normalize each channel individually
        epsilon = 1e-10
        rgb_image[:, :, 0] = red_data / (red_max + epsilon)
        rgb_image[:, :, 1] = green_data / (green_max + epsilon)
        rgb_image[:, :, 2] = blue_data / (blue_max + epsilon)
        
        # Clip values to 0-1 range
        rgb_image = np.clip(rgb_image, 0, 1)
        
        # Apply contrast enhancement for better visualization
        p2 = np.percentile(rgb_image, 2)
        p98 = np.percentile(rgb_image, 98)
        rgb_image_enhanced = np.clip((rgb_image - p2) / (p98 - p2), 0, 1)
        
        # Convert to uint8 for display
        rgb_display = (rgb_image_enhanced * 255).astype(np.uint8)
        
        return rgb_display
    except Exception as e:
        print(f"Error generating RGB preview: {e}")
        return None


def visualize_jp2(file_path):
    """
    Visualize a single JP2 file.
    """
    with rasterio.open(file_path) as src:
        # Read the data
        data = src.read(1)
        
        # Normalize the data for visualization
        data = (data - np.min(data)) / (np.max(data) - np.min(data))
        
        # Apply a colormap for better visualization
        cmap = plt.get_cmap('viridis')
        colored_image = cmap(data)
        
        # Convert to 8-bit for display
        return (colored_image[:, :, :3] * 255).astype(np.uint8)

def load_band(file_path, resample=False, target_height=None, target_width=None):
    """
    Load a single band from a raster file with optional resampling.
    """
    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(red_file, green_file, blue_file, nir_file):
    """
    Prepare a stacked array of satellite bands for cloud mask prediction.
    """
    # Get dimensions from red band to use for resampling
    with rasterio.open(red_file) as src:
        target_height = src.height
        target_width = src.width
    
    # Load bands (resample NIR band to match 10m resolution)
    blue_data = load_band(blue_file)
    green_data = load_band(green_file)
    red_data = load_band(red_file)
    nir_data = load_band(
        nir_file, 
        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}")
    
    # Compute max values for each channel for dynamic normalization
    red_max = np.max(red_data)
    green_max = np.max(green_data)
    blue_max = np.max(blue_data)
    
    print(f"Max values - Red: {red_max}, Green: {green_max}, Blue: {blue_max}")
    
    # Create RGB image for visualization with dynamic normalization
    rgb_image = np.zeros((red_data.shape[0], red_data.shape[1], 3), dtype=np.float32)
    
    # Normalize each channel individually
    # Add a small epsilon to avoid division by zero
    epsilon = 1e-10
    rgb_image[:, :, 0] = red_data / (red_max + epsilon)
    rgb_image[:, :, 1] = green_data / (green_max + epsilon)
    rgb_image[:, :, 2] = blue_data / (blue_max + epsilon)
    
    # Clip values to 0-1 range
    rgb_image = np.clip(rgb_image, 0, 1)
    
    # Optional: Apply contrast enhancement for better visualization
    p2 = np.percentile(rgb_image, 2)
    p98 = np.percentile(rgb_image, 98)
    rgb_image_enhanced = np.clip((rgb_image - p2) / (p98 - p2), 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_enhanced


def visualize_cloud_mask(rgb_image, pred_mask):
    """
    Create a visualization of the cloud mask overlaid on the RGB image.
    """
    # Ensure pred_mask has the right dimensions
    if pred_mask.ndim > 2:
        pred_mask = np.squeeze(pred_mask)
    
    print(f"RGB image shape: {rgb_image.shape}, Pred mask shape: {pred_mask.shape}")
    
    # Ensure mask has the same spatial dimensions as the image
    if pred_mask.shape != rgb_image.shape[:2]:
        pred_mask = cv2.resize(
            pred_mask.astype(np.float32), 
            (rgb_image.shape[1], rgb_image.shape[0]),
            interpolation=cv2.INTER_NEAREST
        ).astype(np.uint8)
        print(f"Resized mask shape: {pred_mask.shape}")
    
    # Define colors for each class
    colors = {
        0: [0, 255, 0],      # Clear - Green
        1: [255, 255, 255],  # Thick Cloud - White
        2: [200, 200, 200],  # Thin Cloud - Light Gray
        3: [100, 100, 100]   # Cloud Shadow - Dark Gray
    }
    
    # Create a color-coded mask
    mask_vis = np.zeros((pred_mask.shape[0], pred_mask.shape[1], 3), dtype=np.uint8)
    for class_idx, color in colors.items():
        mask_vis[pred_mask == class_idx] = color
    
    # Create a blended visualization
    alpha = 0.5
    blended = cv2.addWeighted((rgb_image * 255).astype(np.uint8), 1-alpha, mask_vis, alpha, 0)
    
    # Get the width of the blended image for the legend
    image_width = blended.shape[1]
    
    # Create a legend with the same width as the image
    legend = np.ones((100, image_width, 3), dtype=np.uint8) * 255
    legend_text = ["Clear", "Thick Cloud", "Thin Cloud", "Cloud Shadow"]
    legend_colors = [colors[i] for i in range(4)]
    
    for i, (text, color) in enumerate(zip(legend_text, legend_colors)):
        cv2.rectangle(legend, (10, 10 + i*20), (30, 30 + i*20), color, -1)
        cv2.putText(legend, text, (40, 25 + i*20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
    
    # Combine image and legend
    final_output = np.vstack([blended, legend])
    
    return final_output
    
def process_satellite_images(red_file, green_file, blue_file, nir_file, batch_size, patch_size, patch_overlap):
    """
    Process the satellite images and detect clouds.
    """
    if not all([red_file, green_file, blue_file, nir_file]):
        return None, None, "Please upload all four channel files (Red, Green, Blue, NIR)"

    # Prepare input array and RGB image for visualization
    input_array, rgb_image = prepare_input_array(red_file, green_file, blue_file, nir_file)
    
    # Convert RGB image to format suitable for display
    rgb_display = (rgb_image * 255).astype(np.uint8)
    
    # Predict cloud mask using omnicloudmask
    pred_mask = predict_from_array(
        input_array, 
        batch_size=batch_size,
        patch_size=patch_size,
        patch_overlap=patch_overlap
    )
    
    # Calculate class distribution
    if pred_mask.ndim > 2:
        flat_mask = np.squeeze(pred_mask)
    else:
        flat_mask = pred_mask
    
    clear_pixels = np.sum(flat_mask == 0)
    thick_cloud_pixels = np.sum(flat_mask == 1)
    thin_cloud_pixels = np.sum(flat_mask == 2)
    cloud_shadow_pixels = np.sum(flat_mask == 3)
    total_pixels = flat_mask.size
    
    stats = f"""
    Cloud Mask Statistics:
    - Clear: {clear_pixels} pixels ({clear_pixels/total_pixels*100:.2f}%)
    - Thick Cloud: {thick_cloud_pixels} pixels ({thick_cloud_pixels/total_pixels*100:.2f}%)
    - Thin Cloud: {thin_cloud_pixels} pixels ({thin_cloud_pixels/total_pixels*100:.2f}%)
    - Cloud Shadow: {cloud_shadow_pixels} pixels ({cloud_shadow_pixels/total_pixels*100:.2f}%)
    - Total Cloud Cover: {(thick_cloud_pixels + thin_cloud_pixels)/total_pixels*100:.2f}%
    """
    
    # Visualize the cloud mask on the original image
    visualization = visualize_cloud_mask(rgb_image, flat_mask)
    
    return rgb_display, visualization, stats

def update_cpu():
    return f"CPU Usage: {psutil.cpu_percent()}%"

with gr.Blocks() as demo:
    cpu_text = gr.Textbox(label="CPU Usage")
    check_cpu_btn = gr.Button("Check CPU")
    
    # Attach the event handler using the click method
    check_cpu_btn.click(fn=update_cpu, inputs=None, outputs=cpu_text)
    

# Define the CPU check function
def check_cpu_usage():
    """Check and return the current CPU usage."""
    return f"CPU Usage: {psutil.cpu_percent()}%"

# Create the Gradio application with Blocks
with gr.Blocks(title="Satellite Cloud Detection") as demo:
    # Add the description
    gr.Markdown("""
    # Satellite Cloud Detection
    
    Upload separate JP2 files for Red, Green, Blue, and NIR channels to detect clouds in satellite imagery.
    
    This application uses the OmniCloudMask model to classify each pixel as:
    - Clear (0)
    - Thick Cloud (1)
    - Thin Cloud (2)
    - Cloud Shadow (3)
    
    The model works best with imagery at 10-50m resolution. For higher resolution imagery, downsampling is recommended.
    """)
    
    # Main cloud detection interface
    with gr.Row():
        with gr.Column():
            # Input components
            red_input = gr.Image(type="filepath", label="Red Channel (JP2)")
            green_input = gr.Image(type="filepath", label="Green Channel (JP2)")
            blue_input = gr.Image(type="filepath", label="Blue Channel (JP2)")
            nir_input = gr.Image(type="filepath", label="NIR Channel (JP2)")
            
            batch_size = gr.Slider(minimum=1, maximum=32, value=1, step=1, 
                                  label="Batch Size", 
                                  info="Higher values use more memory but process faster")
            patch_size = gr.Slider(minimum=500, maximum=2000, value=1000, step=100, 
                                  label="Patch Size", 
                                  info="Size of image patches for processing")
            patch_overlap = gr.Slider(minimum=100, maximum=500, value=300, step=50, 
                                     label="Patch Overlap", 
                                     info="Overlap between patches to avoid edge artifacts")
            
            process_btn = gr.Button("Process Cloud Detection")
            
        with gr.Column():
            # Output components
            rgb_output = gr.Image(label="Original RGB Image")
            cloud_output = gr.Image(label="Cloud Detection Visualization")
            stats_output = gr.Textbox(label="Statistics")
    
    # CPU usage monitoring section
    with gr.Row():
        with gr.Column():
            gr.Markdown("## System Monitoring")
            cpu_button = gr.Button("Check CPU Usage")
            cpu_output = gr.Textbox(label="CPU Usage")
    
    # Set up event handlers
    process_btn.click(
        fn=process_satellite_images,
        inputs=[red_input, green_input, blue_input, nir_input, batch_size, patch_size, patch_overlap],
        outputs=[rgb_output, cloud_output, stats_output]
    )
    
    cpu_button.click(
        fn=check_cpu_usage,
        inputs=None,
        outputs=cpu_output
    )
    
    # Add examples
    gr.Examples(
        examples=[["jp2s/B04.jp2", "jp2s/B03.jp2", "jp2s/B02.jp2", "jp2s/B8A.jp2", 1, 1000, 300]],
        inputs=[red_input, green_input, blue_input, nir_input, batch_size, patch_size, patch_overlap]
    )

# Launch the app
demo.queue(default_concurrency_limit=8).launch(debug=True)