Upload burn_scar_batch_inference_script.py

README.md

@@ -33,6 +33,42 @@ Code for Finetuning is available through [github](https://github.com/NASA-IMPACT
 Configuration used for finetuning is available through [config](https://github.com/NASA-IMPACT/hls-foundation-os/blob/main/fine-tuning-examples/configs/firescars_config.py
 )
 
+To run inference, first install dependencies 
+
+```
+mamba create -n prithvi-burn-scar python=3.10 pycocotools ncurses
+mamba activate prithvi-burn-scar
+pip install --upgrade pip && \
+    pip install -r requirements.txt && \
+    mim install mmcv-full==1.5.0
+```
+
+#### Instructions for downloading from [HuggingFace datasets](https://huggingface.co/datasets)
+
+1. Create account on https://huggingface.co/join
+2. Install `git` following https://git-scm.com/downloads
+3. Install git-lfs with `sudo apt install git-lfs` and `git lfs install`
+4. Run the following command to download the HLS datasets. You may need to
+   enter your HuggingFace username/password to do the `git clone`.
+
+   ```
+   mkdir -p data
+   cd data/
+   git clone https://huggingface.co/datasets/ibm-nasa-geospatial/hls_burn_scars burn_scars
+   tar -xzvf burn_scars/hls_burn_scars.tar.gz -C ./
+   ```
+
+
+With the datasets and the environment, you can now run the inference script.
+
+```
+python burn_scar_batch_inference_script.py \
+-config burn_scars_Prithvi_100M.py \
+-ckpt burn_scars_Prithvi_100M.pth \
+-input data/burn_scars/validation \
+-output data/burn_scars/inference_output \
+-input_type tif
+```
 
 ### Results
 

burn_scar_batch_inference_script.py

@@ -0,0 +1,219 @@
+import argparse
+from mmcv import Config
+from mmcv.runner import (get_dist_info, init_dist, load_checkpoint,wrap_fp16_model)
+from mmseg.models import build_segmentor
+
+import matplotlib.pyplot as plt
+import mmcv
+import torch
+from mmcv.parallel import collate, scatter
+from mmcv.runner import load_checkpoint
+
+from mmseg.datasets.pipelines import Compose
+from mmseg.models import build_segmentor
+
+from mmseg.datasets import build_dataloader, build_dataset, load_flood_test_data
+import rasterio
+import torch
+import torch.nn.functional as F
+
+from torchvision import transforms
+from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
+
+from mmseg.apis import multi_gpu_test, single_gpu_test, init_segmentor
+from . import custom # custom preprocessing for hls
+import pdb
+
+import numpy as np
+import glob
+import os
+
+import time
+
+def parse_args():
+
+    parser = argparse.ArgumentParser(description="Inference on burn scar fine-tuned model")
+    parser.add_argument('-config', help='path to model configuration file')
+    parser.add_argument('-ckpt', help='path to model checkpoint')
+    parser.add_argument('-input', help='path to input images folder for inference')
+    parser.add_argument('-output', help='directory path to save output images')
+    parser.add_argument('-input_type', help='file type of input images',default="tif")
+
+    args = parser.parse_args()
+
+    return args
+
+def open_tiff(fname):
+
+    with rasterio.open(fname, "r") as src:
+
+        data = src.read()
+
+    return data
+
+def write_tiff(img_wrt, filename, metadata):
+
+    """
+    It writes a raster image to file.
+
+    :param img_wrt: numpy array containing the data (can be 2D for single band or 3D for multiple bands)
+    :param filename: file path to the output file
+    :param metadata: metadata to use to write the raster to disk
+    :return:
+    """
+
+    with rasterio.open(filename, "w", **metadata) as dest:
+
+        if len(img_wrt.shape) == 2:
+
+            img_wrt = img_wrt[None]
+
+        for i in range(img_wrt.shape[0]):
+            dest.write(img_wrt[i, :, :], i + 1)
+
+
+def get_meta(fname):
+
+    with rasterio.open(fname, "r") as src:
+
+        meta = src.meta
+
+    return meta
+
+def preprocess_image(data, means, stds, nodata=-9999):
+
+    data=np.where(data == nodata, 0, data)
+    data = data.astype(np.float32)
+
+    if len(data)==2:
+        (x, y) = data
+    else:
+        x=data
+        y=np.full((x.shape[-2], x.shape[-1]), -1)
+
+    im, label = x.copy(), y.copy()
+    label = label.astype(np.float64)
+
+    im1 = im[0]  # red
+    im2 = im[1]  # green
+    im3 = im[2]  # blue
+    im4 = im[3]  # NIR narrow
+    im5 = im[4]  # swir 1
+    im6 = im[5]  # swir 2
+
+    dim = x.shape[-1]
+    label = label.squeeze()
+    norm = transforms.Normalize(means, stds)
+    ims = [torch.stack((transforms.ToTensor()(im1).squeeze(),
+                        transforms.ToTensor()(im2).squeeze(),
+                        transforms.ToTensor()(im3).squeeze(),
+                        transforms.ToTensor()(im4).squeeze(),
+                        transforms.ToTensor()(im5).squeeze(),
+                        transforms.ToTensor()(im6).squeeze()))]
+    ims = [norm(im) for im in ims]
+    ims = torch.stack(ims)
+
+    label = transforms.ToTensor()(label).squeeze()
+
+    _img_metas = {
+        'ori_shape': (dim, dim),
+        'img_shape': (dim, dim),
+        'pad_shape': (dim, dim),
+        'scale_factor': [1., 1., 1., 1.],
+        'flip': False, # needs flip direction specified
+    }
+
+    img_metas = [_img_metas] * 1
+    return {"img": ims,
+            "img_metas": img_metas,
+            "gt_semantic_seg": label}
+
+
+def load_model(config, ckpt):
+
+    print('Loading configuration...')
+    cfg = Config.fromfile(config)
+    print('Building model...')
+    model = build_segmentor(cfg.model, test_cfg=cfg.get('test_cfg'))
+    print('Loading checkpoint...')
+    checkpoint = load_checkpoint(model,ckpt, map_location='cpu')
+    print('Evaluating model...')
+    model = MMDataParallel(model, device_ids=[0])
+    model.eval()
+
+    return model
+
+
+def inference_on_file(model, target_image, output_image, means, stds):
+
+    try:
+        st = time.time()
+        data_orig = open_tiff(target_image)
+        meta = get_meta(target_image)
+        nodata = meta['nodata'] if meta['nodata'] is not None else -9999
+
+        data = preprocess_image(data_orig, means, stds, nodata)
+
+        small_fixed_size_arrs =  custom.split_and_pad(data['img'][:,:,None,:,:], (1, 6, 1, 224, 224))
+        single_chip_batch = [torch.vstack([torch.tensor(t) for t in small_fixed_size_arrs])]
+        print('Running inference...')
+        with torch.no_grad():
+            result = model(single_chip_batch, data['img_metas'], return_loss=False, rescale=False)
+            print("Result: Unique Values: ",np.unique(result))
+
+        print("Output has shape: " + str(result[0].shape))
+        #### TO DO: Post process (e.g. morphological operations)
+
+        result = custom.merge_and_unpad(result, (data_orig.shape[-2],data_orig.shape[-1]), (224, 224))
+
+        print("Result: Unique Values: ",np.unique(result))
+
+        ##### Save file to disk
+        meta["count"] = 1
+        meta["dtype"] = "int16"
+        meta["compress"] = "lzw"
+        meta["nodata"] = -1
+        meta["nodata"] = nodata
+        print('Saving output...')
+        # pdb.set_trace()
+        result = np.where(data_orig[0] == nodata, nodata, result)
+
+        write_tiff(result, output_image, meta)
+        et = time.time()
+        print(f'Inference completed in {str(np.round(et - st, 1))} seconds. Output available at: ' + output_image)
+
+    except:
+        print(f'Error on image {target_image} \nContinue to next input')
+
+def main():
+
+    args = parse_args()
+
+    model = load_model(args.config, args.ckpt)
+    image_pattern = "*merged"
+    target_images = glob.glob(os.path.join(args.input, image_pattern + "." + args.input_type))
+
+    print('Identified images to predict on: ' + str(len(target_images)))
+
+    if not os.path.isdir(args.output):
+        os.mkdir(args.output)
+
+    means, stds = custom.calculate_band_statistics(args.input, image_pattern, bands=[0, 1, 2, 3, 4, 5])
+
+    for i, target_image in enumerate(target_images):
+
+        print(f'Working on Image {i}')
+        output_image = os.path.join(args.output,target_image.split("/")[-1].split(f"_{image_pattern[1:]}.")[0]+'_pred.'+args.input_type)
+
+        inference_on_file(model, target_image, output_image, means, stds)
+
+    print("Running metric eval")
+
+    gt_dir = "/home/workdir/hls-foundation/data/burn_scars/validation"
+    pred_dir = args.output
+    avg_dice_score = custom.compute_metrics(gt_dir, pred_dir)
+    print("Average Dice score:", avg_dice_score)
+
+
+if __name__ == "__main__":
+    main()

custom.py

@@ -0,0 +1,191 @@
+# utils.py
+
+import numpy as np
+import glob
+import rasterio
+from torchvision import transforms
+import torch
+import re
+from torchmetrics import Dice
+import os
+
+def calculate_band_statistics(image_directory, image_pattern, bands=[0, 1, 2, 3, 4, 5]):
+    """
+    Calculate the mean and standard deviation of each band in a folder of GeoTIFF files.
+
+    Args:
+        image_directory (str): Directory where the source GeoTIFF files are stored that are passed to model for training.
+        image_pattern (str): Pattern of the GeoTIFF file names that globs files for computing stats.
+        bands (list, optional): List of bands to calculate statistics for. Defaults to [0, 1, 2, 3, 4, 5].
+
+    Raises:
+        Exception: If no images are found in the given directory.
+
+    Returns:
+        tuple: Two lists containing the means and standard deviations of each band.
+    """
+    # Initialize lists to store the means and standard deviations
+    all_means = []
+    all_stds = []
+
+    # Use glob to get a list of all .tif images in the directory
+    all_images = glob.glob(f"{image_directory}/{image_pattern}.tif")
+
+    # Make sure there are images to process
+    if not all_images:
+        raise Exception("No images found")
+
+    # Get the number of bands
+    num_bands = len(bands)
+
+    # Initialize arrays to hold sums and sum of squares for each band
+    band_sums = np.zeros(num_bands)
+    band_sq_sums = np.zeros(num_bands)
+    pixel_counts = np.zeros(num_bands)
+
+    # Iterate over each image
+    for image_file in all_images:
+        with rasterio.open(image_file) as src:
+            # For each band, calculate the sum, square sum, and pixel count
+            for band in bands:
+                data = src.read(band + 1)  # rasterio band index starts from 1
+                band_sums[band] += np.nansum(data)
+                band_sq_sums[band] += np.nansum(data**2)
+                pixel_counts[band] += np.count_nonzero(~np.isnan(data))
+
+    # Calculate means and standard deviations for each band
+    for i in bands:
+        mean = band_sums[i] / pixel_counts[i]
+        std = np.sqrt((band_sq_sums[i] / pixel_counts[i]) - (mean**2))
+        all_means.append(mean)
+        all_stds.append(std)
+
+    return all_means, all_stds
+
+
+def split_and_pad(array, target_shape):
+    """
+    Splits the input array into smaller arrays of the target shape, padding if necessary.
+
+    Args:
+        array (numpy.ndarray): The input array. Must be shape (batch, band, time, height, width)
+        target_shape (tuple): The target shape of the smaller arrays. Must be of shape
+        (batch, band, time, height, width)
+
+    Raises:
+        ValueError: If target shape is larger than the array shape.
+
+    Returns:
+        list[numpy.ndarray]: A list of the smaller arrays.
+    """
+    # Check if the target shape is smaller or equal to the array shape
+    if target_shape[-2:] > array.shape[-2:]:
+        raise ValueError('Target shape must be smaller or equal to the array shape.')
+
+    # Calculate how much padding is needed
+    pad_h = (target_shape[-2] - array.shape[-2] % target_shape[-2]) % target_shape[-2]
+    pad_w = (target_shape[-1] - array.shape[-1] % target_shape[-1]) % target_shape[-1]
+
+    # Apply padding to the array
+    padded_array = np.pad(array, ((0, 0), (0, 0), (0, 0), (0, pad_h), (0, pad_w)))
+
+    # Split the array into smaller arrays of the target shape
+    result = []
+    for i in range(0, padded_array.shape[-2], target_shape[-2]):
+        for j in range(0, padded_array.shape[-1], target_shape[-1]):
+            result.append(padded_array[..., i:i+target_shape[-2], j:j+target_shape[-1]])
+
+    return result
+
+def merge_and_unpad(np_array_list, original_shape, target_shape):
+    """
+    Assembles smaller numpy arrays back into the original larger numpy array, removing padding if necessary.
+
+    Args:
+        np_array_list (list[numpy.ndarray]): The list of smaller numpy arrays derived from split_and_pad.
+        original_shape (tuple): The original shape of the larger numpy array. Must be shape (Height, Width).
+        target_shape (tuple): The target shape of the smaller numpy arrays. Must be shape (Height, Width).
+
+    Returns:
+        numpy.ndarray: The original larger numpy array.
+    """
+    # Calculate how much padding was added
+    pad_h = (target_shape[0] - original_shape[0] % target_shape[0]) % target_shape[0]
+    pad_w = (target_shape[1] - original_shape[1] % target_shape[1]) % target_shape[1]
+
+    # Calculate the shape of the padded larger array
+    padded_shape = (original_shape[0] + pad_h, original_shape[1] + pad_w)
+
+    # Calculate the number of smaller arrays in each dimension
+    num_arrays_h = padded_shape[0] // target_shape[0]
+    num_arrays_w = padded_shape[1] // target_shape[1]
+
+    # Reshape the list of smaller arrays back into the shape of the padded larger array
+    merged_array = np.stack(np_array_list).reshape(num_arrays_h, num_arrays_w, *target_shape)
+
+    # Rearrange the array dimensions
+    merged_array = merged_array.transpose(0, 2, 1, 3).reshape(*padded_shape)
+
+    # Remove the padding
+    unpadded_array = merged_array[:original_shape[0], :original_shape[1]]
+
+    return unpadded_array
+
+def compute_metrics(gt_dir, pred_dir):
+    """
+    Compute the Dice similarity coefficient between the predicted and ground truth images.
+
+    Args:
+        gt_dir (str): Directory where the ground truth images are stored.
+        pred_dir (str): Directory where the predicted images are stored.
+
+    Returns:
+        Tensor: Dice similarity coefficient score.
+    """
+    dice_metric = Dice()
+
+    # find all .tif files in the prediction directory
+    pred_files = glob.glob(os.path.join(pred_dir, "*.tif"))
+
+    # iterate over each prediction file
+    for pred_file in pred_files:
+        # extract the unique_id from the file name
+        unique_id = re.search('HLS\..*\.v1\.4', os.path.basename(pred_file))
+
+        if unique_id is not None:
+            unique_id = unique_id.group()
+
+            # create the unique pattern for the gt directory
+            gt_file_pattern = os.path.join(gt_dir, f"*{unique_id}*mask.tif")
+
+            # glob the file pattern
+            gt_files = glob.glob(gt_file_pattern)
+
+            # if we found a matching gt file
+            if len(gt_files) == 1:
+                gt_file = gt_files[0]
+
+                # read the .tif files
+                with rasterio.open(gt_file) as src:
+                    gt_img = src.read(1) # ground truth image
+
+                with rasterio.open(pred_file) as src:
+                    pred_img = src.read(1) # predicted image
+
+                # make sure the images are binary (values are 0 or 1)
+                gt_img = (gt_img > 0).astype(np.uint8)
+                pred_img = (pred_img > 0).astype(np.uint8)
+
+                # convert numpy arrays to PyTorch tensors
+                gt_img_tensor = torch.from_numpy(gt_img).long().flatten()
+                pred_img_tensor = torch.from_numpy(pred_img).long().flatten()
+
+                # update dice_metric
+                dice_metric.update(pred_img_tensor, gt_img_tensor)
+
+            else:
+                print(f"No matching ground truth file for prediction file {pred_file}.")
+
+    # compute the dice score
+    dice_score = dice_metric.compute()
+    return dice_score

requirements.txt

@@ -0,0 +1,47 @@
+boxsdk==3.6.2
+cityscapesscripts==2.2.1
+codecov
+detail==0.2.2
+docutils==0.16.0
+einops==0.6.0
+flake8
+interrogate
+jupyterlab==4.0.1
+matplotlib==3.5.1
+mmcls>=0.20.1
+mmdet==2.22.0
+model_archiver==1.0.3
+myst-parser
+-e git+https://github.com/gaotongxiao/pytorch_sphinx_theme.git#egg=pytorch_sphinx_theme
+natsort==8.3.1
+numpy==1.21.6
+onnx==1.13.1
+onnxruntime==1.14.1
+onnx2torch
+opencv-python==4.7.0.72
+openmim
+packaging==21.3
+pandas==1.3.5
+pavi==0.0.1
+Pillow==9.4.0
+pip-tools
+prettytable==3.6.0
+pytest==7.1.3
+rasterio==1.3.4
+requests==2.28.2
+scikit-learn
+scipy==1.7.3
+scikit-image
+seaborn==0.12.2
+sphinx==4.0.2
+sphinx_copybutton
+sphinx_markdown_tables
+tensorrt==8.5.3.1
+timm==0.4.12
+torch==1.9.0+cu111
+-f https://download.pytorch.org/whl/torch_stable.html
+torchvision==0.10.0
+torchmetrics
+ts==0.5.1
+xdoctest>=0.10.0
+yapf