load app & the rest

.streamlit/config.toml

@@ -0,0 +1,10 @@
+[server]
+
+# Max size, in megabytes, for files uploaded with the file_uploader.
+# Default: 200
+maxUploadSize = 1024
+
+# Max size, in megabytes, of messages that can be sent via the WebSocket
+# connection.
+# Default: 200
+maxMessageSize = 1024

README.md

@@ -1,14 +1,40 @@
 ---
-title: Gaia Growseg Demo
-emoji: 🌖
-colorFrom: pink
-colorTo: pink
+title: GRowSeg demo
+emoji: 🍇
+colorFrom: indigo
+colorTo: gray
 sdk: streamlit
 sdk_version: 1.43.2
+suggested_hardware: t4-small
 app_file: app.py
+short_description: Vineyard row segmentation from UAV imagery
 pinned: false
 license: mit
-short_description: Vineyard row segmentation from UAV imagery
+models:
+- links-ads/gaia-growseg
+datasets:
+- links-ads/gaia-vineyard-uav-dataset
+preload_from_hub:
+- links-ads/gaia-growseg
+tags:
+- agriculture
+- viticulture
+- remote-sensing
+- image-segmentation
+- segmentation
+- semantic-segmentation
+- grapevines
+- grapes
+- vineyard
+- uav
+- drone
+- aerial-imagery
+- aerial-photography
+- aerial-photos
+- aerial-images
+- crop
+- field
+- links-ads
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+GRowSeg, a deep learning model for vineyard row segmentation from UAV imagery

app.py

@@ -0,0 +1,173 @@
+import os
+from datetime import datetime
+from pathlib import Path
+import torch
+import folium
+import streamlit as st
+from loguru import logger
+from tqdm import tqdm
+from streamlit_folium import st_folium
+from transformers import SegformerForSemanticSegmentation
+from lib.folium import (
+    get_clean_rendering_container,
+    create_map,
+    process_raster_and_overlays,
+)
+import streamlit.components.v1 as components
+
+# Page configs
+st.set_page_config(page_title="GrowSeg Demo", page_icon="🍇", layout="wide")
+
+# BUGFIX (https://discuss.streamlit.io/t/message-error-about-torch/90886/6)
+torch.classes.__path__ = []
+
+# Interoperability with tqdm (https://loguru.readthedocs.io/en/stable/resources/recipes.html#interoperability-with-tqdm-iterations)
+logger.remove()
+logger.add(lambda msg: tqdm.write(msg, end=""), colorize=True, format="<green>{message}</green>")
+
+@st.cache_resource
+def load_model(hf_path='links-ads/gaia-growseg'):
+    # logger.info(f'Loading GAIA GRowSeg on {device}...')
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = SegformerForSemanticSegmentation.from_pretrained(
+        hf_path,
+        num_labels=1,
+        num_channels=3,
+        id2label={1: 'vine'},
+        label2id={'vine': 1},
+        token=os.getenv('hf_read_access_token')
+    )
+    return model.to(device).eval()
+
+# Load GAIA GRowSeg model
+model = load_model()
+
+
+def change_key():
+    st.session_state["key_map"] = str(datetime.now())
+
+
+# Create selection menu
+container_predictions = st.container(border=True)
+with container_predictions:
+    col1, col2 = st.columns([0.3, 0.7])
+    with col1:
+        # raster_path = st.text_input(
+        #     "Enter the path to your local file: ",
+        #     key="raster_path_block",
+        # )
+        # raster_path = st.file_uploader(
+        #     "Upload a raster file",
+        #     type=["tif", "tiff"],
+        #     key="raster_path_block",
+        # )
+        precomputed_map_path = None
+        raster_path = None
+        raster_selection = st.selectbox(
+            "Select an example or your own raster...",
+            options=[
+                "Italy",
+                "Portugal",
+                "Spain",
+                "Upload file...",
+            ],
+            key="raster_selection_block",
+            index=None,
+            placeholder="Choose an example or upload your own raster",
+        )
+        if raster_selection == "Italy":
+            st.markdown("At this stage, only Portugal is available due to the WebSocket payload limit.")
+            # TODO GEOSERVER
+            #precomputed_map_path = "data/italy_2022-06-13_cropped.html"
+        elif raster_selection == "Portugal":
+            precomputed_map_path = "data/portugal_2023-08-01.html"
+        elif raster_selection == "Spain":
+            st.markdown("At this stage, only Portugal is available due to the WebSocket payload limit.")
+            #precomputed_map_path = "data/spain_2022-07-29_cropped.html"
+        elif raster_selection == "Upload file...":
+            uploaded_file = st.file_uploader(
+                "Upload a raster file",
+                type=["tif"],
+                key="uploaded_file_block",
+            )
+            if uploaded_file is not None:
+                fn = Path(uploaded_file.name).name
+                print(fn)
+                raster_path = os.path.join("temp", fn)
+                with open(raster_path, "wb") as f:
+                    f.write(uploaded_file.getbuffer())
+
+    is_raster_path_selected = raster_path is not None
+    is_precomputed_map_selected = precomputed_map_path is not None
+
+    with col2:
+        with st.container():
+            st.write("######")
+            with st.expander("More info on the model"):
+                st.write("""
+                    Under the hood, this model is a SegFormer-b5, trained on
+                    UAV-acquired vineyard orthoimages and their ground-truth
+                    delineation masks. Paper will be available soon. Stay tuned!
+                """)
+
+    if not is_precomputed_map_selected and is_raster_path_selected:
+        progress_bar = st.progress(0, text="Begin processing...")
+        # Process raster and get overlays
+        overlays = process_raster_and_overlays(raster_path, model, _progress_bar=progress_bar)
+        #progress_bar.empty()
+
+
+#container = get_clean_rendering_container(raster_path)
+container = st.empty()
+
+# draw map
+interactive_map = create_map()
+
+if is_raster_path_selected:
+    # Add overlays to map
+    for overlay in overlays:
+        overlay.add_to(interactive_map)
+
+with container.form(key="form1"):
+
+    if is_precomputed_map_selected:
+        # Load precomputed map
+        # interactive_map = folium.Map(location=[35, -10], zoom_start=6)
+        # folium.IFrame(
+        #     precomputed_map_path,
+        #     width=1000,
+        #     height=500,
+        # ).add_to(interactive_map)
+
+        with open(precomputed_map_path, 'r') as f:
+            html_content = f.read()
+        interactive_map = components.html(html_content, height=500)
+
+
+    else:
+
+        if is_raster_path_selected:
+            # Center map on overlays
+            bounds = overlays[0].get_bounds()
+            interactive_map.fit_bounds(bounds)
+        else:
+            # Center map on Europe
+            interactive_map.fit_bounds([[35, -10], [60, 40]])
+
+        # Add Layer Control (first remove existing one)
+        for key, child in list(interactive_map._children.items()):
+            if isinstance(child, folium.map.LayerControl):
+                del interactive_map._children[key]
+        folium.LayerControl().add_to(interactive_map)
+
+        # Folium Map component
+        output_map = st_folium(
+            interactive_map,
+            width=None,
+            height=500,
+            returned_objects=["all_drawings"],
+            key=st.session_state.get("key_map", "key_map"), # This is a workaround to force the map to recenter
+        )
+
+    # Recenter map
+    submit = st.form_submit_button("Recenter map")

lib/folium.py

@@ -0,0 +1,246 @@
+import branca
+import folium
+import geopandas as gpd
+import streamlit as st
+from loguru import logger
+import rioxarray as rxr
+import numpy as np
+import xarray as xr
+import torch
+from .utils import compute_mask, compute_vndvi, compute_vdi
+import os
+
+
+
+@st.cache_resource
+def create_map(location=[41.9099533, 12.3711879], zoom_start=5, crs=3857, max_zoom=23):
+    """Create a folium map with OpenStreetMap tiles and optional Esri.WorldImagery basemap."""
+    if isinstance(crs, int):
+        crs = f"EPSG{crs}"
+    assert crs in ["EPSG3857"], f"Only EPSG:3857 supported for now. Got {crs}."
+    
+    m = folium.Map(
+        location=location,
+        zoom_start=zoom_start,
+        crs=crs,
+        max_zoom=max_zoom,
+        tiles="OpenStreetMap",  # Esri.WorldImagery
+        attributionControl=False,
+        prefer_canvas=True,
+    )
+
+    # Add Esri.WorldImagery as optional basemap (radio button)
+    folium.TileLayer(
+        tiles="Esri.WorldImagery",
+        show=False,
+        overlay=False,
+        control=True,
+    ).add_to(m)
+
+    return m
+
+
+
+def get_clean_rendering_container(app_state: str):
+    """Makes sure we can render from a clean slate on state changes."""
+    slot_in_use = st.session_state.slot_in_use = st.session_state.get(
+        "slot_in_use", "a"
+    )
+    if app_state != st.session_state.get("previous_state", app_state):
+        if slot_in_use == "a":
+            slot_in_use = st.session_state.slot_in_use = "b"
+        else:
+            slot_in_use = st.session_state.slot_in_use = "a"
+
+    st.session_state.previous_state = app_state
+
+    slot = {
+        "a": st.empty(),
+        "b": st.empty(),
+    }[slot_in_use]
+
+    return slot.container()
+
+
+
+def create_image_overlay(raster_path_or_array, name="Raster", opacity=1.0, to_crs=4326, show=True):
+    """ Create a folium image overlay from a raster filepath or xarray.DataArray. """
+    if isinstance(raster_path_or_array, str):
+        # Open the raster and its metadata
+        r = rxr.open_rasterio(raster_path_or_array)
+    else:
+        r = raster_path_or_array
+    nodata = r.rio.nodata or 0
+    if r.rio.crs.to_epsg() != to_crs:
+        r = r.rio.reproject(to_crs, nodata=nodata) # nodata default: 255
+    r = r.transpose("y", "x", "band")
+    bounds = r.rio.bounds()   # (left, bottom, right, top)
+
+    # Create a folium image overlay
+    overlay = folium.raster_layers.ImageOverlay(
+        image=r.to_numpy(),
+        name=name,
+        bounds=[[bounds[1], bounds[0]], [bounds[3], bounds[2]]],    # format for folium: ((bottom,left),(top,right))
+        opacity=opacity,
+        interactive=True,
+        cross_origin=False,
+        zindex=1,
+        show=show,
+    )
+
+    return overlay
+
+
+
+@st.cache_resource
+def process_raster_and_overlays(
+        raster_path: str,
+        _model: torch.nn.Module,
+        patch_size=512,
+        stride=256,
+        scaling_factor=None,
+        rotate=False,
+        batch_size=16,
+        window_size=360,
+        dilate_rows=False,
+        _progress_bar=None,
+        ):
+    
+    # Define paths for mask, vNDVI, and VDI
+    mask_path = raster_path.replace('.tif', '_mask.tif')
+    vndvi_rows_path = raster_path.replace('.tif', '_vndvi_rows.tif')
+    vndvi_interrows_path = raster_path.replace('.tif', '_vndvi_interrows.tif')
+    vdi_path = raster_path.replace('.tif', '_vdi.tif')
+    if os.path.exists(mask_path):
+        assert os.path.exists(vndvi_rows_path)
+        assert os.path.exists(vndvi_interrows_path)
+        assert os.path.exists(vdi_path)
+        logger.info(f"Found mask at {mask_path!r}, vNDVI at {vndvi_rows_path!r} and {vndvi_interrows_path!r}, and VDI at {vdi_path!r}. Loading...")
+
+    # Read raster
+    logger.info(f'Reading raster image {raster_path!r}...')
+    if _progress_bar: _progress_bar.progress(0, text=f'Reading raster image {raster_path!r}...')
+    raster = rxr.open_rasterio(raster_path)
+
+    # Compute mask
+    logger.info('### Computing mask...')
+    if _progress_bar: _progress_bar.progress(10, text='### Computing mask...')
+
+    
+    if os.path.exists(mask_path):
+        mask_raster = rxr.open_rasterio(mask_path)    # mask is RGBA (red for vine)
+    else:
+        mask = compute_mask(
+            raster.to_numpy(),
+            _model,
+            patch_size=patch_size,
+            stride=stride,
+            scaling_factor=scaling_factor,
+            rotate=rotate,
+            batch_size=batch_size
+        )   # mask is a HxW uint8 array in with 0=background, 255=vine, 1=nodata
+
+        # Convert mask from grayscale to RGBA, with red pixels for vine
+        alpha = ((mask != 1)*255).astype(np.uint8)
+        mask_colored = np.stack([mask, np.zeros_like(mask), np.zeros_like(mask), alpha], axis=0)  # now, mask is a 4xHxW uint8 array in with 0=background, 255=vine
+
+        # Georef mask like raster
+        logger.info('Georeferencing mask...')
+        if _progress_bar: _progress_bar.progress(30, text='Georeferencing mask...')
+        mask_raster = xr.DataArray(
+            mask_colored,
+            dims=('band', 'y', 'x'),
+            coords={'x': raster.x, 'y': raster.y, 'band': raster.band}
+            )
+        mask_raster.rio.write_crs(raster.rio.crs, inplace=True)  # Copy CRS
+        mask_raster.rio.write_transform(raster.rio.transform(), inplace=True)  # Copy affine transform
+
+    # Compute vNDVI
+    logger.info('### Computing vNDVI...')
+    if _progress_bar: _progress_bar.progress(35, text='### Computing vNDVI...')
+
+    if os.path.exists(vndvi_rows_path) and os.path.exists(vndvi_interrows_path):
+        vndvi_rows_raster = rxr.open_rasterio(vndvi_rows_path)    # vNDVI is RGBA
+        vndvi_interrows_raster = rxr.open_rasterio(vndvi_interrows_path)  # vNDVI is RGBA
+    else:
+        vndvi_rows, vndvi_interrows = compute_vndvi(
+            raster.to_numpy(),
+            mask,
+            dilate_rows=dilate_rows,
+            window_size=window_size
+            )    # vNDVI is RGBA
+
+        # Georef vNDVI like raster
+        logger.info('Georeferencing vNDVI...')
+        if _progress_bar: _progress_bar.progress(55, text='Georeferencing vNDVI...')
+        vndvi_rows_raster = xr.DataArray(
+            vndvi_rows.transpose(2, 0, 1),
+            dims=('band', 'y', 'x'),
+            coords={'x': raster.x, 'y': raster.y, 'band': raster.band}
+            )
+        vndvi_rows_raster.rio.write_crs(raster.rio.crs, inplace=True)
+        vndvi_rows_raster.rio.write_transform(raster.rio.transform(), inplace=True)
+
+        vndvi_interrows_raster = xr.DataArray(
+            vndvi_interrows.transpose(2, 0, 1),
+            dims=('band', 'y', 'x'),
+            coords={'x': raster.x, 'y': raster.y, 'band': raster.band}
+            )
+        vndvi_interrows_raster.rio.write_crs(raster.rio.crs, inplace=True)
+        vndvi_interrows_raster.rio.write_transform(raster.rio.transform(), inplace=True)
+
+    # Compute VDI
+    logger.info('### Computing VDI...')
+    if _progress_bar: _progress_bar.progress(60, text='### Computing VDI...')
+
+    if os.path.exists(vdi_path):
+        vdi_raster = rxr.open_rasterio(vdi_path)    # VDI is RGBA
+    else:
+        vdi = compute_vdi(
+            raster.to_numpy(),
+            mask,
+            window_size=window_size
+            )    # VDI is RGBA
+
+        # Georef VDI like raster
+        logger.info('Georeferencing VDI...')
+        if _progress_bar: _progress_bar.progress(80, text='Georeferencing VDI...')
+        vdi_raster = xr.DataArray(
+            vdi.transpose(2, 0, 1),
+            dims=('band', 'y', 'x'),
+            coords={'x': raster.x, 'y': raster.y, 'band': raster.band}
+            )
+        vdi_raster.rio.write_crs(raster.rio.crs, inplace=True)
+        vdi_raster.rio.write_transform(raster.rio.transform(), inplace=True)
+
+    # Reproject all rasters to EPSG:4326
+    if raster.rio.crs.to_epsg() != 4326:
+        logger.info(f"Reprojecting rasters to EPSG:4326 with NODATA value 0...")
+        if _progress_bar: _progress_bar.progress(82, text=f"Reprojecting rasters to EPSG:4326 with NODATA value 0...")
+        raster = raster.rio.reproject("EPSG:4326", nodata=0)    # nodata default: 255
+        mask_raster = mask_raster.rio.reproject("EPSG:4326", nodata=0)
+        vndvi_rows_raster = vndvi_rows_raster.rio.reproject("EPSG:4326", nodata=0)
+        vndvi_interrows_raster = vndvi_interrows_raster.rio.reproject("EPSG:4326", nodata=0)
+        vdi_raster = vdi_raster.rio.reproject("EPSG:4326", nodata=0)
+
+    # Create overlays
+    logger.info(f'Creating RGB raster overlay...')
+    if _progress_bar: _progress_bar.progress(85, text='Creating overlays: drone image...')
+    raster_overlay = create_image_overlay(raster, name="Orthoimage", opacity=1.0, show=True)
+    logger.info(f'Creating mask overlay...')
+    if _progress_bar: _progress_bar.progress(88, text='Creating overlays: mask...')
+    mask_overlay = create_image_overlay(mask_raster, name="Mask", opacity=1.0, show=False)
+    logger.info(f'Creating vNDVI rows overlay...')
+    if _progress_bar: _progress_bar.progress(91, text='Creating overlays: vNDVI (rows)...')
+    vndvi_rows_overlay = create_image_overlay(vndvi_rows_raster, name="vNDVI Rows", opacity=1.0, show=False)
+    logger.info(f'Creating vNDVI interrows overlay...')
+    if _progress_bar: _progress_bar.progress(94, text='Creating overlays: vNDVI (interrows)...')
+    vndvi_interrows_overlay = create_image_overlay(vndvi_interrows_raster, name="vNDVI Interrows", opacity=1.0, show=False)
+    logger.info(f'Creating VDI overlay...')
+    if _progress_bar: _progress_bar.progress(97, text='Creating overlays: VDI...')
+    vdi_overlay = create_image_overlay(vdi_raster, name="VDI", opacity=1.0, show=False)
+
+    logger.info('Done!')
+    if _progress_bar: _progress_bar.progress(100, text='Done!')
+
+    return [raster_overlay, mask_overlay, vndvi_rows_overlay, vndvi_interrows_overlay, vdi_overlay]

lib/utils.py

@@ -0,0 +1,587 @@
+import numpy as np
+import torch
+import rasterio
+import xarray as xr
+import rioxarray as rxr
+import cv2
+from transformers import SegformerForSemanticSegmentation
+from tqdm import tqdm
+from scipy.ndimage import grey_dilation
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from .viz_utils import alpha_composite
+from loguru import logger
+
+
+
+def resize(img, shape=None, scaling_factor=1., order='CHW'):
+    """Resize an image by a given scaling factor"""
+    assert order in ['HWC', 'CHW'], f"Got unknown order '{order}', expected one of ['HWC','CHW']"
+    assert shape is None or scaling_factor == 1., "Got both shape and scaling_factor. Please provide only one of them"
+
+    # resize image
+    if order == 'CHW':
+        img = np.moveaxis(img, 0, -1)   # CHW -> HWC
+
+    if shape is not None:
+        img = cv2.resize(img, shape[::-1], interpolation=cv2.INTER_LINEAR)
+    else:
+        img = cv2.resize(img, None, fx=scaling_factor, fy=scaling_factor, interpolation=cv2.INTER_LINEAR)
+    
+    # NB: cv2.resize returns a HW image if the input image is HW1: restore the C dimension
+    if len(img.shape) == 2:
+        img = img[..., None]
+
+    if order == 'CHW':
+        img = np.moveaxis(img, -1, 0)   # HWC -> CHW
+
+    return img
+
+
+def minimum_needed_padding(img_size, patch_size: int, stride: int):
+    """
+    Compute the minimum padding needed to make an image divisible by a patch size with a given stride.
+    Args:
+        image_shape (tuple): the shape (H,W) of the image tensor
+        patch_size (int): the size of the patches to extract
+        stride (int): the stride to use when extracting patches
+    Returns:
+        tuple: the padding needed to make the image tensor divisible by the patch size with the given stride
+    """
+    
+    img_size = np.array(img_size)
+    pad = np.where(
+        img_size <= patch_size,
+        (patch_size - img_size) % patch_size,   # the % patch_size is to handle the case img_size = (0,0)
+        (stride - (img_size - patch_size)) % stride
+        )
+    pad_t, pad_l = pad // 2
+    pad_b, pad_r = pad[0] - pad_t, pad[1] - pad_l
+
+    return pad_t, pad_b, pad_l, pad_r
+
+
+def pad(img, pad, order='CHW'):
+    """Pad an image by the given pad values, in the format (pad_t, pad_b, pad_l, pad_r)"""
+    assert order in ['HWC', 'CHW'], f"Got unknown order '{order}', expected one of ['HWC','CHW']"
+
+    pad_t, pad_b, pad_l, pad_r = pad
+
+    # pad image
+    if order == 'HWC':
+        padded_img = np.pad(img, ((pad_t,pad_b), (pad_l,pad_r), (0,0)), mode='constant', constant_values=0) # can also try mode='reflect'
+    else:
+        padded_img = np.pad(img, ((0,0), (pad_t,pad_b), (pad_l,pad_r)), mode='constant', constant_values=0) # can also try mode='reflect'
+
+    if isinstance(img, torch.Tensor):
+        padded_img = torch.tensor(padded_img)
+
+    return padded_img
+
+
+def extract_patches(img, patch_size=512, stride=256, order='CHW', only_return_idx=True):
+    """Extract patches from an image, in the format (h_start, h_end, w_start, w_end)"""
+    assert order in ['HWC', 'CHW'], f"Got unknown order '{order}', expected one of ['HWC','CHW']"
+
+    if order == 'HWC':
+        H, W = img.shape[:2]
+    else:
+        H, W = img.shape[1:]
+
+    # compute the number of patches
+    n_patches = ((H - patch_size) // stride + 1) * ((W - patch_size) // stride + 1)
+
+    # extract patches
+    patches = []
+    patches_idx = []
+    for i in range(0, H-patch_size+1, stride):
+        for j in range(0, W-patch_size+1, stride):
+            
+            patches_idx.append((i, i+patch_size, j, j+patch_size))
+
+            if not only_return_idx:
+                if order == 'HWC':
+                    patch = img[i:i+patch_size, j:j+patch_size, :]
+                else:
+                    patch = img[:, i:i+patch_size, j:j+patch_size]
+                patches.append(patch)
+            
+    if only_return_idx:
+        return patches_idx
+    return patches, patches_idx
+
+
+def segment_batch(batch, model):
+
+    # perform prediction
+    with torch.no_grad():
+        out = model(batch)  # (n_patches, 1, H, W) logits
+        if isinstance(model, SegformerForSemanticSegmentation):
+            out = upsample(out.logits, size=batch.shape[-2:])
+
+        # apply sigmoid
+        out = torch.sigmoid(out)    # logits -> confidence scores
+    
+    return out
+
+
+def upsample(x, size):
+    """Upsample a 3D/4D/5D tensor"""
+    return torch.nn.functional.interpolate(x, size=size, mode='bilinear', align_corners=False)
+
+
+def merge_patches(patches, patches_idx, rotate=False, canvas_shape=None, order='CHW'): # TODO
+    """Merge patches into a single image"""
+    assert order in ['HWC', 'CHW'], f"Got unknown order '{order}', expected one of ['HWC','CHW']"
+    if rotate:
+        axes_to_rotate = (0,1) if order == 'HWC' else (1,2)
+        patches = [np.rot90(p, -i, axes=axes_to_rotate) for i,p in enumerate(patches)]
+    else:
+        assert len(patches) == len(patches_idx), f"Got {len(patches)} patches and {len(patches_idx)} indexes"
+
+    # if canvas_shape is None, infer it from patches_idx
+    if canvas_shape is None:
+        patches_idx_zipped = list(zip(*patches_idx))
+        canvas_H = max(patches_idx_zipped[1])
+        canvas_W = max(patches_idx_zipped[3])
+    else:
+        canvas_H, canvas_W = canvas_shape
+    
+    # initialize canvas
+    dtype = patches[0].dtype
+    if order == 'HWC':
+        canvas_C = patches[0].shape[-1]
+        canvas = np.zeros((canvas_H, canvas_W, canvas_C), dtype=dtype)  # HWC
+        n_overlapping_patches = np.zeros((canvas_H, canvas_W, 1))
+    else:
+        canvas_C = patches[0].shape[0]
+        canvas = np.zeros((canvas_C, canvas_H, canvas_W, ), dtype=dtype)  # CHW
+        n_overlapping_patches = np.zeros((1, canvas_H, canvas_W))
+    
+    # merge patches     
+    for p, (t,b,l,r) in zip(patches, patches_idx):
+        if order == 'HWC':
+            canvas[t:b, l:r, :] += p
+            n_overlapping_patches[t:b, l:r, 0] += 1
+        else:
+            canvas[:, t:b, l:r] += p
+            n_overlapping_patches[0, t:b, l:r] += 1
+        
+    
+    # compute average
+    canvas = np.divide(canvas, n_overlapping_patches, where=(n_overlapping_patches != 0))
+    
+    return canvas
+
+
+def segment(img, model, patch_size=512, stride=256, scaling_factor=1., rotate=False, device=None, batch_size=16, verbose=False):
+    """Segment an RGB image by using a segmentation model. Returns a probability
+    map (and performance metrics, if requested)"""
+    
+    # some checks
+    assert isinstance(img, np.ndarray), f"Input must be a numpy array. Got {type(img)}"
+    assert img.shape[0] in [3,4], f"Input image must be formatted as CHW, with C = 3,4. Got a shape of {img.shape}"
+    assert img.dtype == np.uint8, f"Input image must be a numpy array with dtype np.uint8. Got {img.dtype}"
+    
+    # prepare model for evaluation
+    model = model.to(device)
+    model.eval()
+
+    # prepare alpha channel
+    original_shape = img.shape
+    if img.shape[0] == 3:
+        # create dummy alpha channel
+        alpha = np.full(original_shape[1:], 255, dtype=np.uint8)
+    else:
+        # extract alpha channel
+        img, alpha = img[:3], img[3]
+    
+    # resize image
+    img = resize(img, scaling_factor=scaling_factor)
+
+    # pad image
+    pad_t, pad_b, pad_l, pad_r = minimum_needed_padding(img.shape[1:], patch_size, stride)
+    padded_img = pad(img, pad=(pad_t, pad_b, pad_l, pad_r))
+    padded_shape = padded_img.shape
+
+    # extract patches indexes
+    patches_idx = extract_patches(padded_img, patch_size=patch_size, stride=stride)
+
+    ### segment
+    masks = []
+    masks_idx = []
+
+    batch = []
+    for i, p_idx in enumerate(tqdm(patches_idx, disable=not verbose, desc="Predicting...", total=len(patches_idx))):
+        t, b, l, r = p_idx
+
+        # extract patch
+        patch = padded_img[:, t:b, l:r]
+
+        # consider patch only if it is valid (i.e. not all black or all white)
+        if np.any(patch != 0) and np.any(patch != 255):
+
+            # convert patch to torch.tensor with float32 values in [0,1] (as required by torch)
+            patch = torch.tensor(patch).float() / 255.
+
+            # normalize patch with ImageNet mean and std
+            patch = (patch - torch.tensor([0.485, 0.456, 0.406]).view(3,1,1)) / torch.tensor([0.229, 0.224, 0.225]).view(3,1,1)
+
+            # add patch to batch
+            batch.append(patch)
+            masks_idx.append(p_idx)
+
+            # (optional) for each patch extracted, consider also its rotated versions
+            if rotate:
+                for rot in range(1,4):
+                    patch = torch.rot90(patch, rot, dims=[1,2])
+                    batch.append(patch)
+                    masks_idx.append(p_idx)
+        
+        # if the batch is full, perform prediction
+        if len(batch) >= batch_size or i == len(patches_idx)-1:
+
+            # move batch to GPU
+            batch = torch.stack(batch).to(device)
+
+            # perform prediction
+            out = segment_batch(batch, model)
+
+            # append predictions to masks
+            masks.append(out.cpu().numpy())
+
+            # reset batch
+            batch = []
+
+    # concatenate predictions
+    masks = np.concatenate(masks)  # (n_patches, 1, H, W)
+
+    # merge patches
+    mask = merge_patches(masks, masks_idx, rotate=rotate, canvas_shape=padded_shape[1:])    # (1, H, W)
+
+    # undo padding
+    mask = mask[:, pad_t:padded_shape[1]-pad_b, pad_l:padded_shape[2]-pad_r]
+
+    # resize mask to original shape
+    mask = resize(mask, shape=original_shape[1:])
+
+    # apply alpha channel, i.e. set to -1 the pixels where alpha is 0
+    mask = np.where(alpha == 0, -1, mask)
+
+    return mask.squeeze()
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+def sliding_window_avg_pooling(img, window, granularity, alpha=None, min_nonblank_pixels=0., order="HWC", normalize=False, return_min_max=False, verbose=False):
+    assert isinstance(img, np.ndarray), f'Input image must be a numpy array. Got {type(img)}'
+    if order == "HWC":
+        assert img.shape[2] == 1, f'Input image must be formatted as HWC, with C = 1. Got a shape of {img.shape}'
+    elif order == "CHW":
+        assert img.shape[0] == 1, f'Input image must be formatted as CHW, with C = 1. Got a shape of {img.shape}'
+
+    # check if alpha channel was given, and cast it to np.float32 with values in [0,1]
+    if alpha is not None:
+        assert img.shape == alpha.shape, f'The shape of input image {img.shape} and alpha channel {alpha.shape} do not match'
+        if alpha.dtype == np.uint8:
+            alpha = (alpha / 255).astype(np.float32)
+        elif alpha.dtype == bool:
+            alpha = alpha.astype(np.float32)
+    else:
+        alpha = np.ones_like(img, dtype=np.float32)
+
+    # compute threshold
+    thresh = min_nonblank_pixels * window**2
+
+    # extract patches idxs
+    patches_idx = extract_patches(img, patch_size=window, stride=granularity, order=order, only_return_idx=True)
+
+    # initialize canvas
+    canvas = np.zeros_like(img, dtype=np.float32)
+    n_overlapping_patches = np.zeros_like(img, dtype=np.float32)
+
+    # cycle through patches idxs
+    for t,b,l,r in tqdm(patches_idx, disable=not verbose):
+        p_a = alpha[t:b,l:r]
+        n_valid_pixels = p_a.sum()
+        # keep only if it has more than min_nonblank_pixels
+        if n_valid_pixels <= thresh:
+            continue
+
+        # compute average patch value (i.e. density inside the patch)
+        p = img[t:b,l:r]
+        p_density = (p * p_a).sum() / n_valid_pixels
+
+        # add to canvas
+        canvas[t:b,l:r] += p_density
+        n_overlapping_patches[t:b,l:r] += 1
+
+    # compute average density
+    density_map = np.divide(canvas, n_overlapping_patches, where=(n_overlapping_patches != 0))
+
+    # apply alpha
+    density_map = density_map * alpha
+
+    if normalize:
+        # [0,1]-normalize
+        density_map_min = density_map.min()
+        density_map_max = density_map.max()
+        density_map = (density_map - density_map_min) / (density_map_max - density_map_min)
+        
+        if return_min_max:
+            return density_map, density_map_min, density_map_max
+
+    return density_map
+
+
+
+def compute_vndvi(
+        raster: np.ndarray,
+        mask: np.ndarray,
+        dilate_rows=True,
+        window_size=360,
+        granularity=45,
+        ):
+    assert isinstance(raster, np.ndarray)
+    assert isinstance(mask, np.ndarray)
+    assert len(raster.shape) == 3   # CHW
+    assert len(mask.shape) == 2 # HW
+    assert raster.shape[0] in [3,4] # RGB or RGBA
+
+    # CHW -> HWC
+    raster = raster.transpose(1,2,0)
+
+    # Extract channels
+    _raster = raster.astype(np.float32) / 255 # convert to float32 in [0,1]
+    R, G, B = _raster[:,:,0], _raster[:,:,1], _raster[:,:,2]
+
+    # To avoid division by 0 due to negative power, we replace 0 with 1 in R and B channels
+    R = np.where(R == 0, 1, R)
+    B = np.where(B == 0, 1, B)
+
+    # Mask has values: 0=interrows, 255=rows, 1=nodata
+    # Get mask for the rows and interrows
+    mask_rows = (mask == 255)
+    mask_interrows = (mask == 0)
+    mask_valid = mask_rows | mask_interrows
+
+    # Compute vndvi
+    vndvi = 0.5268 * (R**(-0.1294) * G**(0.3389) * B**(-0.3118))
+
+    # Clip values to [0,1]
+    vndvi = np.clip(vndvi, 0, 1)
+
+    # Compute 10th and 90th percentile on whole vineyard vndvi heatmap
+    vndvi_perc10, vndvi_perc90 = np.percentile(vndvi[mask_valid], [10,90])
+
+    # Clip values between 10th and 90th percentile
+    vndvi_clipped = np.clip(vndvi, vndvi_perc10, vndvi_perc90)
+
+    # Perform sliding window average pooling to smooth the heatmap
+    # NB: the window takes into account only the rows
+    vndvi_rows_clipped_pooled = sliding_window_avg_pooling(
+        np.where(mask_rows, vndvi_clipped, 0)[..., None],
+        window = int(window_size / 4),
+        granularity = granularity,
+        alpha = mask_rows[..., None],
+        min_nonblank_pixels = 0.0,
+        verbose=True,
+        )
+    # Same, but for interrows
+    vndvi_interrows_clipped_pooled = sliding_window_avg_pooling(
+        np.where(mask_interrows, vndvi_clipped, 0)[..., None],
+        window = int(window_size / 4),
+        granularity = granularity,
+        alpha = mask_interrows[..., None],
+        min_nonblank_pixels = 0.0,
+        verbose=True,
+        )
+
+    # Apply dilation to rows mask
+    dil_factor = int(window_size / 60)
+    mask_rows_dilated = grey_dilation(mask_rows, size=(dil_factor, dil_factor))
+    vndvi_rows_clipped_pooled_dilated = grey_dilation(vndvi_rows_clipped_pooled, size=(dil_factor, dil_factor, 1))
+            
+    # For visualization purposes, normalize with vndvi_perc10 and
+    # vndvi_perc90 (because we want vndvi_perc10 to be the first color of
+    # the colormap and vndvi_perc90 to be the last)
+    vndvi_rows_clipped_pooled_normalized = (vndvi_rows_clipped_pooled - vndvi_perc10) / (vndvi_perc90 - vndvi_perc10)
+    vndvi_rows_clipped_pooled_dilated_normalized = (vndvi_rows_clipped_pooled_dilated - vndvi_perc10) / (vndvi_perc90 - vndvi_perc10)
+    vndvi_interrows_clipped_pooled_normalized = (vndvi_interrows_clipped_pooled - vndvi_perc10) / (vndvi_perc90 - vndvi_perc10)
+
+    # for visualization
+    vndvi_rows_img = alpha_composite(
+        raster,
+        vndvi_rows_clipped_pooled_dilated_normalized if dilate_rows else vndvi_rows_clipped_pooled_normalized,
+        opacity = 1.0,
+        colormap = 'RdYlGn',
+        alpha_image = np.zeros_like(raster[:,:,[0]]),
+        alpha_mask = mask_rows_dilated[...,None] if dilate_rows else mask_rows[...,None],
+        )   # HW4 RGBA
+
+    vndvi_interrows_img = alpha_composite(
+        raster,
+        vndvi_interrows_clipped_pooled_normalized,
+        opacity = 1.0,
+        colormap = 'RdYlGn',
+        alpha_image = np.zeros_like(raster[:,:,[0]]),
+        alpha_mask = mask_interrows[...,None],
+        )   # HW4 RGBA
+
+    # add colorbar
+    # fig_rows, ax = plt.subplots(1, 1, figsize=(10, 10))
+    # divider = make_axes_locatable(ax)
+    # cax = divider.append_axes('right', size='5%', pad=0.15)
+    # ax.imshow(vndvi_rows_img)
+    # fig_rows.colorbar(
+    #     mappable = mpl.cm.ScalarMappable(
+    #         norm = mpl.colors.Normalize(
+    #             vmin = vndvi_perc10,
+    #             vmax = vndvi_perc90),
+    #         cmap = 'RdYlGn'),
+    #     cax = cax,
+    #     orientation = 'vertical',
+    #     label = 'vNDVI',
+    #     shrink = 1)
+
+    # fig_interrows, ax = plt.subplots(1, 1, figsize=(10, 10))
+    # divider = make_axes_locatable(ax)
+    # cax = divider.append_axes('right', size='5%', pad=0.15)
+    # ax.imshow(vndvi_interrows_img)
+    # fig_interrows.colorbar(
+    #     mappable = mpl.cm.ScalarMappable(
+    #         norm = mpl.colors.Normalize(
+    #             vmin = vndvi_perc10,
+    #             vmax = vndvi_perc90),
+    #         cmap = 'RdYlGn'),
+    #     cax = cax,
+    #     orientation = 'vertical',
+    #     label = 'vNDVI',
+    #     shrink = 1)
+    
+    # return fig_rows, fig_interrows
+    return vndvi_rows_img, vndvi_interrows_img
+
+
+
+def compute_vdi(
+        raster: np.ndarray,
+        mask: np.ndarray,
+        window_size=360,
+        granularity=40,
+        ):
+
+    # CHW -> HWC
+    raster = raster.transpose(1,2,0)
+
+    # Mask has values: 0=interrows, 255=rows, 1=nodata
+    # Get mask for the rows and interrows
+    mask_rows = (mask == 255)
+    mask_interrows = (mask == 0)
+    mask_valid = mask_rows | mask_interrows
+
+    # compute vdi
+    vdi, vdi_min, vdi_max = sliding_window_avg_pooling(
+        mask_rows[...,None],
+        window=window_size,
+        granularity=granularity,
+        alpha=mask_valid[...,None],
+        min_nonblank_pixels=0.9,
+        normalize=True,
+        return_min_max=True,
+        verbose=True,
+        )
+
+    # for visualization
+    vdi_img = alpha_composite(
+        raster,
+        vdi,
+        opacity = 1,
+        colormap = 'jet_r',
+        alpha_image = mask_valid[...,None],
+        alpha_mask = mask_valid[...,None],
+        )
+
+    # add colorbar
+    # fig, ax = plt.subplots(1, 1, figsize=(10, 10))
+    # divider = make_axes_locatable(ax)
+    # cax = divider.append_axes('right', size='5%', pad=0.15)
+    # ax.imshow(vdi_img)
+    # fig.colorbar(
+    #     mappable = mpl.cm.ScalarMappable(
+    #         norm = mpl.colors.Normalize(
+    #             vmin = vdi_min,
+    #             vmax = vdi_max),
+    #         cmap = 'jet_r'),
+    #     cax = cax,
+    #     orientation = 'vertical',
+    #     label = 'VDI',
+    #     shrink = 1)
+
+    # return fig
+    return vdi_img
+
+
+
+def compute_mask(
+        raster: np.ndarray,
+        model: torch.nn.Module,
+        patch_size=512,
+        stride=256,
+        scaling_factor=None,
+        rotate=False,
+        batch_size=16
+        ):
+    assert isinstance(raster, np.ndarray), f'Input raster must be a numpy array. Got {type(raster)}'
+    assert len(raster.shape) == 3, f'Input raster must have 3 dimensions (bands, rows, cols). Got shape {raster.shape}'
+    assert raster.shape[0] in [3,4], f'Input raster must have 3 bands (RGB) or 4 bands (RGBA). Got {raster.shape[0]} bands'
+    assert isinstance(model, torch.nn.Module), 'Model must be a torch.nn.Module'
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    # Infer GSD
+    #gsd = abs(raster.rio.transform()[0])  # ground sampling distance (NB: valid only if image is a GeoTIFF)
+
+    # Growseg works best on orthoimages with gsd in [1, 1.7] cm/px. You may want to
+    # specify a scaling factor different from 1 if your image has a different gsd.
+    # E.g.: SCALING_FACTOR = gsd / 0.015
+    # logger.info(f'Image GSD: {gsd*100:.2f} cm/px')
+    # scaling_factor = scaling_factor or (gsd / 0.015)
+    scaling_factor = scaling_factor or 1
+    logger.info(f'Applying scaling factor: {scaling_factor:.2f}')
+
+    # segment
+    logger.info('Segmenting image...')
+    score_map = segment(
+        raster,
+        model,
+        patch_size=patch_size,
+        stride=stride,
+        scaling_factor=scaling_factor,
+        rotate=rotate,
+        device=device,
+        batch_size=batch_size,
+        verbose=True
+        )   # mask is a HxW float32 array in [0, 1]
+    
+    # apply threshold on confidence scores
+    alpha = (score_map == -1)
+    mask = (score_map > 0.5)
+
+    # convert to uint8
+    mask = (mask * 255).astype(np.uint8)
+
+    # set nodata pixels to 1
+    mask[alpha] = 1
+
+    return mask

lib/viz_utils.py

@@ -0,0 +1,125 @@
+import sys
+import functools
+import numpy as np
+import cv2
+import cmapy
+from PIL import Image
+import matplotlib
+
+
+
+# BUGFIX in cmapy.py
+def cmap(cmap_name, rgb_order=False):
+    """
+    Extract colormap color information as a LUT compatible with cv2.applyColormap().
+    Default channel order is BGR.
+
+    Args:
+        cmap_name: string, name of the colormap.
+        rgb_order: boolean, if false or not set, the returned array will be in
+                   BGR order (standard OpenCV format). If true, the order
+                   will be RGB.
+
+    Returns:
+        A numpy array of type uint8 containing the colormap.
+    """
+
+    c_map = matplotlib.colormaps.get_cmap(cmap_name)
+    rgba_data = matplotlib.cm.ScalarMappable(cmap=c_map).to_rgba(
+        np.arange(0, 1.0, 1.0 / 256.0), bytes=True
+    )
+    rgba_data = rgba_data[:, 0:-1].reshape((256, 1, 3))
+
+    # Convert to BGR (or RGB), uint8, for OpenCV.
+    cmap = np.zeros((256, 1, 3), np.uint8)
+
+    if not rgb_order:
+        cmap[:, :, :] = rgba_data[:, :, ::-1]
+    else:
+        cmap[:, :, :] = rgba_data[:, :, :]
+
+    return cmap
+
+# If python 3, redefine cmap() to use lru_cache.
+if sys.version_info > (3, 0):
+    cmap = functools.lru_cache(maxsize=200)(cmap)
+
+
+
+def alpha_composite(img, msk, opacity=0.5, colormap=None, alpha_image=None, alpha_mask=None, red_mask=False):
+    """Alpha composite an RGBA image (img) and a grayscale mask (msk).
+    - If alpha_image is None, img's alpha channel is used (or, if not present,
+    initialized to all 255).
+    - If alpha_mask is None, msk is overlaid on img only where img's alpha
+    channel is not 0.
+    - If alpha_mask is not None, the above behavior is overridden and msk is
+    overlaid on img only where alpha_mask is not 0."""
+    # only HWC numpy arrays allowed
+    assert isinstance(img, np.ndarray), f'Input image must be a numpy array. Got {type(img)}'
+    assert isinstance(msk, np.ndarray), f'Input mask must be a numpy array. Got {type(msk)}'
+    if alpha_mask is not None:
+        assert isinstance(alpha_mask, np.ndarray), f'Alpha mask must be a numpy array. Got {type(alpha_mask)}'
+        assert alpha_mask.dtype in [np.float32, bool], f'Alpha mask must be of type np.float32 or bool. Got {alpha_mask.dtype}'
+        assert alpha_mask.shape[2] == 1, f'Alpha mask must be formatted as HWC, with C = 1. Got a shape of {msk.shape}'
+    assert img.shape[2] in [3,4], f'Input image must be formatted as HWC, with C = 3,4. Got a shape of {img.shape}'
+    assert msk.shape[2] == 1, f'Input mask must be formatted as HWC, with C = 1. Got a shape of {msk.shape}'
+    assert (opacity >= 0) and (opacity <= 1), f'Mask opacity must be between 0 and 1. Got {opacity}'
+    
+    # to avoid modifying the original arrays
+    img = img.copy()
+    msk = msk.copy()
+
+    if img.shape[2] == 3:
+        # add alpha channel to img
+        img = np.concatenate([
+            img,
+            np.full((img.shape[0], img.shape[1], 1), 255, dtype=np.uint8)
+            ], axis=-1)
+
+    if alpha_image is None:
+        # initialize alpha_image to all Trues
+        alpha_image = img[:,:,[3]]
+    # convert alpha image to bool
+    alpha_image = alpha_image.astype(bool)
+
+    if alpha_mask is None:
+        # initialize alpha_mask to alpha_image
+        alpha_mask = alpha_image    # so that alpha_mask is AT LEAST as restrictive as alpha_image
+    # convert alpha mask to bool
+    alpha_mask = alpha_mask.astype(bool)
+
+
+    if msk.dtype != np.uint8:
+        # convert mask to a uint8 grayscale image ([0,1] -> [0,255])
+        # NB: normalize the pixels of the mask we are interested in to [0,1]
+        # before passing it as input!!!
+        msk = (msk * 255).astype(np.uint8)
+
+    # convert mask from grayscale to RGBA
+    msk = cv2.cvtColor(msk, cv2.COLOR_GRAY2RGBA)
+
+    if colormap is not None:
+        # apply specified colormap to msk
+        # NB: values near 0 will be converted to the first colors of the chosen
+        # colormap, whereas values near 255 will be converted to the last colors
+        msk[:,:,:3] = cmapy.colorize(msk[:,:,:3], colormap, rgb_order=True)
+    elif red_mask:
+        # convert white to red
+        msk[:,:,[1,2]] = 0
+
+
+    # apply alpha_image to img's alpha channel
+    img[:,:,[3]] = (alpha_image * img[:,:,[3]]).astype(np.uint8)
+
+    # apply alpha_mask and opacity to msk's alpha channel
+    msk[:,:,[3]] = (alpha_mask * opacity * msk[:,:,[3]]).astype(np.uint8)
+
+    # alpha compositing
+    img_pil = Image.fromarray(img)
+    msk_pil = Image.fromarray(msk)
+    img_pil.alpha_composite(msk_pil)
+
+    return np.array(img_pil)
+
+
+

precompute_examples.ipynb

@@ -0,0 +1,358 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "import geopandas as gpd\n",
+    "import rioxarray as rxr\n",
+    "import xarray as xr\n",
+    "import numpy as np\n",
+    "import os\n",
+    "import torch\n",
+    "from transformers import SegformerForSemanticSegmentation\n",
+    "from lib.utils import compute_mask, compute_vndvi, compute_vdi"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# # Read raster data\n",
+    "# raster_path = \"data/spain_2022-07-29.tif\"\n",
+    "# raster = rxr.open_rasterio(raster_path)\n",
+    "\n",
+    "# # Crop raster with GeoJSON geometry, if available\n",
+    "# geom_path = raster_path.replace(\".tif\", \".geojson\")\n",
+    "# if os.path.exists(geom_path):\n",
+    "#     geom = gpd.read_file(geom_path)\n",
+    "#     raster = raster.rio.clip(geom.geometry)\n",
+    "#     raster.rio.to_raster(raster_path.replace(\".tif\", \"_cropped.tif\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def load_model(hf_path='links-ads/gaia-growseg'):\n",
+    "    # logger.info(f'Loading GAIA GRowSeg on {device}...')\n",
+    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
+    "    model = SegformerForSemanticSegmentation.from_pretrained(\n",
+    "        hf_path,\n",
+    "        num_labels=1,\n",
+    "        num_channels=3,\n",
+    "        id2label={1: 'vine'},\n",
+    "        label2id={'vine': 1},\n",
+    "        token=os.getenv('hf_read_access_token')\n",
+    "    )\n",
+    "    return model.to(device).eval()\n",
+    "\n",
+    "# Load GAIA GRowSeg model\n",
+    "model = load_model()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-03-20 12:39:09.921\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mlib.utils\u001b[0m:\u001b[36msliding_window_avg_pooling\u001b[0m:\u001b[36m308\u001b[0m - \u001b[1mExtracting patches idx...\u001b[0m\n",
+      "100%|█████████████████████████████████████████████| 67848/67848 [00:03<00:00, 20745.29it/s]\n",
+      "\u001b[32m2025-03-20 12:39:14.795\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mlib.utils\u001b[0m:\u001b[36msliding_window_avg_pooling\u001b[0m:\u001b[36m308\u001b[0m - \u001b[1mExtracting patches idx...\u001b[0m\n",
+      "100%|█████████████████████████████████████████████| 67848/67848 [00:03<00:00, 19329.36it/s]\n",
+      "\u001b[32m2025-03-20 12:39:56.011\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mlib.utils\u001b[0m:\u001b[36msliding_window_avg_pooling\u001b[0m:\u001b[36m308\u001b[0m - \u001b[1mExtracting patches idx...\u001b[0m\n",
+      "100%|██████████████████████████████████████████████| 64758/64758 [00:20<00:00, 3203.45it/s]\n"
+     ]
+    }
+   ],
+   "source": [
+    "raster_path = \"data/italy_2022-06-13_cropped.tif\"\n",
+    "patch_size = 512\n",
+    "stride = 256\n",
+    "scaling_factor = 1.0\n",
+    "dilate_rows = False\n",
+    "window_size = 360\n",
+    "granularity = int(window_size/8)\n",
+    "\n",
+    "# raster_path = \"data/spain_2022-07-29_cropped.tif\"\n",
+    "# patch_size = 512\n",
+    "# stride = 256\n",
+    "# scaling_factor = 1.0\n",
+    "# dilate_rows = False\n",
+    "# window_size = 400\n",
+    "# granularity = int(window_size/8)\n",
+    "\n",
+    "# raster_path = \"data/portugal_2023-08-01.tif\"\n",
+    "# patch_size = 512\n",
+    "# stride = 256\n",
+    "# scaling_factor = 1.25\n",
+    "# dilate_rows = False\n",
+    "# window_size = 80\n",
+    "# granularity = int(window_size/8)\n",
+    "\n",
+    "raster = rxr.open_rasterio(raster_path)\n",
+    "\n",
+    "# Compute mask\n",
+    "mask_path = raster_path.replace(\".tif\", \"_mask.tif\")\n",
+    "if not os.path.exists(mask_path):\n",
+    "    mask = compute_mask(\n",
+    "        raster.to_numpy(),\n",
+    "        model,\n",
+    "        patch_size=patch_size,\n",
+    "        stride=stride,\n",
+    "        scaling_factor=scaling_factor,\n",
+    "        rotate=False,\n",
+    "        batch_size=16,\n",
+    "    )   # mask is a HxW uint8 array in with 0=background, 255=vine, 1=nodata\n",
+    "\n",
+    "    # Convert mask from grayscale to RGBA, with red pixels for vine\n",
+    "    alpha = ((mask != 1)*255).astype(np.uint8)\n",
+    "    mask_colored = np.stack([mask, np.zeros_like(mask), np.zeros_like(mask), alpha], axis=0)  # now, mask is a 4xHxW uint8 array in with 0=background, 255=vine\n",
+    "\n",
+    "    # Georef mask like raster\n",
+    "    mask_raster = xr.DataArray(\n",
+    "        mask_colored,\n",
+    "        dims=('band', 'y', 'x'),\n",
+    "        coords={'x': raster.x, 'y': raster.y, 'band': raster.band}\n",
+    "        )\n",
+    "    mask_raster.rio.write_crs(raster.rio.crs, inplace=True)  # Copy CRS\n",
+    "    mask_raster.rio.write_transform(raster.rio.transform(), inplace=True)  # Copy affine transform\n",
+    "\n",
+    "    # Save mask\n",
+    "    mask_raster.rio.to_raster(raster_path.replace(\".tif\", \"_mask.tif\"), compress='lzw')\n",
+    "else:\n",
+    "    mask = rxr.open_rasterio(mask_path).sel(band=1).squeeze().to_numpy()\n",
+    "\n",
+    "# Compute vNDVI\n",
+    "vndvi_rows_path = raster_path.replace(\".tif\", \"_vndvi_rows.tif\")\n",
+    "vndvi_interrows_path = raster_path.replace(\".tif\", \"_vndvi_interrows.tif\")\n",
+    "if not os.path.exists(vndvi_rows_path) or not os.path.exists(vndvi_interrows_path):\n",
+    "    vndvi_rows, vndvi_interrows = compute_vndvi(\n",
+    "        raster.to_numpy(),\n",
+    "        mask,\n",
+    "        dilate_rows=dilate_rows,\n",
+    "        window_size=window_size,\n",
+    "        granularity=granularity,\n",
+    "        )    # vNDVI is RGBA\n",
+    "\n",
+    "    # Georef vNDVI like raster\n",
+    "    vndvi_rows_raster = xr.DataArray(\n",
+    "        vndvi_rows.transpose(2, 0, 1),\n",
+    "        dims=('band', 'y', 'x'),\n",
+    "        coords={'x': raster.x, 'y': raster.y, 'band': raster.band}\n",
+    "        )\n",
+    "    vndvi_rows_raster.rio.write_crs(raster.rio.crs, inplace=True)\n",
+    "    vndvi_rows_raster.rio.write_transform(raster.rio.transform(), inplace=True)\n",
+    "\n",
+    "    vndvi_interrows_raster = xr.DataArray(\n",
+    "        vndvi_interrows.transpose(2, 0, 1),\n",
+    "        dims=('band', 'y', 'x'),\n",
+    "        coords={'x': raster.x, 'y': raster.y, 'band': raster.band}\n",
+    "        )\n",
+    "    vndvi_interrows_raster.rio.write_crs(raster.rio.crs, inplace=True)\n",
+    "    vndvi_interrows_raster.rio.write_transform(raster.rio.transform(), inplace=True)\n",
+    "\n",
+    "    # Save vNDVI\n",
+    "    vndvi_rows_raster.rio.to_raster(raster_path.replace(\".tif\", \"_vndvi_rows.tif\"), compress='lzw')\n",
+    "    vndvi_interrows_raster.rio.to_raster(raster_path.replace(\".tif\", \"_vndvi_interrows.tif\"), compress='lzw')\n",
+    "\n",
+    "# Compute VDI\n",
+    "vdi_path = raster_path.replace(\".tif\", \"_vdi.tif\")\n",
+    "if not os.path.exists(vdi_path):\n",
+    "    vdi = compute_vdi(\n",
+    "        raster.to_numpy(),\n",
+    "        mask,\n",
+    "        window_size=window_size,\n",
+    "        granularity=granularity,\n",
+    "        )    # VDI is RGBA\n",
+    "\n",
+    "    # Georef VDI like raster\n",
+    "    vdi_raster = xr.DataArray(\n",
+    "        vdi.transpose(2, 0, 1),\n",
+    "        dims=('band', 'y', 'x'),\n",
+    "        coords={'x': raster.x, 'y': raster.y, 'band': raster.band}\n",
+    "        )\n",
+    "    vdi_raster.rio.write_crs(raster.rio.crs, inplace=True)\n",
+    "    vdi_raster.rio.write_transform(raster.rio.transform(), inplace=True)\n",
+    "\n",
+    "    # Save results\n",
+    "    vdi_raster.rio.to_raster(raster_path.replace(\".tif\", \"_vdi.tif\"), compress='lzw')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-03-20 12:40:30.816\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m76\u001b[0m - \u001b[1mReprojecting rasters to EPSG:4326 with NODATA value 0...\u001b[0m\n",
+      "\u001b[32m2025-03-20 12:40:52.371\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m84\u001b[0m - \u001b[1mCreating RGB raster overlay...\u001b[0m\n",
+      "\u001b[32m2025-03-20 12:40:52.373\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mcreate_image_overlay\u001b[0m:\u001b[36m46\u001b[0m - \u001b[1mCreating overlay: 'Orthoimage'...\u001b[0m\n",
+      "\u001b[32m2025-03-20 12:40:58.801\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m86\u001b[0m - \u001b[1mCreating mask overlay...\u001b[0m\n",
+      "\u001b[32m2025-03-20 12:40:58.806\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mcreate_image_overlay\u001b[0m:\u001b[36m46\u001b[0m - \u001b[1mCreating overlay: 'Mask'...\u001b[0m\n",
+      "\u001b[32m2025-03-20 12:41:05.006\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m88\u001b[0m - \u001b[1mCreating vNDVI rows overlay...\u001b[0m\n",
+      "\u001b[32m2025-03-20 12:41:05.008\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mcreate_image_overlay\u001b[0m:\u001b[36m46\u001b[0m - \u001b[1mCreating overlay: 'vNDVI Rows'...\u001b[0m\n",
+      "\u001b[32m2025-03-20 12:41:10.988\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m90\u001b[0m - \u001b[1mCreating vNDVI interrows overlay...\u001b[0m\n",
+      "\u001b[32m2025-03-20 12:41:10.990\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mcreate_image_overlay\u001b[0m:\u001b[36m46\u001b[0m - \u001b[1mCreating overlay: 'vNDVI Interrows'...\u001b[0m\n",
+      "\u001b[32m2025-03-20 12:41:16.558\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m92\u001b[0m - \u001b[1mCreating VDI overlay...\u001b[0m\n",
+      "\u001b[32m2025-03-20 12:41:16.560\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mcreate_image_overlay\u001b[0m:\u001b[36m46\u001b[0m - \u001b[1mCreating overlay: 'VDI'...\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "import folium\n",
+    "from loguru import logger\n",
+    "\n",
+    "def create_map(location=[41.9099533, 12.3711879], zoom_start=5, crs=3857, max_zoom=23):\n",
+    "    \"\"\"Create a folium map with OpenStreetMap tiles and optional Esri.WorldImagery basemap.\"\"\"\n",
+    "    if isinstance(crs, int):\n",
+    "        crs = f\"EPSG{crs}\"\n",
+    "    assert crs in [\"EPSG3857\"], f\"Only EPSG:3857 supported for now. Got {crs}.\"\n",
+    "    \n",
+    "    m = folium.Map(\n",
+    "        location=location,\n",
+    "        zoom_start=zoom_start,\n",
+    "        crs=crs,\n",
+    "        max_zoom=max_zoom,\n",
+    "        tiles=\"OpenStreetMap\",  # Esri.WorldImagery\n",
+    "        attributionControl=False,\n",
+    "        prefer_canvas=True,\n",
+    "    )\n",
+    "\n",
+    "    # Add Esri.WorldImagery as optional basemap (radio button)\n",
+    "    folium.TileLayer(\n",
+    "        tiles=\"Esri.WorldImagery\",\n",
+    "        show=False,\n",
+    "        overlay=False,\n",
+    "        control=True,\n",
+    "    ).add_to(m)\n",
+    "\n",
+    "    return m\n",
+    "\n",
+    "def create_image_overlay(raster_path_or_array, name=\"Raster\", opacity=1.0, to_crs=4326, show=True):\n",
+    "    \"\"\" Create a folium image overlay from a raster filepath or xarray.DataArray. \"\"\"\n",
+    "    if isinstance(raster_path_or_array, str):\n",
+    "        # Open the raster and its metadata\n",
+    "        logger.info(f\"Opening raster: {raster_path_or_array!r}...\")\n",
+    "        r = rxr.open_rasterio(raster_path_or_array)\n",
+    "    else:\n",
+    "        r = raster_path_or_array\n",
+    "    nodata = r.rio.nodata or 0\n",
+    "    if r.rio.crs.to_epsg() != to_crs:\n",
+    "        logger.info(f\"Reprojecting raster to EPSG:{to_crs} with NODATA value {nodata}...\")\n",
+    "        r = r.rio.reproject(to_crs, nodata=nodata) # nodata default: 255\n",
+    "    r = r.transpose(\"y\", \"x\", \"band\")\n",
+    "    bounds = r.rio.bounds()   # (left, bottom, right, top)\n",
+    "\n",
+    "    # Create a folium image overlay\n",
+    "    logger.info(f\"Creating overlay: {name!r}...\")\n",
+    "    overlay = folium.raster_layers.ImageOverlay(\n",
+    "        image=r.to_numpy(),\n",
+    "        name=name,\n",
+    "        bounds=[[bounds[1], bounds[0]], [bounds[3], bounds[2]]],    # format for folium: ((bottom,left),(top,right))\n",
+    "        opacity=opacity,\n",
+    "        interactive=True,\n",
+    "        cross_origin=False,\n",
+    "        zindex=1,\n",
+    "        show=show,\n",
+    "    )\n",
+    "\n",
+    "    return overlay\n",
+    "\n",
+    "# Define paths\n",
+    "raster_path = \"data/portugal_2023-08-01.tif\"\n",
+    "mask_path = raster_path.replace('.tif', '_mask.tif')\n",
+    "vndvi_rows_path = raster_path.replace('.tif', '_vndvi_rows.tif')\n",
+    "vndvi_interrows_path = raster_path.replace('.tif', '_vndvi_interrows.tif')\n",
+    "vdi_path = raster_path.replace('.tif', '_vdi.tif')\n",
+    "\n",
+    "# Load rasters\n",
+    "raster = rxr.open_rasterio(raster_path)\n",
+    "mask_raster = rxr.open_rasterio(mask_path)\n",
+    "vndvi_rows_raster = rxr.open_rasterio(vndvi_rows_path)\n",
+    "vndvi_interrows_raster = rxr.open_rasterio(vndvi_interrows_path)\n",
+    "vdi_raster = rxr.open_rasterio(vdi_path)\n",
+    "\n",
+    "# Reproject all rasters to EPSG:4326\n",
+    "if raster.rio.crs.to_epsg() != 4326:\n",
+    "    logger.info(f\"Reprojecting rasters to EPSG:4326 with NODATA value 0...\")\n",
+    "    raster = raster.rio.reproject(\"EPSG:4326\", nodata=0)    # nodata default: 255\n",
+    "    mask_raster = mask_raster.rio.reproject(\"EPSG:4326\", nodata=0)\n",
+    "    vndvi_rows_raster = vndvi_rows_raster.rio.reproject(\"EPSG:4326\", nodata=0)\n",
+    "    vndvi_interrows_raster = vndvi_interrows_raster.rio.reproject(\"EPSG:4326\", nodata=0)\n",
+    "    vdi_raster = vdi_raster.rio.reproject(\"EPSG:4326\", nodata=0)\n",
+    "\n",
+    "# Create overlays\n",
+    "logger.info(f'Creating RGB raster overlay...')\n",
+    "raster_overlay = create_image_overlay(raster, name=\"Orthoimage\", opacity=1.0, show=True)\n",
+    "logger.info(f'Creating mask overlay...')\n",
+    "mask_overlay = create_image_overlay(mask_raster, name=\"Mask\", opacity=1.0, show=False)\n",
+    "logger.info(f'Creating vNDVI rows overlay...')\n",
+    "vndvi_rows_overlay = create_image_overlay(vndvi_rows_raster, name=\"vNDVI Rows\", opacity=1.0, show=False)\n",
+    "logger.info(f'Creating vNDVI interrows overlay...')\n",
+    "vndvi_interrows_overlay = create_image_overlay(vndvi_interrows_raster, name=\"vNDVI Interrows\", opacity=1.0, show=False)\n",
+    "logger.info(f'Creating VDI overlay...')\n",
+    "vdi_overlay = create_image_overlay(vdi_raster, name=\"VDI\", opacity=1.0, show=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "m = create_map()\n",
+    "raster_overlay.add_to(m)\n",
+    "mask_overlay.add_to(m)\n",
+    "vndvi_rows_overlay.add_to(m)\n",
+    "vndvi_interrows_overlay.add_to(m)\n",
+    "vdi_overlay.add_to(m)\n",
+    "\n",
+    "# Add layer control\n",
+    "folium.LayerControl().add_to(m)\n",
+    "\n",
+    "# Fit map to bounds\n",
+    "m.fit_bounds(raster_overlay.get_bounds())\n",
+    "\n",
+    "# Save map\n",
+    "map_path = raster_path.replace('.tif', '.html')\n",
+    "m.save(map_path)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

requirements.txt

@@ -0,0 +1,14 @@
+numpy
+scipy
+rasterio
+torch
+transformers
+tqdm
+loguru
+opencv-python-headless
+pillow
+matplotlib
+cmapy
+python-dotenv
+rioxarray
+geopandas