app first commit

README.md

@@ -1,6 +1,6 @@
 ---
-title: Gaia Growseg Demo
-emoji: 🏃
+title: Growseg Demo
+emoji: 🍇
 colorFrom: blue
 colorTo: gray
 sdk: gradio

app.py

@@ -0,0 +1,138 @@
+import gradio as gr
+import torch
+import numpy as np
+from transformers import SegformerForSemanticSegmentation
+from loguru import logger
+import rasterio as rio
+from lib.utils import segment, compute_vndvi, compute_vdi
+import os
+import os.path as osp
+
+# set temp dir for gradio
+os.environ["GRADIO_TEMP_DIR"] = "/nfs/home/monopoli/VITIGEOSS/gaia-growseg-demo/temp"
+
+# set device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+logger.info(f'Using device: {device}')
+
+# load model architecture
+logger.info('Loading model architecture...')
+model = SegformerForSemanticSegmentation.from_pretrained(
+    'nvidia/mit-b5',
+    num_labels = 1, # binary segmentation
+    num_channels = 3,   # RGB
+    id2label = {1: 'vine'},
+    label2id = {'vine': 1},
+)
+
+# load model weights
+logger.info('Loading model weights...')
+device = torch.device(f'cuda:0') if torch.cuda.is_available() else torch.device('cpu')
+model.load_state_dict(torch.load(f"model/growseg.pt", map_location=device))
+model = model.to(device)
+model.eval()
+
+
+def main(input, output, patch_size=512, stride=256, scaling_factor=1., rotate=False, batch_size=16, verbose=False, return_vndvi=True, return_vdi=True, window_size=360):
+    assert osp.splitext(output)[1].lower() in ['.tif', '.tiff', '.png', '.jpg', '.jpeg'], 'Output file format not supported'
+    assert osp.splitext(input)[1].lower() == osp.splitext(output)[1].lower(), f'Input and output file formats must match. Got {osp.splitext(input)[1].lower()} and {osp.splitext(output)[1].lower()} respectively.'
+
+    # read image
+    logger.info(f'Reading image {input}...')
+    with rio.open(input, 'r') as src:
+        image = src.read()
+        profile = src.profile
+        shape = src.shape
+
+    if profile['driver'] == 'GTiff':
+        gsd = profile['transform'][0]  # ground sampling distance (NB: valid only if image is a GeoTIFF)
+    else:
+        gsd = None
+
+    # 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 / 1.5
+
+    # segment
+    logger.info('Segmenting image...')
+    mask = segment(
+        image,
+        model,
+        patch_size=patch_size,
+        stride=stride,
+        scaling_factor=scaling_factor,
+        rotate=rotate,
+        device=device,
+        batch_size=batch_size,
+        verbose=verbose
+        )   # mask is a HxW float32 array in [0, 1]
+    
+    # apply threshold on confidence scores
+    alpha = (mask == -1)
+    mask = (mask > 0.5)
+
+    # convert to uint8
+    mask = (mask * 255).astype(np.uint8)
+
+    # set nodata pixels to 1
+    mask[alpha] = 1
+
+    # if requested, compute additional
+    if return_vndvi:
+        logger.info('Computing VNDVI...')
+        vndvi_rows_fig, vndvi_interrows_fig = compute_vndvi(image, mask, window_size)
+    else:
+        vndvi_rows_fig = vndvi_interrows_fig = None
+
+    if return_vdi:
+        logger.info('Computing VDI...')
+        vdi_fig = compute_vdi(image, mask, window_size)
+    else:
+        vdi_fig = None
+    
+    # save mask
+    """
+    logger.info('Saving mask...')
+    profile.update(
+        dtype=rio.uint8,
+        count=1,
+        compress='lzw',
+    )
+
+    with rio.open(output, 'w', **profile) as dst:
+        dst.write(mask[None, ...])
+
+    logger.info(f'Mask saved to {output}')
+    """
+
+    # return mask and eventually additional outputs
+    return image.transpose(1,2,0), mask, vndvi_rows_fig, vndvi_interrows_fig, vdi_fig
+
+demo = gr.Interface(
+    fn=main,
+    inputs=[
+        gr.File(type='filepath', file_types=['image','.tif','.tiff'], label='Input image path'),
+        gr.Textbox(value='mask.tif', label='Output mask path'),
+        gr.Slider(minimum=128, maximum=512, value=512, step=128, label='Patch size'),
+        gr.Slider(minimum=0, maximum=256, value=256, step=64, label='Stride'),
+        gr.Slider(minimum=0.1, maximum=10, value=1, step=0.05, label='Scaling factor'),
+        gr.Checkbox(value=False, label='Rotate patches'),
+        gr.Slider(minimum=4, maximum=128, value=16, step=4, label='Batch size'),
+        gr.Checkbox(value=False, label='Verbose'),
+        gr.Checkbox(value=True, label='Return VNDVI map'),
+        gr.Checkbox(value=True, label='Return VDI map'),
+        gr.Slider(minimum=10, maximum=600, value=360, step=1, label='Moving window size for computing vNDVI/VDI (suggestion: inversely proportional to the GSD [px/m])'),
+    ],
+    outputs=[
+        gr.Image(type='numpy', format='png', label='Input image'),
+        gr.Image(type='numpy', format='png', label='Predicted mask'),
+        gr.Plot(format='png', label='VNDVI rows (dilated for visibility)'),
+        gr.Plot(format='png', label='VNDVI interrows'),
+        gr.Plot(format='png', label='VDI'),
+    ], # NB: if one of the outputs is None, it will not be displayed in the interface (https://github.com/gradio-app/gradio/issues/500#issuecomment-1046877766)
+    title='Growseg',
+    description='Segment vineyards in orthoimages',
+    delete_cache=[3600,3600],
+)
+
+demo.launch(share=True)

lib/utils.py

@@ -0,0 +1,521 @@
+import numpy as np
+import torch
+import rasterio
+import cv2
+from transformers import SegformerForSemanticSegmentation
+from tqdm import tqdm
+from PIL import Image
+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
+
+
+def read_raster(path, order='CHW'):
+    """Read a raster file and return a numpy array"""
+    assert order in ['HWC', 'CHW'], f"Got unknown order '{order}', expected one of ['HWC','CHW']"
+
+    with rasterio.open(path) as src:
+        img = src.read()
+    
+    if order == 'HWC':
+        img = np.moveaxis(img, 0, -1)
+    
+    return img
+
+def write_raster(path, img, profile, order='CHW'):
+    """Write a numpy array to a raster file"""
+    assert order in ['HWC', 'CHW'], f"Got unknown order '{order}', expected one of ['HWC','CHW']"
+
+    if order == 'HWC':
+        img = np.moveaxis(img, -1, 0)
+    
+    with rasterio.open(path, 'w', **profile) as dst:
+        dst.write(img)
+
+
+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., normalize=False, return_min_max=False, verbose=False):
+    assert isinstance(img, np.ndarray), f'Input image must be a numpy array. Got {type(img)}'
+    assert img.shape[2] == 1, f'Input image must be formatted as HWC, 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 isinstance(alpha, np.ndarray), f'Alpha channel must be a numpy array. Got {type(alpha)}'
+        assert alpha.shape[2] == 1, f'Alpha channel must be formatted as HWC, with C = 1. Got a shape of {alpha.shape}'
+        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)
+
+    # extract patches
+    patches, patches_idx = extract_patches(img, patch_size=window, stride=granularity, order='HWC', only_return_idx=False)
+    patches_alpha, _ = extract_patches(alpha, patch_size=window, stride=granularity, order='HWC', only_return_idx=False)
+
+    # keep only patches with more than min_nonblank_pixels
+    kept_patches = []
+    for i, p_a in tqdm(enumerate(patches_alpha), total=len(patches), disable=not verbose):
+        if p_a.sum() > min_nonblank_pixels * window**2:
+            kept_patches.append(i)
+    patches = [patches[i] for i in kept_patches]
+    patches_idx = [patches_idx[i] for i in kept_patches]
+    patches_alpha = [patches_alpha[i] for i in kept_patches]
+
+    # compute average patch value (i.e. density inside the patch)
+    patches_density = [np.full_like(p_a, (p * p_a).sum() / p_a.sum()) for p, p_a in zip(patches, patches_alpha)]
+
+    # merge patches
+    pooled_img = merge_patches(patches_density, patches_idx, canvas_shape=img.shape[:2], order='HWC')
+    
+    # apply alpha
+    pooled_img = pooled_img * alpha
+
+    if normalize:
+        # [0,1]-normalize
+        pooled_img_min = pooled_img.min()
+        pooled_img_max = pooled_img.max()
+        pooled_img = (pooled_img - pooled_img_min) / (pooled_img_max - pooled_img_min)
+        
+        if return_min_max:
+            return pooled_img, pooled_img_min, pooled_img_max
+
+    return pooled_img
+
+
+
+def compute_vndvi(image, mask, dilate_rows=True, window_size=360):
+    assert image.dtype == np.uint8, f"Input image must be a numpy array with dtype np.uint8. Got {image.dtype}"
+    assert mask.dtype == np.uint8, f"Input mask must be a numpy array with dtype np.uint8. Got {mask.dtype}"
+
+    # CHW -> HWC
+    image = image.transpose(1,2,0)
+
+    # extract channels
+    _image = image.astype(np.float32) / 255 # convert to float32 in [0,1]
+    R, G, B = _image[:,:,0], _image[:,:,1], _image[:,:,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)
+
+    # 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 vndvi rows heatmap
+    #vndvi_rows = np.where(mask == 255, vndvi, np.nan)
+
+    # compute vndvi interrows heatmap
+    #vndvi_interrows = np.where(mask == 0, vndvi, np.nan)
+
+    # compute 10th and 90th percentile on whole vineyard vndvi heatmap
+    vndvi_perc10, vndvi_perc90 = np.percentile(vndvi[mask != 1], [10,90])   # mask is 1 for nodata, 0 or 255 for valid pixels
+
+    # 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 == 255, vndvi_clipped, 0)[...,None],
+        window = int(window_size / 4),
+        granularity = 10,
+        alpha = (mask == 255)[...,None],
+        min_nonblank_pixels = 0.0,
+        )
+    # same, but for interrows
+    vndvi_interrows_clipped_pooled = sliding_window_avg_pooling(
+        np.where(mask == 0, vndvi_clipped, 0)[...,None],
+        window = int(window_size / 4),
+        granularity = 10,
+        alpha = (mask == 0)[...,None],
+        min_nonblank_pixels = 0.0,
+        )
+
+    # apply dilation to rows mask
+    dilate_rows = True
+    if dilate_rows:
+        dil_factor = int(window_size / 60)
+        mask_rows_dilated = grey_dilation(mask == 255, 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(
+        image,
+        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(image[:,:,[0]]),
+        alpha_mask = mask_rows_dilated[...,None] if dilate_rows else (mask == 255)[...,None],
+        )
+
+    vndvi_interrows_img = alpha_composite(
+        image,
+        vndvi_interrows_clipped_pooled_normalized,
+        opacity = 1.0,
+        colormap = 'RdYlGn',
+        alpha_image = np.zeros_like(image[:,:,[0]]),
+        alpha_mask = (mask == 0)[...,None],
+        )
+
+    # 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
+
+
+
+def compute_vdi(image, mask, window_size=360):
+    assert image.dtype == np.uint8, f"Input image must be a numpy array with dtype np.uint8. Got {image.dtype}"
+    assert mask.dtype == np.uint8, f"Input mask must be a numpy array with dtype np.uint8. Got {mask.dtype}"
+
+    # CHW -> HWC
+    image = image.transpose(1,2,0)
+
+    # compute vdi
+    vdi, vdi_min, vdi_max = sliding_window_avg_pooling(
+        (mask == 255)[...,None],
+        window = window_size,
+        granularity = 10,
+        alpha = (mask != 1)[...,None],  # mask is 1 for nodata, 0 or 255 for valid pixels
+        min_nonblank_pixels = 0.9,
+        normalize=True,
+        return_min_max=True
+        )
+
+    # for visualization
+    vdi_img = alpha_composite(
+        image,
+        vdi,
+        opacity = 0.5,
+        colormap = 'jet_r',
+        alpha_image = (mask != 1)[...,None],
+        alpha_mask = (mask != 1)[...,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

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)
+
+
+

model/growseg.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:293f45cfd45e402b2f2e4398aa32e32e2866431e7fa582293718c51bd1317182
+size 338870239

requirements.txt

@@ -0,0 +1,11 @@
+numpy
+scipy
+rasterio
+torch
+transformers
+tqdm
+loguru
+opencv-python-headless
+pillow
+matplotlib
+cmapy