canopy module

Canopy height estimation module using Meta's HighResCanopyHeight model.

This module provides canopy height estimation from RGB imagery using DINOv2 backbone with DPT decoder, based on Meta's HighResCanopyHeight research (https://github.com/facebookresearch/HighResCanopyHeight).

Reference

Tolan et al., "Very high resolution canopy height maps from RGB imagery using self-supervised vision transformer and convolutional decoder trained on Aerial Lidar," Remote Sensing of Environment, 2024. https://doi.org/10.1016/j.rse.2023.113888

CanopyHeightEstimation

Estimate canopy height from RGB imagery using Meta's HighResCanopyHeight model.

This class provides canopy height estimation from RGB aerial or satellite imagery using a DINOv2 backbone with DPT decoder architecture. The model predicts canopy height in meters for each pixel.

The model was developed by Meta AI Research (FAIR) and is described in: Tolan et al., "Very high resolution canopy height maps from RGB imagery using self-supervised vision transformer and convolutional decoder trained on Aerial Lidar," Remote Sensing of Environment, 2023.

Attributes:

Name	Type	Description
model_name	str	Name of the model variant being used.
device	str	Device the model is running on.

Example

estimator = CanopyHeightEstimation() estimator.predict("input.tif", output_path="canopy_height.tif")

Source code in geoai/canopy.py

class CanopyHeightEstimation:
    """Estimate canopy height from RGB imagery using Meta's HighResCanopyHeight model.

    This class provides canopy height estimation from RGB aerial or satellite
    imagery using a DINOv2 backbone with DPT decoder architecture. The model
    predicts canopy height in meters for each pixel.

    The model was developed by Meta AI Research (FAIR) and is described in:
    Tolan et al., "Very high resolution canopy height maps from RGB imagery
    using self-supervised vision transformer and convolutional decoder trained
    on Aerial Lidar," Remote Sensing of Environment, 2023.

    Attributes:
        model_name (str): Name of the model variant being used.
        device (str): Device the model is running on.

    Example:
        >>> estimator = CanopyHeightEstimation()
        >>> estimator.predict("input.tif", output_path="canopy_height.tif")
    """

    def __init__(
        self,
        model_name: str = "compressed_SSLhuge",
        checkpoint_path: Optional[str] = None,
        device: Optional[str] = None,
        cache_dir: str = DEFAULT_CACHE_DIR,
    ):
        """Initialize the CanopyHeightEstimation model.

        Args:
            model_name: Model variant to use. Options:
                - "compressed_SSLhuge" (default): Compressed huge model (749M),
                  CPU-friendly. Best general-purpose model.
                - "SSLhuge_satellite": Full huge model (2.9G), GPU required.
                  Best results on satellite imagery.
                - "compressed_SSLlarge": Compressed large model (400M), CPU-friendly.
                - "compressed_SSLhuge_aerial": Compressed huge model fine-tuned on
                  aerial imagery (749M), CPU-friendly.
            checkpoint_path: Optional path to a local checkpoint file.
                If None, the checkpoint will be downloaded automatically.
            device: Device to run inference on ('cpu', 'cuda', 'cuda:0', etc.).
                If None, automatically selects CUDA if available.  Compressed
                (quantized) models are loaded without quantization when placed
                on GPU, which is faster than CPU inference for large images.
            cache_dir: Directory to cache downloaded model checkpoints.
                Defaults to ~/.cache/geoai/canopy/.
        """
        self.model_name = model_name

        if model_name not in MODEL_VARIANTS:
            raise ValueError(
                f"Unknown model variant '{model_name}'. "
                f"Available: {list(MODEL_VARIANTS.keys())}"
            )

        info = MODEL_VARIANTS[model_name]

        # Determine device – prefer CUDA when available for all model variants
        if device is None:
            if torch.cuda.is_available():
                device = "cuda"
            else:
                device = "cpu"

        self.device = device

        # Download or use provided checkpoint
        if checkpoint_path is None:
            checkpoint_path = _download_checkpoint(model_name, cache_dir)
        self.checkpoint_path = checkpoint_path

        # Load model
        print(f"Loading {model_name} model on {device}...")
        self.model = _load_model(model_name, checkpoint_path, device)
        self.model.eval()
        print("Model loaded successfully.")

        # Image normalization (from Meta's inference.py)
        self._norm_mean = [0.420, 0.411, 0.296]
        self._norm_std = [0.213, 0.156, 0.143]

    def _normalize_image(self, img_tensor: torch.Tensor) -> torch.Tensor:
        """Apply normalization to an image tensor.

        Args:
            img_tensor: Image tensor of shape (B, C, H, W) with values in [0, 1].

        Returns:
            Normalized image tensor.
        """
        mean = torch.tensor(self._norm_mean, device=img_tensor.device).view(1, 3, 1, 1)
        std = torch.tensor(self._norm_std, device=img_tensor.device).view(1, 3, 1, 1)
        return (img_tensor - mean) / std

    @staticmethod
    def _make_weight_map(tile_size: int, overlap: int) -> np.ndarray:
        """Create a 2D raised-cosine weight map for smooth tile blending.

        Produces a weight map that is 1.0 in the non-overlapping centre and
        tapers smoothly to 0 at the edges using a cosine ramp over the
        overlap region.  When multiple overlapping tiles are accumulated
        with these weights the result is seamless.

        Args:
            tile_size: Size of each square tile.
            overlap: Number of pixels that overlap between adjacent tiles.

        Returns:
            2D float32 array of shape (tile_size, tile_size).
        """
        if overlap <= 0:
            return np.ones((tile_size, tile_size), dtype=np.float32)

        ramp = np.linspace(0.0, 1.0, overlap, dtype=np.float32)
        ramp = 0.5 * (1.0 - np.cos(np.pi * ramp))  # raised cosine

        w = np.ones(tile_size, dtype=np.float32)
        w[:overlap] = ramp
        w[-overlap:] = ramp[::-1]

        return np.outer(w, w)

    def predict(
        self,
        input_path: str,
        output_path: Optional[str] = None,
        tile_size: int = 256,
        overlap: int = 128,
        batch_size: int = 4,
        scale_factor: float = 10.0,
        **kwargs,
    ) -> np.ndarray:
        """Predict canopy height from a GeoTIFF image.

        Processes the input image in tiles of the specified size, runs
        inference on each tile, and assembles the results into a full
        canopy height map.  Adjacent tiles overlap and are blended using
        raised-cosine weights for seamless output.

        Args:
            input_path: Path to the input GeoTIFF file (RGB imagery).
            output_path: Optional path to save the output canopy height
                map as a GeoTIFF. If None, only returns the numpy array.
            tile_size: Size of tiles for processing (default: 256).
                The model expects 256x256 tiles.
            overlap: Number of pixels of overlap between tiles (default: 128).
                Using overlap with blending weights eliminates tile-boundary
                artefacts.  Higher values (up to tile_size // 2) give
                smoother results at the cost of more computation.
            batch_size: Number of tiles to process at once (default: 4).
                Larger values use more memory but process faster.
            scale_factor: Factor to multiply model output by to get height
                in meters (default: 10.0). This is set by the original model.
            **kwargs: Additional keyword arguments (reserved for future use).

        Returns:
            numpy.ndarray: 2D array of canopy height values in meters.

        Raises:
            FileNotFoundError: If the input file does not exist.
            ValueError: If the input file is not a valid GeoTIFF.

        Example:
            >>> estimator = CanopyHeightEstimation()
            >>> heights = estimator.predict("aerial_image.tif",
            ...                              output_path="canopy_height.tif",
            ...                              overlap=128)
        """
        if not os.path.exists(input_path):
            raise FileNotFoundError(f"Input file not found: {input_path}")

        if overlap < 0 or overlap >= tile_size:
            raise ValueError(
                f"overlap must be >= 0 and < tile_size ({tile_size}), got {overlap}"
            )

        with rasterio.open(input_path) as src:
            # Read the image
            img = src.read()  # (bands, height, width)
            profile = src.profile.copy()
            height, width = src.height, src.width

            # Handle different band counts
            if img.shape[0] >= 3:
                img = img[:3]  # Use first 3 bands (RGB)
            else:
                raise ValueError(
                    f"Input image has {img.shape[0]} bands, but at least 3 (RGB) are required."
                )

            # Convert to float32 in [0, 1]
            if img.dtype == np.uint8:
                img = img.astype(np.float32) / 255.0
            elif img.dtype == np.uint16:
                img = img.astype(np.float32) / 65535.0
            elif img.dtype in (np.float32, np.float64):
                img = img.astype(np.float32)
                # Clip to [0, 1] if needed
                img_max = img.max()
                if img_max > 1.0:
                    img = img / img_max
            else:
                img = img.astype(np.float32)
                img_max = img.max()
                if img_max > 0:
                    img = img / img_max

        # Calculate tile grid
        step = tile_size - overlap
        n_rows = max(1, math.ceil((height - overlap) / step))
        n_cols = max(1, math.ceil((width - overlap) / step))

        # Pad image to fit tile grid
        padded_h = n_rows * step + overlap
        padded_w = n_cols * step + overlap
        padded_img = np.zeros((3, padded_h, padded_w), dtype=np.float32)
        padded_img[:, :height, :width] = img

        # Blending weight map (raised-cosine taper in overlap regions)
        weight_map = self._make_weight_map(tile_size, overlap)

        # Output arrays
        output = np.zeros((padded_h, padded_w), dtype=np.float32)
        count = np.zeros((padded_h, padded_w), dtype=np.float32)

        # Process tiles in batches
        tiles = []
        positions = []
        for row in range(n_rows):
            for col in range(n_cols):
                y = row * step
                x = col * step
                tile = padded_img[:, y : y + tile_size, x : x + tile_size]
                tiles.append(tile)
                positions.append((y, x))

                if len(tiles) == batch_size:
                    self._process_batch(
                        tiles, positions, output, count, scale_factor, weight_map
                    )
                    tiles = []
                    positions = []

        # Process remaining tiles
        if tiles:
            self._process_batch(
                tiles, positions, output, count, scale_factor, weight_map
            )

        # Average overlapping regions
        count[count == 0] = 1
        output = output / count

        # Crop back to original size
        result = output[:height, :width]

        # Save output if requested
        if output_path is not None:
            out_profile = profile.copy()
            out_profile.update(
                dtype="float32",
                count=1,
                compress="lzw",
                nodata=0,
            )
            # Remove alpha band related settings
            if "photometric" in out_profile:
                del out_profile["photometric"]

            with rasterio.open(output_path, "w", **out_profile) as dst:
                dst.write(result, 1)
            print(f"Canopy height map saved to: {output_path}")

        return result

    def _process_batch(
        self,
        tiles: list,
        positions: list,
        output: np.ndarray,
        count: np.ndarray,
        scale_factor: float,
        weight_map: Optional[np.ndarray] = None,
    ):
        """Process a batch of tiles through the model.

        Args:
            tiles: List of tile arrays (C, H, W).
            positions: List of (y, x) positions for each tile.
            output: Output array to accumulate predictions.
            count: Count array for averaging overlapping tiles.
            scale_factor: Scale factor for model output.
            weight_map: Optional 2D weight array for blending overlapping
                tiles.  If None, uniform weights (1.0) are used.
        """
        tile_size = tiles[0].shape[1]
        batch = torch.from_numpy(np.stack(tiles)).to(self.device)
        batch = self._normalize_image(batch)

        with torch.no_grad():
            pred = self.model(batch)
            pred = pred * scale_factor
            pred = pred.relu()

        pred = pred.cpu().numpy()

        if weight_map is None:
            weight_map = np.ones((tile_size, tile_size), dtype=np.float32)

        for i, (y, x) in enumerate(positions):
            output[y : y + tile_size, x : x + tile_size] += pred[i, 0] * weight_map
            count[y : y + tile_size, x : x + tile_size] += weight_map

    def visualize(
        self,
        input_path: str,
        height_map: Optional[np.ndarray] = None,
        output_path: Optional[str] = None,
        figsize: Tuple[int, int] = (16, 6),
        cmap: str = "viridis",
        vmin: float = 0,
        vmax: Optional[float] = None,
        title: Optional[str] = None,
        **kwargs,
    ) -> plt.Figure:
        """Visualize the input image alongside the canopy height map.

        Args:
            input_path: Path to the input GeoTIFF image, or path to the
                canopy height GeoTIFF if height_map is None.
            height_map: Canopy height map array. If None, will attempt to
                read from input_path (assuming it's a height map).
            output_path: Optional path to save the figure.
            figsize: Figure size as (width, height).
            cmap: Colormap for the height map.
            vmin: Minimum value for colormap.
            vmax: Maximum value for colormap. If None, uses the 98th
                percentile of the height values.
            title: Optional title for the figure.
            **kwargs: Additional keyword arguments for matplotlib.

        Returns:
            matplotlib.figure.Figure: The visualization figure.

        Example:
            >>> estimator = CanopyHeightEstimation()
            >>> heights = estimator.predict("input.tif")
            >>> fig = estimator.visualize("input.tif", heights)
        """
        if height_map is None:
            # Assume input_path is the height map
            with rasterio.open(input_path) as src:
                height_map = src.read(1)
            fig, ax = plt.subplots(1, 1, figsize=(figsize[0] // 2, figsize[1]))
            if vmax is None:
                vmax = (
                    np.percentile(height_map[height_map > 0], 98)
                    if np.any(height_map > 0)
                    else 1
                )
            im = ax.imshow(height_map, cmap=cmap, vmin=vmin, vmax=vmax)
            ax.set_title(title or "Canopy Height Map")
            ax.set_xlabel("Pixels")
            ax.set_ylabel("Pixels")
            cbar = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
            cbar.set_label("Height (meters)")
            ax.axis("off")
        else:
            # Show input image and height map side by side
            with rasterio.open(input_path) as src:
                img = src.read()
                if img.shape[0] >= 3:
                    img = img[:3]
                if img.dtype == np.uint8:
                    img_display = np.moveaxis(img, 0, 2)
                elif img.dtype == np.uint16:
                    img_float = img.astype(np.float32)
                    img_max = img_float.max()
                    if img_max > 0:
                        img_float = img_float / img_max
                    img_display = np.moveaxis(img_float, 0, 2)
                else:
                    img_float = img.astype(np.float32)
                    img_max = img_float.max()
                    if img_max > 1.0:
                        img_float = img_float / img_max
                    img_display = np.moveaxis(img_float, 0, 2)

            fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)

            ax1.imshow(img_display)
            ax1.set_title("Input Image")
            ax1.axis("off")

            if vmax is None:
                vmax = (
                    np.percentile(height_map[height_map > 0], 98)
                    if np.any(height_map > 0)
                    else 1
                )
            im = ax2.imshow(height_map, cmap=cmap, vmin=vmin, vmax=vmax)
            ax2.set_title("Canopy Height Map")
            ax2.axis("off")
            cbar = fig.colorbar(im, ax=ax2, fraction=0.046, pad=0.04)
            cbar.set_label("Height (meters)")

        if title:
            fig.suptitle(title, fontsize=14, fontweight="bold")

        plt.tight_layout()

        if output_path:
            fig.savefig(output_path, dpi=150, bbox_inches="tight")
            print(f"Figure saved to: {output_path}")

        return fig

__init__(model_name='compressed_SSLhuge', checkpoint_path=None, device=None, cache_dir=DEFAULT_CACHE_DIR)

Initialize the CanopyHeightEstimation model.

Parameters:

Name	Type	Description	Default
model_name	str	Model variant to use. Options: - "compressed_SSLhuge" (default): Compressed huge model (749M), CPU-friendly. Best general-purpose model. - "SSLhuge_satellite": Full huge model (2.9G), GPU required. Best results on satellite imagery. - "compressed_SSLlarge": Compressed large model (400M), CPU-friendly. - "compressed_SSLhuge_aerial": Compressed huge model fine-tuned on aerial imagery (749M), CPU-friendly.	'compressed_SSLhuge'
checkpoint_path	Optional[str]	Optional path to a local checkpoint file. If None, the checkpoint will be downloaded automatically.	None
device	Optional[str]	Device to run inference on ('cpu', 'cuda', 'cuda:0', etc.). If None, automatically selects CUDA if available. Compressed (quantized) models are loaded without quantization when placed on GPU, which is faster than CPU inference for large images.	None
cache_dir	str	Directory to cache downloaded model checkpoints. Defaults to ~/.cache/geoai/canopy/.	DEFAULT_CACHE_DIR

Source code in geoai/canopy.py

def __init__(
    self,
    model_name: str = "compressed_SSLhuge",
    checkpoint_path: Optional[str] = None,
    device: Optional[str] = None,
    cache_dir: str = DEFAULT_CACHE_DIR,
):
    """Initialize the CanopyHeightEstimation model.

    Args:
        model_name: Model variant to use. Options:
            - "compressed_SSLhuge" (default): Compressed huge model (749M),
              CPU-friendly. Best general-purpose model.
            - "SSLhuge_satellite": Full huge model (2.9G), GPU required.
              Best results on satellite imagery.
            - "compressed_SSLlarge": Compressed large model (400M), CPU-friendly.
            - "compressed_SSLhuge_aerial": Compressed huge model fine-tuned on
              aerial imagery (749M), CPU-friendly.
        checkpoint_path: Optional path to a local checkpoint file.
            If None, the checkpoint will be downloaded automatically.
        device: Device to run inference on ('cpu', 'cuda', 'cuda:0', etc.).
            If None, automatically selects CUDA if available.  Compressed
            (quantized) models are loaded without quantization when placed
            on GPU, which is faster than CPU inference for large images.
        cache_dir: Directory to cache downloaded model checkpoints.
            Defaults to ~/.cache/geoai/canopy/.
    """
    self.model_name = model_name

    if model_name not in MODEL_VARIANTS:
        raise ValueError(
            f"Unknown model variant '{model_name}'. "
            f"Available: {list(MODEL_VARIANTS.keys())}"
        )

    info = MODEL_VARIANTS[model_name]

    # Determine device – prefer CUDA when available for all model variants
    if device is None:
        if torch.cuda.is_available():
            device = "cuda"
        else:
            device = "cpu"

    self.device = device

    # Download or use provided checkpoint
    if checkpoint_path is None:
        checkpoint_path = _download_checkpoint(model_name, cache_dir)
    self.checkpoint_path = checkpoint_path

    # Load model
    print(f"Loading {model_name} model on {device}...")
    self.model = _load_model(model_name, checkpoint_path, device)
    self.model.eval()
    print("Model loaded successfully.")

    # Image normalization (from Meta's inference.py)
    self._norm_mean = [0.420, 0.411, 0.296]
    self._norm_std = [0.213, 0.156, 0.143]

predict(input_path, output_path=None, tile_size=256, overlap=128, batch_size=4, scale_factor=10.0, **kwargs)

Predict canopy height from a GeoTIFF image.

Processes the input image in tiles of the specified size, runs inference on each tile, and assembles the results into a full canopy height map. Adjacent tiles overlap and are blended using raised-cosine weights for seamless output.

Parameters:

Name	Type	Description	Default
input_path	str	Path to the input GeoTIFF file (RGB imagery).	required
output_path	Optional[str]	Optional path to save the output canopy height map as a GeoTIFF. If None, only returns the numpy array.	None
tile_size	int	Size of tiles for processing (default: 256). The model expects 256x256 tiles.	256
overlap	int	Number of pixels of overlap between tiles (default: 128). Using overlap with blending weights eliminates tile-boundary artefacts. Higher values (up to tile_size // 2) give smoother results at the cost of more computation.	128
batch_size	int	Number of tiles to process at once (default: 4). Larger values use more memory but process faster.	4
scale_factor	float	Factor to multiply model output by to get height in meters (default: 10.0). This is set by the original model.	10.0
**kwargs		Additional keyword arguments (reserved for future use).	{}

Returns:

Type	Description
ndarray	numpy.ndarray: 2D array of canopy height values in meters.

Raises:

Type	Description
FileNotFoundError	If the input file does not exist.
ValueError	If the input file is not a valid GeoTIFF.

Example

estimator = CanopyHeightEstimation() heights = estimator.predict("aerial_image.tif", ... output_path="canopy_height.tif", ... overlap=128)

Source code in geoai/canopy.py

def predict(
    self,
    input_path: str,
    output_path: Optional[str] = None,
    tile_size: int = 256,
    overlap: int = 128,
    batch_size: int = 4,
    scale_factor: float = 10.0,
    **kwargs,
) -> np.ndarray:
    """Predict canopy height from a GeoTIFF image.

    Processes the input image in tiles of the specified size, runs
    inference on each tile, and assembles the results into a full
    canopy height map.  Adjacent tiles overlap and are blended using
    raised-cosine weights for seamless output.

    Args:
        input_path: Path to the input GeoTIFF file (RGB imagery).
        output_path: Optional path to save the output canopy height
            map as a GeoTIFF. If None, only returns the numpy array.
        tile_size: Size of tiles for processing (default: 256).
            The model expects 256x256 tiles.
        overlap: Number of pixels of overlap between tiles (default: 128).
            Using overlap with blending weights eliminates tile-boundary
            artefacts.  Higher values (up to tile_size // 2) give
            smoother results at the cost of more computation.
        batch_size: Number of tiles to process at once (default: 4).
            Larger values use more memory but process faster.
        scale_factor: Factor to multiply model output by to get height
            in meters (default: 10.0). This is set by the original model.
        **kwargs: Additional keyword arguments (reserved for future use).

    Returns:
        numpy.ndarray: 2D array of canopy height values in meters.

    Raises:
        FileNotFoundError: If the input file does not exist.
        ValueError: If the input file is not a valid GeoTIFF.

    Example:
        >>> estimator = CanopyHeightEstimation()
        >>> heights = estimator.predict("aerial_image.tif",
        ...                              output_path="canopy_height.tif",
        ...                              overlap=128)
    """
    if not os.path.exists(input_path):
        raise FileNotFoundError(f"Input file not found: {input_path}")

    if overlap < 0 or overlap >= tile_size:
        raise ValueError(
            f"overlap must be >= 0 and < tile_size ({tile_size}), got {overlap}"
        )

    with rasterio.open(input_path) as src:
        # Read the image
        img = src.read()  # (bands, height, width)
        profile = src.profile.copy()
        height, width = src.height, src.width

        # Handle different band counts
        if img.shape[0] >= 3:
            img = img[:3]  # Use first 3 bands (RGB)
        else:
            raise ValueError(
                f"Input image has {img.shape[0]} bands, but at least 3 (RGB) are required."
            )

        # Convert to float32 in [0, 1]
        if img.dtype == np.uint8:
            img = img.astype(np.float32) / 255.0
        elif img.dtype == np.uint16:
            img = img.astype(np.float32) / 65535.0
        elif img.dtype in (np.float32, np.float64):
            img = img.astype(np.float32)
            # Clip to [0, 1] if needed
            img_max = img.max()
            if img_max > 1.0:
                img = img / img_max
        else:
            img = img.astype(np.float32)
            img_max = img.max()
            if img_max > 0:
                img = img / img_max

    # Calculate tile grid
    step = tile_size - overlap
    n_rows = max(1, math.ceil((height - overlap) / step))
    n_cols = max(1, math.ceil((width - overlap) / step))

    # Pad image to fit tile grid
    padded_h = n_rows * step + overlap
    padded_w = n_cols * step + overlap
    padded_img = np.zeros((3, padded_h, padded_w), dtype=np.float32)
    padded_img[:, :height, :width] = img

    # Blending weight map (raised-cosine taper in overlap regions)
    weight_map = self._make_weight_map(tile_size, overlap)

    # Output arrays
    output = np.zeros((padded_h, padded_w), dtype=np.float32)
    count = np.zeros((padded_h, padded_w), dtype=np.float32)

    # Process tiles in batches
    tiles = []
    positions = []
    for row in range(n_rows):
        for col in range(n_cols):
            y = row * step
            x = col * step
            tile = padded_img[:, y : y + tile_size, x : x + tile_size]
            tiles.append(tile)
            positions.append((y, x))

            if len(tiles) == batch_size:
                self._process_batch(
                    tiles, positions, output, count, scale_factor, weight_map
                )
                tiles = []
                positions = []

    # Process remaining tiles
    if tiles:
        self._process_batch(
            tiles, positions, output, count, scale_factor, weight_map
        )

    # Average overlapping regions
    count[count == 0] = 1
    output = output / count

    # Crop back to original size
    result = output[:height, :width]

    # Save output if requested
    if output_path is not None:
        out_profile = profile.copy()
        out_profile.update(
            dtype="float32",
            count=1,
            compress="lzw",
            nodata=0,
        )
        # Remove alpha band related settings
        if "photometric" in out_profile:
            del out_profile["photometric"]

        with rasterio.open(output_path, "w", **out_profile) as dst:
            dst.write(result, 1)
        print(f"Canopy height map saved to: {output_path}")

    return result

visualize(input_path, height_map=None, output_path=None, figsize=(16, 6), cmap='viridis', vmin=0, vmax=None, title=None, **kwargs)

Visualize the input image alongside the canopy height map.

Parameters:

Name	Type	Description	Default
input_path	str	Path to the input GeoTIFF image, or path to the canopy height GeoTIFF if height_map is None.	required
height_map	Optional[ndarray]	Canopy height map array. If None, will attempt to read from input_path (assuming it's a height map).	None
output_path	Optional[str]	Optional path to save the figure.	None
figsize	Tuple[int, int]	Figure size as (width, height).	(16, 6)
cmap	str	Colormap for the height map.	'viridis'
vmin	float	Minimum value for colormap.	0
vmax	Optional[float]	Maximum value for colormap. If None, uses the 98th percentile of the height values.	None
title	Optional[str]	Optional title for the figure.	None
**kwargs		Additional keyword arguments for matplotlib.	{}

Returns:

Type	Description
Figure	matplotlib.figure.Figure: The visualization figure.

Example

estimator = CanopyHeightEstimation() heights = estimator.predict("input.tif") fig = estimator.visualize("input.tif", heights)

Source code in geoai/canopy.py

def visualize(
    self,
    input_path: str,
    height_map: Optional[np.ndarray] = None,
    output_path: Optional[str] = None,
    figsize: Tuple[int, int] = (16, 6),
    cmap: str = "viridis",
    vmin: float = 0,
    vmax: Optional[float] = None,
    title: Optional[str] = None,
    **kwargs,
) -> plt.Figure:
    """Visualize the input image alongside the canopy height map.

    Args:
        input_path: Path to the input GeoTIFF image, or path to the
            canopy height GeoTIFF if height_map is None.
        height_map: Canopy height map array. If None, will attempt to
            read from input_path (assuming it's a height map).
        output_path: Optional path to save the figure.
        figsize: Figure size as (width, height).
        cmap: Colormap for the height map.
        vmin: Minimum value for colormap.
        vmax: Maximum value for colormap. If None, uses the 98th
            percentile of the height values.
        title: Optional title for the figure.
        **kwargs: Additional keyword arguments for matplotlib.

    Returns:
        matplotlib.figure.Figure: The visualization figure.

    Example:
        >>> estimator = CanopyHeightEstimation()
        >>> heights = estimator.predict("input.tif")
        >>> fig = estimator.visualize("input.tif", heights)
    """
    if height_map is None:
        # Assume input_path is the height map
        with rasterio.open(input_path) as src:
            height_map = src.read(1)
        fig, ax = plt.subplots(1, 1, figsize=(figsize[0] // 2, figsize[1]))
        if vmax is None:
            vmax = (
                np.percentile(height_map[height_map > 0], 98)
                if np.any(height_map > 0)
                else 1
            )
        im = ax.imshow(height_map, cmap=cmap, vmin=vmin, vmax=vmax)
        ax.set_title(title or "Canopy Height Map")
        ax.set_xlabel("Pixels")
        ax.set_ylabel("Pixels")
        cbar = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
        cbar.set_label("Height (meters)")
        ax.axis("off")
    else:
        # Show input image and height map side by side
        with rasterio.open(input_path) as src:
            img = src.read()
            if img.shape[0] >= 3:
                img = img[:3]
            if img.dtype == np.uint8:
                img_display = np.moveaxis(img, 0, 2)
            elif img.dtype == np.uint16:
                img_float = img.astype(np.float32)
                img_max = img_float.max()
                if img_max > 0:
                    img_float = img_float / img_max
                img_display = np.moveaxis(img_float, 0, 2)
            else:
                img_float = img.astype(np.float32)
                img_max = img_float.max()
                if img_max > 1.0:
                    img_float = img_float / img_max
                img_display = np.moveaxis(img_float, 0, 2)

        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)

        ax1.imshow(img_display)
        ax1.set_title("Input Image")
        ax1.axis("off")

        if vmax is None:
            vmax = (
                np.percentile(height_map[height_map > 0], 98)
                if np.any(height_map > 0)
                else 1
            )
        im = ax2.imshow(height_map, cmap=cmap, vmin=vmin, vmax=vmax)
        ax2.set_title("Canopy Height Map")
        ax2.axis("off")
        cbar = fig.colorbar(im, ax=ax2, fraction=0.046, pad=0.04)
        cbar.set_label("Height (meters)")

    if title:
        fig.suptitle(title, fontsize=14, fontweight="bold")

    plt.tight_layout()

    if output_path:
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
        print(f"Figure saved to: {output_path}")

    return fig

canopy_height_estimation(input_path, output_path=None, model_name='compressed_SSLhuge', checkpoint_path=None, device=None, tile_size=256, overlap=128, batch_size=4, cache_dir=DEFAULT_CACHE_DIR, **kwargs)

Convenience function for canopy height estimation.

Creates a CanopyHeightEstimation instance and runs prediction in one call.

Parameters:

Name	Type	Description	Default
input_path	str	Path to the input GeoTIFF file (RGB imagery).	required
output_path	Optional[str]	Optional path to save the output as a GeoTIFF.	None
model_name	str	Model variant to use. Default: "compressed_SSLhuge". See CanopyHeightEstimation for available options.	'compressed_SSLhuge'
checkpoint_path	Optional[str]	Optional path to a local checkpoint file.	None
device	Optional[str]	Device to run inference on. If None, auto-selects.	None
tile_size	int	Size of tiles for processing (default: 256).	256
overlap	int	Overlap between tiles in pixels (default: 128). Higher values give smoother results.	128
batch_size	int	Number of tiles per batch (default: 4).	4
cache_dir	str	Directory to cache model checkpoints.	DEFAULT_CACHE_DIR
**kwargs		Additional keyword arguments passed to predict().	{}

Returns:

Type	Description
ndarray	numpy.ndarray: 2D array of canopy height values in meters.

Example

heights = canopy_height_estimation( ... "aerial_image.tif", ... output_path="canopy_height.tif", ... model_name="compressed_SSLhuge", ... )

Source code in geoai/canopy.py

def canopy_height_estimation(
    input_path: str,
    output_path: Optional[str] = None,
    model_name: str = "compressed_SSLhuge",
    checkpoint_path: Optional[str] = None,
    device: Optional[str] = None,
    tile_size: int = 256,
    overlap: int = 128,
    batch_size: int = 4,
    cache_dir: str = DEFAULT_CACHE_DIR,
    **kwargs,
) -> np.ndarray:
    """Convenience function for canopy height estimation.

    Creates a CanopyHeightEstimation instance and runs prediction in one call.

    Args:
        input_path: Path to the input GeoTIFF file (RGB imagery).
        output_path: Optional path to save the output as a GeoTIFF.
        model_name: Model variant to use. Default: "compressed_SSLhuge".
            See CanopyHeightEstimation for available options.
        checkpoint_path: Optional path to a local checkpoint file.
        device: Device to run inference on. If None, auto-selects.
        tile_size: Size of tiles for processing (default: 256).
        overlap: Overlap between tiles in pixels (default: 128).
            Higher values give smoother results.
        batch_size: Number of tiles per batch (default: 4).
        cache_dir: Directory to cache model checkpoints.
        **kwargs: Additional keyword arguments passed to predict().

    Returns:
        numpy.ndarray: 2D array of canopy height values in meters.

    Example:
        >>> heights = canopy_height_estimation(
        ...     "aerial_image.tif",
        ...     output_path="canopy_height.tif",
        ...     model_name="compressed_SSLhuge",
        ... )
    """
    estimator = CanopyHeightEstimation(
        model_name=model_name,
        checkpoint_path=checkpoint_path,
        device=device,
        cache_dir=cache_dir,
    )
    return estimator.predict(
        input_path=input_path,
        output_path=output_path,
        tile_size=tile_size,
        overlap=overlap,
        batch_size=batch_size,
        **kwargs,
    )

list_canopy_models()

List available canopy height estimation model variants.

Returns:

Type	Description
Dict[str, str]	Dictionary mapping model names to their descriptions.

Example

models = list_canopy_models() for name, desc in models.items(): ... print(f"{name}: {desc}")

Source code in geoai/canopy.py

def list_canopy_models() -> Dict[str, str]:
    """List available canopy height estimation model variants.

    Returns:
        Dictionary mapping model names to their descriptions.

    Example:
        >>> models = list_canopy_models()
        >>> for name, desc in models.items():
        ...     print(f"{name}: {desc}")
    """
    return {name: info["description"] for name, info in MODEL_VARIANTS.items()}