Table of Contents
Shortcuts

# Source code for mmseg.datasets.pipelines.misc_utils

    from pathlib import Path
    from osgeo import gdal
    import numpy as np
    import json
    import os
    from tqdm import tqdm
    from PIL import Image

    UNITS_PER_METER_CONVERSION_FACTORS = {"cm": 100.0, "m": 1.0}

    def save_image(img_path, img):
        """Note this function does not utilize conversion factors to meters, so implicitly assumes units are in meters."""
        rows = img.shape[0]
        cols = img.shape[1]
        bands = 1
        if np.ndim(img) == 3:
            bands = img.shape[2]
        if img.dtype == np.uint8:
            data_type = gdal.GDT_Byte
        else:
            data_type = gdal.GDT_Float32
        driver = gdal.GetDriverByName("GTiff").Create(
            str(img_path), rows, cols, bands, data_type, ["COMPRESS=LZW"]
        )
        if bands == 1:
            driver.GetRasterBand(1).WriteArray(img[:, :])
        else:
            for i in range(bands):
                driver.GetRasterBand(i + 1).WriteArray(img[:, :, i])
        driver.FlushCache()
        driver = None

    def load_image(
        image_path,
        args,
        dtype_out="float32",
        units_per_meter_conversion_factors=UNITS_PER_METER_CONVERSION_FACTORS,
    ):

        image_path = Path(image_path)
        if not image_path.exists():
            return None
        image = gdal.Open(str(image_path))
        image = image.ReadAsArray()

        # convert AGL units and fill nan placeholder with nan
        if "AGL" in image_path.name:
            image = image.astype(dtype_out)
            np.putmask(image, image == args.nan_placeholder, np.nan)
            # e.g., (cm) / (cm / m) = m
            units_per_meter = units_per_meter_conversion_factors[args.unit]
            image = (image / units_per_meter).astype(dtype_out)

        # transpose if RGB
        if len(image.shape) == 3:
            image = np.transpose(image, [1, 2, 0])

        return image

    def load_vflow(
        vflow_path,
        agl,
        args,
        dtype_out="float32",
        units_per_meter_conversion_factors=UNITS_PER_METER_CONVERSION_FACTORS,
        return_vflow_pred_mat=False,
    ):

        vflow_path = Path(vflow_path)
        vflow_data = json.load(vflow_path.open("r"))

        # e.g., (pixels / cm) * (cm / m) = (pixels / m)
        units_per_meter = units_per_meter_conversion_factors[args.unit]
        vflow_data["scale"] = vflow_data["scale"] * units_per_meter

        xdir, ydir = np.sin(vflow_data["angle"]), np.cos(vflow_data["angle"])
        mag = agl * vflow_data["scale"]

        vflow_items = [mag.astype(dtype_out), xdir.astype(dtype_out), ydir.astype(dtype_out), vflow_data]

        if return_vflow_pred_mat:
            vflow = np.zeros((agl.shape[0],agl.shape[1],2))
            vflow[:,:,0] = mag * xdir
            vflow[:,:,1] = mag * ydir
            vflow_items.insert(0, vflow)

        return vflow_items

    def get_r2(error_sum, gt_sq_sum, data_sum, count):
        return 1 - error_sum / (gt_sq_sum - (data_sum ** 2) / count)

    def get_rms(errors):
        return np.sqrt(np.mean(np.square(errors)))

    def get_angle_error(dir_pred, dir_gt):
        dir_pred /= np.linalg.norm(dir_pred)
        dir_gt /= np.linalg.norm(dir_gt)
        cos_ang = np.dot(dir_pred, dir_gt)
        sin_ang = np.linalg.norm(np.cross(dir_pred, dir_gt))
        rad_diff = np.arctan2(sin_ang, cos_ang)
        angle_error = np.degrees(rad_diff)
        return angle_error

    def get_r2_info(data_gt, data_pred):

        data_pred = np.squeeze(data_pred)

        not_nan = ~np.isnan(data_gt)
        count = np.sum(not_nan)
        diff = data_pred - data_gt

        error_sum = np.sum(np.square(diff[not_nan]))
        rms = get_rms(diff[not_nan])
        data_sum = np.sum(data_gt[not_nan])
        gt_sq_sum = np.sum(np.square(data_gt[not_nan]))

        return count, error_sum, rms, data_sum, gt_sq_sum

    def convert_and_compress_prediction_dir(
        predictions_dir,
        to_unit="cm",
        agl_dtype="uint16",
        compression_type="tiff_adobe_deflate",
        conversion_factors=UNITS_PER_METER_CONVERSION_FACTORS,
    ):
        """Convert and compress prediction directory.
        Default parameters are consistent with DrivenData platform submission requirements.
        """
        predictions_dir = Path(predictions_dir)
        converted_predictions_dir = predictions_dir.with_name(
            f"{predictions_dir.name}_converted_{to_unit}_compressed_{agl_dtype}"
        )
        converted_predictions_dir.mkdir(exist_ok=True, parents=True)

        conversion_factor = conversion_factors[to_unit]

        agl_paths = list(predictions_dir.glob("*_AGL.tif"))
        json_paths = list(
            pth.with_name(pth.name.replace("_AGL", "_VFLOW")).with_suffix(".json")
            for pth in agl_paths
        )
        for agl_path, json_path in tqdm(zip(agl_paths, json_paths), total=len(agl_paths)):
            # convert and compress agl tif
            imarray = np.array(Image.open(agl_path))
            imarray = np.round(imarray * conversion_factor).astype(agl_dtype)
            new_image = Image.fromarray(imarray)
            new_image_path = converted_predictions_dir / agl_path.name
            new_image.save(str(new_image_path), "TIFF", compression=compression_type)

            # convert and compress vflow json
            vflow = json.load(json_path.open("r"))
            vflow["scale"] = vflow["scale"] / conversion_factor
            new_json_path = converted_predictions_dir / json_path.name
            json.dump(vflow, new_json_path.open("w"))