apps/canada_radar_gradio.py

import gradio as gr
from PIL import Image
import numpy as np

from pydardraw.sources.msc_geomet import fetch_wms_png
from pydardraw.sources.msc_colormap import get_msc_reflectivity_colormap, get_msc_reflectivity_gradient_map
from pydardraw.colormaps.canada_msc import CANADA_MSC_REF
from pydardraw.colormaps.nws_us import NWS_REF
from pydardraw.processing.convert import image_to_dbz, dbz_to_image
from pydardraw.processing.components import estimate_pixel_resolution, label_components, component_areas_km2
from pydardraw.viz.folium_overlay import make_map_with_overlay


CANADA_BBOX = (-141.0, 41.6751, -52.6194, 83.1139)  # lon_min, lat_min, lon_max, lat_max


def fetch_and_recolor(
    size_w: int,
    size_h: int,
    layer: str,
    min_dbz: float = 20.0,
    derive_palette: bool = True,
    overlay_opacity: float = 0.8,
    fractional_dbz: bool = True,
    use_gradient_mapping: bool = True,
    grad_min: float = 5.0,
    grad_max: float = 75.0,
    grad_orientation: str = "vertical",
    grad_low_at_bottom: bool = True,
    grad_stride: int = 1,
    grad_manual_crop: bool = False,
    crop_x0: int = 0,
    crop_y0: int = 0,
    crop_x1: int = 0,
    crop_y1: int = 0,
):
    try:
        img = fetch_wms_png(layer=layer, bbox=CANADA_BBOX, size=(size_w, size_h))
    except Exception as e:
        return None, None, f"Fetch error: {e}", ""

    # Choose mapping: gradient (pixel-by-pixel) or discrete stops
    if use_gradient_mapping:
        manual_bbox = (int(crop_x0), int(crop_y0), int(crop_x1), int(crop_y1)) if grad_manual_crop else None
        grad_map = get_msc_reflectivity_gradient_map(
            layer=layer,
            min_dbz=grad_min,
            max_dbz=grad_max,
            orientation=grad_orientation,
            low_at_bottom=grad_low_at_bottom,
            sample_stride=int(grad_stride),
            manual_bbox=manual_bbox,
        )
        # Use gradient map's fractional evaluator for precise decimals
        dbz, alpha = image_to_dbz(img, grad_map, fractional=True)
    else:
        # Build exact palette (derived) if requested; fallback to built-in
        src_map = get_msc_reflectivity_colormap(layer, derive_from_legend=derive_palette)
        dbz, alpha = image_to_dbz(img, src_map, fractional=fractional_dbz)
    recolored = dbz_to_image(dbz, alpha, NWS_REF)

    res = estimate_pixel_resolution(CANADA_BBOX, (size_w, size_h))
    labels, stats = label_components(dbz, min_dbz=float(min_dbz))
    areas = component_areas_km2(labels, res)

    # Summary text
    comp_count = len(stats)
    total_km2 = sum(areas.values())
    # Compute basic stats (mean/median dBZ over valid)
    valid = ~np.isnan(dbz)
    mean_dbz = float(np.nanmean(dbz)) if np.any(valid) else float("nan")
    median_dbz = float(np.nanmedian(dbz)) if np.any(valid) else float("nan")
    summary = (
        f"Resolution: {res.km_per_px_x:.2f} km/px (x), {res.km_per_px_y:.2f} km/px (y)\n"
        f"Components (>= {min_dbz:.2f} dBZ): {comp_count}\n"
        f"Total masked area: {total_km2:.0f} km^2\n"
        f"Mean dBZ: {mean_dbz:.2f}  Median dBZ: {median_dbz:.2f}"
    )
    # Also produce a Folium overlay for the recolored image
    sw = (CANADA_BBOX[1], CANADA_BBOX[0])
    ne = (CANADA_BBOX[3], CANADA_BBOX[2])
    fmap = make_map_with_overlay(recolored, bounds=(sw, ne), opacity=overlay_opacity)
    fmap_html = fmap._repr_html_()

    return img, recolored, fmap_html, summary


with gr.Blocks(title="Canada Radar Recolor → US NWS") as demo:
    gr.Markdown("# 🇨🇦→🇺🇸 Radar Recolor (MSC → NWS)")
    gr.Markdown(
        "Fetch MSC composite, quantize to dBZ, and recolor to NWS scale.\n"
        "Note: On Hugging Face Spaces, enable ‘Allow internet’ in the Space Settings to fetch WMS/legend."
    )
    with gr.Row():
        with gr.Column(scale=1):
            layer = gr.Dropdown(
                choices=["RADAR_1KM_RDBR", "RADAR_1KM_RRAI", "RADAR_1KM_RSNO"],
                value="RADAR_1KM_RDBR",
                label="WMS Layer",
            )
            size_w = gr.Slider(256, 2048, value=1024, step=64, label="Width (px)")
            size_h = gr.Slider(256, 2048, value=768, step=64, label="Height (px)")
            min_dbz = gr.Slider(0, 75, value=20.0, step=0.25, label="Min dBZ for components")
            derive_palette = gr.Checkbox(value=True, label="Derive MSC palette from legend (exact)")
            overlay_opacity = gr.Slider(0.1, 1.0, value=0.8, step=0.05, label="Overlay opacity")
            fractional_dbz = gr.Checkbox(value=True, label="Estimate fractional dBZ (xx.xx)")
            use_gradient_mapping = gr.Checkbox(value=True, label="Use legend gradient mapping (pixel-by-pixel)")
            grad_min = gr.Number(value=5.0, label="Legend min dBZ")
            grad_max = gr.Number(value=75.0, label="Legend max dBZ")
            grad_orientation = gr.Dropdown(["vertical", "horizontal"], value="vertical", label="Legend orientation")
            grad_low_at_bottom = gr.Checkbox(value=True, label="Low at bottom / left")
            grad_stride = gr.Slider(1, 10, value=1, step=1, label="Legend sample stride (px)")
            grad_manual_crop = gr.Checkbox(value=False, label="Manual legend crop (px)")
            with gr.Row():
                crop_x0 = gr.Number(value=0, label="crop_x0 (left)")
                crop_y0 = gr.Number(value=0, label="crop_y0 (top)")
                crop_x1 = gr.Number(value=0, label="crop_x1 (right)")
                crop_y1 = gr.Number(value=0, label="crop_y1 (bottom)")
            btn = gr.Button("Fetch + Recolor", variant="primary")

            summary = gr.Textbox(label="Summary", interactive=False)

        with gr.Column(scale=2):
            orig = gr.Image(label="Original (MSC)", interactive=False)
            out = gr.Image(label="Recolored (NWS)", interactive=False)
            fmap = gr.HTML(label="Folium Map with Recolored Overlay")

    btn.click(
        fetch_and_recolor,
        inputs=[
            size_w,
            size_h,
            layer,
            min_dbz,
            derive_palette,
            overlay_opacity,
            fractional_dbz,
            use_gradient_mapping,
            grad_min,
            grad_max,
            grad_orientation,
            grad_low_at_bottom,
            grad_stride,
            grad_manual_crop,
            crop_x0,
            crop_y0,
            crop_x1,
            crop_y1,
        ],
        outputs=[orig, out, fmap, summary],
    )


if __name__ == "__main__":
    import os
    launch_kwargs = dict(server_name="0.0.0.0", debug=True)
    # Optional port override
    port = os.getenv("GRADIO_SERVER_PORT")
    if port:
        try:
            p = int(port)
            if p > 0:
                launch_kwargs["server_port"] = p
        except Exception:
            pass
    # Optional share flag
    share = os.getenv("GRADIO_SHARE", "0").lower() in {"1", "true", "yes"}
    launch_kwargs["share"] = share
    demo.launch(**launch_kwargs)