geoai module

Main module.

AgricultureFieldDelineator

Bases: ObjectDetector

Agricultural field boundary delineation using a pre-trained Mask R-CNN model.

This class extends the ObjectDetector class to specifically handle Sentinel-2 imagery with 12 spectral bands for agricultural field boundary detection.

Attributes:

Name	Type	Description
band_selection		List of band indices to use for prediction (default: RGB)
sentinel_band_stats		Per-band statistics for Sentinel-2 data
use_ndvi		Whether to calculate and include NDVI as an additional channel

Source code in geoai/extract.py

class AgricultureFieldDelineator(ObjectDetector):
    """
    Agricultural field boundary delineation using a pre-trained Mask R-CNN model.

    This class extends the ObjectDetector class to specifically handle Sentinel-2
    imagery with 12 spectral bands for agricultural field boundary detection.

    Attributes:
        band_selection: List of band indices to use for prediction (default: RGB)
        sentinel_band_stats: Per-band statistics for Sentinel-2 data
        use_ndvi: Whether to calculate and include NDVI as an additional channel
    """

    def __init__(
        self,
        model_path: str = "field_boundary_detector.pth",
        repo_id: Optional[str] = None,
        model: Optional[Any] = None,
        device: Optional[str] = None,
        band_selection: Optional[List[int]] = None,
        use_ndvi: bool = False,
    ) -> None:
        """
        Initialize the field boundary delineator.

        Args:
            model_path: Path to the .pth model file.
            repo_id: Repo ID for loading models from the Hub.
            model: Custom model to use for inference.
            device: Device to use for inference ('cuda:0', 'cpu', etc.).
            band_selection: List of Sentinel-2 band indices to use (None = adapt based on model)
            use_ndvi: Whether to calculate and include NDVI as an additional channel
        """
        # Save parameters before calling parent constructor
        self.custom_band_selection = band_selection
        self.use_ndvi = use_ndvi

        # Set device (copied from parent init to ensure it's set before initialize_model)
        if device is None:
            self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        else:
            self.device = torch.device(device)

        # Initialize model differently for multi-spectral input
        model = self.initialize_sentinel2_model(model)

        # Call parent but with our custom model
        super().__init__(
            model_path=model_path, repo_id=repo_id, model=model, device=device
        )

        # Default Sentinel-2 band statistics (can be overridden with actual stats)
        # Band order: [B1, B2, B3, B4, B5, B6, B7, B8, B8A, B9, B10, B11, B12]
        self.sentinel_band_stats = {
            "means": [
                0.0975,
                0.0476,
                0.0598,
                0.0697,
                0.1077,
                0.1859,
                0.2378,
                0.2061,
                0.2598,
                0.4120,
                0.1956,
                0.1410,
            ],
            "stds": [
                0.0551,
                0.0290,
                0.0298,
                0.0479,
                0.0506,
                0.0505,
                0.0747,
                0.0642,
                0.0782,
                0.1187,
                0.0651,
                0.0679,
            ],
        }

        # Set default band selection (RGB - typically B4, B3, B2 for Sentinel-2)
        self.band_selection = (
            self.custom_band_selection
            if self.custom_band_selection is not None
            else [3, 2, 1]
        )  # R, G, B bands

        # Customize parameters for field delineation
        self.confidence_threshold = 0.5  # Default confidence threshold
        self.overlap = 0.5  # Higher overlap for field boundary detection
        self.min_object_area = 1000  # Minimum area in pixels for field detection
        self.simplify_tolerance = 2.0  # Higher tolerance for field boundaries

    def initialize_sentinel2_model(self, model: Optional[Any] = None) -> Any:
        """
        Initialize a Mask R-CNN model with a modified first layer to accept Sentinel-2 data.

        Args:
            model: Pre-initialized model (optional)

        Returns:
            Modified model with appropriate input channels
        """
        import torchvision
        from torchvision.models.detection import maskrcnn_resnet50_fpn
        from torchvision.models.detection.backbone_utils import resnet_fpn_backbone

        if model is not None:
            return model

        # Determine number of input channels based on band selection and NDVI
        num_input_channels = (
            len(self.custom_band_selection)
            if self.custom_band_selection is not None
            else 3
        )
        if self.use_ndvi:
            num_input_channels += 1

        print(f"Initializing Mask R-CNN model with {num_input_channels} input channels")

        # Create a ResNet50 backbone with modified input channels
        backbone = resnet_fpn_backbone("resnet50", weights=None)

        # Replace the first conv layer to accept multi-spectral input
        original_conv = backbone.body.conv1
        backbone.body.conv1 = torch.nn.Conv2d(
            num_input_channels,
            original_conv.out_channels,
            kernel_size=original_conv.kernel_size,
            stride=original_conv.stride,
            padding=original_conv.padding,
            bias=original_conv.bias is not None,
        )

        # Create Mask R-CNN with our modified backbone
        model = maskrcnn_resnet50_fpn(
            backbone=backbone,
            num_classes=2,  # Background + field
            image_mean=[0.485] * num_input_channels,  # Extend mean to all channels
            image_std=[0.229] * num_input_channels,  # Extend std to all channels
        )

        model.to(self.device)
        return model

    def preprocess_sentinel_bands(
        self,
        image_data: np.ndarray,
        band_selection: Optional[List[int]] = None,
        use_ndvi: Optional[bool] = None,
    ) -> torch.Tensor:
        """
        Preprocess Sentinel-2 band data for model input.

        Args:
            image_data: Raw Sentinel-2 image data as numpy array [bands, height, width]
            band_selection: List of band indices to use (overrides instance default if provided)
            use_ndvi: Whether to include NDVI (overrides instance default if provided)

        Returns:
            Processed tensor ready for model input
        """
        # Use instance defaults if not specified
        band_selection = (
            band_selection if band_selection is not None else self.band_selection
        )
        use_ndvi = use_ndvi if use_ndvi is not None else self.use_ndvi

        # Select bands
        selected_bands = image_data[band_selection]

        # Calculate NDVI if requested (using B8 and B4 which are indices 7 and 3)
        if (
            use_ndvi
            and 7 in range(image_data.shape[0])
            and 3 in range(image_data.shape[0])
        ):
            nir = image_data[7].astype(np.float32)  # B8 (NIR)
            red = image_data[3].astype(np.float32)  # B4 (Red)

            # Avoid division by zero
            denominator = nir + red
            ndvi = np.zeros_like(nir)
            valid_mask = denominator > 0
            ndvi[valid_mask] = (nir[valid_mask] - red[valid_mask]) / denominator[
                valid_mask
            ]

            # Rescale NDVI from [-1, 1] to [0, 1]
            ndvi = (ndvi + 1) / 2

            # Add NDVI as an additional channel
            selected_bands = np.vstack([selected_bands, ndvi[np.newaxis, :, :]])

        # Convert to tensor
        image_tensor = torch.from_numpy(selected_bands).float()

        # Normalize using band statistics
        for i, band_idx in enumerate(band_selection):
            # Make sure band_idx is within range of our statistics
            if band_idx < len(self.sentinel_band_stats["means"]):
                mean = self.sentinel_band_stats["means"][band_idx]
                std = self.sentinel_band_stats["stds"][band_idx]
                image_tensor[i] = (image_tensor[i] - mean) / std

        # If NDVI was added, normalize it too (last channel)
        if use_ndvi:
            # NDVI is already roughly in [0,1] range, just standardize it slightly
            image_tensor[-1] = (image_tensor[-1] - 0.5) / 0.5

        return image_tensor

    def update_band_stats(
        self,
        raster_path: str,
        band_selection: Optional[List[int]] = None,
        sample_size: int = 1000,
    ) -> Dict[str, List[float]]:
        """
        Update band statistics from the input Sentinel-2 raster.

        Args:
            raster_path: Path to the Sentinel-2 raster file
            band_selection: Specific bands to update (None = update all available)
            sample_size: Number of random pixels to sample for statistics calculation

        Returns:
            Updated band statistics dictionary
        """
        with rasterio.open(raster_path) as src:
            # Check if this is likely a Sentinel-2 product
            band_count = src.count
            if band_count < 3:
                print(
                    f"Warning: Raster has only {band_count} bands, may not be Sentinel-2 data"
                )

            # Get dimensions
            height, width = src.height, src.width

            # Determine which bands to analyze
            if band_selection is None:
                band_selection = list(range(1, band_count + 1))  # 1-indexed

            # Initialize arrays for band statistics
            means = []
            stds = []

            # Sample random pixels
            np.random.seed(42)  # For reproducibility
            sample_rows = np.random.randint(0, height, sample_size)
            sample_cols = np.random.randint(0, width, sample_size)

            # Calculate statistics for each band
            for band in band_selection:
                # Read band data
                band_data = src.read(band)

                # Sample values
                sample_values = band_data[sample_rows, sample_cols]

                # Remove invalid values (e.g., nodata)
                valid_samples = sample_values[np.isfinite(sample_values)]

                # Calculate statistics
                mean = float(np.mean(valid_samples))
                std = float(np.std(valid_samples))

                # Store results
                means.append(mean)
                stds.append(std)

                print(f"Band {band}: mean={mean:.4f}, std={std:.4f}")

            # Update instance variables
            self.sentinel_band_stats = {"means": means, "stds": stds}

            return self.sentinel_band_stats

    def process_sentinel_raster(
        self,
        raster_path: str,
        output_path: Optional[str] = None,
        batch_size: int = 4,
        band_selection: Optional[List[int]] = None,
        use_ndvi: Optional[bool] = None,
        filter_edges: bool = True,
        edge_buffer: int = 20,
        **kwargs: Any,
    ) -> gpd.GeoDataFrame:
        """
        Process a Sentinel-2 raster to extract field boundaries.

        Args:
            raster_path: Path to Sentinel-2 raster file
            output_path: Path to output GeoJSON or Parquet file (optional)
            batch_size: Batch size for processing
            band_selection: List of bands to use (None = use instance default)
            use_ndvi: Whether to include NDVI (None = use instance default)
            filter_edges: Whether to filter out objects at the edges of the image
            edge_buffer: Size of edge buffer in pixels to filter out objects
            **kwargs: Additional parameters for processing

        Returns:
            GeoDataFrame with field boundaries
        """
        # Use instance defaults if not specified
        band_selection = (
            band_selection if band_selection is not None else self.band_selection
        )
        use_ndvi = use_ndvi if use_ndvi is not None else self.use_ndvi

        # Get parameters from kwargs or use instance defaults
        confidence_threshold = kwargs.get(
            "confidence_threshold", self.confidence_threshold
        )
        overlap = kwargs.get("overlap", self.overlap)
        chip_size = kwargs.get("chip_size", self.chip_size)
        nms_iou_threshold = kwargs.get("nms_iou_threshold", self.nms_iou_threshold)
        mask_threshold = kwargs.get("mask_threshold", self.mask_threshold)
        min_object_area = kwargs.get("min_object_area", self.min_object_area)
        simplify_tolerance = kwargs.get("simplify_tolerance", self.simplify_tolerance)

        # Update band statistics if not already done
        if kwargs.get("update_stats", True):
            self.update_band_stats(raster_path, band_selection)

        print(f"Processing with parameters:")
        print(f"- Using bands: {band_selection}")
        print(f"- Include NDVI: {use_ndvi}")
        print(f"- Confidence threshold: {confidence_threshold}")
        print(f"- Tile overlap: {overlap}")
        print(f"- Chip size: {chip_size}")
        print(f"- Filter edge objects: {filter_edges}")

        # Create a custom Sentinel-2 dataset class
        class Sentinel2Dataset(torch.utils.data.Dataset):
            def __init__(
                self,
                raster_path: str,
                chip_size: Tuple[int, int],
                stride_x: int,
                stride_y: int,
                band_selection: List[int],
                use_ndvi: bool,
                field_delineator: Any,
            ) -> None:
                self.raster_path = raster_path
                self.chip_size = chip_size
                self.stride_x = stride_x
                self.stride_y = stride_y
                self.band_selection = band_selection
                self.use_ndvi = use_ndvi
                self.field_delineator = field_delineator

                with rasterio.open(self.raster_path) as src:
                    self.height = src.height
                    self.width = src.width
                    self.count = src.count
                    self.crs = src.crs
                    self.transform = src.transform

                    # Calculate row_starts and col_starts
                    self.row_starts = []
                    self.col_starts = []

                    # Normal row starts using stride
                    for r in range((self.height - 1) // self.stride_y):
                        self.row_starts.append(r * self.stride_y)

                    # Add a special last row that ensures we reach the bottom edge
                    if self.height > self.chip_size[0]:
                        self.row_starts.append(max(0, self.height - self.chip_size[0]))
                    else:
                        # If the image is smaller than chip size, just start at 0
                        if not self.row_starts:
                            self.row_starts.append(0)

                    # Normal column starts using stride
                    for c in range((self.width - 1) // self.stride_x):
                        self.col_starts.append(c * self.stride_x)

                    # Add a special last column that ensures we reach the right edge
                    if self.width > self.chip_size[1]:
                        self.col_starts.append(max(0, self.width - self.chip_size[1]))
                    else:
                        # If the image is smaller than chip size, just start at 0
                        if not self.col_starts:
                            self.col_starts.append(0)

                # Calculate number of tiles
                self.rows = len(self.row_starts)
                self.cols = len(self.col_starts)

                print(
                    f"Dataset initialized with {self.rows} rows and {self.cols} columns of chips"
                )
                print(f"Image dimensions: {self.width} x {self.height} pixels")
                print(f"Chip size: {self.chip_size[1]} x {self.chip_size[0]} pixels")

            def __len__(self) -> int:
                return self.rows * self.cols

            def __getitem__(self, idx: int) -> Dict[str, Any]:
                # Convert flat index to grid position
                row = idx // self.cols
                col = idx % self.cols

                # Get pre-calculated starting positions
                j = self.row_starts[row]
                i = self.col_starts[col]

                # Read window from raster
                with rasterio.open(self.raster_path) as src:
                    # Make sure we don't read outside the image
                    width = min(self.chip_size[1], self.width - i)
                    height = min(self.chip_size[0], self.height - j)

                    window = Window(i, j, width, height)

                    # Read all bands
                    image = src.read(window=window)

                    # Handle partial windows at edges by padding
                    if (
                        image.shape[1] != self.chip_size[0]
                        or image.shape[2] != self.chip_size[1]
                    ):
                        temp = np.zeros(
                            (image.shape[0], self.chip_size[0], self.chip_size[1]),
                            dtype=image.dtype,
                        )
                        temp[:, : image.shape[1], : image.shape[2]] = image
                        image = temp

                # Preprocess bands for the model
                image_tensor = self.field_delineator.preprocess_sentinel_bands(
                    image, self.band_selection, self.use_ndvi
                )

                # Get geographic bounds for the window
                with rasterio.open(self.raster_path) as src:
                    window_transform = src.window_transform(window)
                    minx, miny = window_transform * (0, height)
                    maxx, maxy = window_transform * (width, 0)
                    bbox = [minx, miny, maxx, maxy]

                return {
                    "image": image_tensor,
                    "bbox": bbox,
                    "coords": torch.tensor([i, j], dtype=torch.long),
                    "window_size": torch.tensor([width, height], dtype=torch.long),
                }

        # Calculate stride based on overlap
        stride_x = int(chip_size[1] * (1 - overlap))
        stride_y = int(chip_size[0] * (1 - overlap))

        # Create dataset
        dataset = Sentinel2Dataset(
            raster_path=raster_path,
            chip_size=chip_size,
            stride_x=stride_x,
            stride_y=stride_y,
            band_selection=band_selection,
            use_ndvi=use_ndvi,
            field_delineator=self,
        )

        # Define custom collate function
        def custom_collate(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
            elem = batch[0]
            if isinstance(elem, dict):
                result = {}
                for key in elem:
                    if key == "bbox":
                        # Don't collate bbox objects, keep as list
                        result[key] = [d[key] for d in batch]
                    else:
                        # For tensors and other collatable types
                        try:
                            result[key] = (
                                torch.utils.data._utils.collate.default_collate(
                                    [d[key] for d in batch]
                                )
                            )
                        except TypeError:
                            # Fall back to list for non-collatable types
                            result[key] = [d[key] for d in batch]
                return result
            else:
                # Default collate for non-dict types
                return torch.utils.data._utils.collate.default_collate(batch)

        # Create dataloader
        dataloader = torch.utils.data.DataLoader(
            dataset,
            batch_size=batch_size,
            shuffle=False,
            num_workers=0,
            collate_fn=custom_collate,
        )

        # Process batches (call the parent class's process_raster method)
        # We'll adapt the process_raster method to work with our Sentinel2Dataset
        results = super().process_raster(
            raster_path=raster_path,
            output_path=output_path,
            batch_size=batch_size,
            filter_edges=filter_edges,
            edge_buffer=edge_buffer,
            confidence_threshold=confidence_threshold,
            overlap=overlap,
            chip_size=chip_size,
            nms_iou_threshold=nms_iou_threshold,
            mask_threshold=mask_threshold,
            min_object_area=min_object_area,
            simplify_tolerance=simplify_tolerance,
        )

        return results

__init__(model_path='field_boundary_detector.pth', repo_id=None, model=None, device=None, band_selection=None, use_ndvi=False)

Initialize the field boundary delineator.

Parameters:

Name	Type	Description	Default
model_path	str	Path to the .pth model file.	'field_boundary_detector.pth'
repo_id	Optional[str]	Repo ID for loading models from the Hub.	None
model	Optional[Any]	Custom model to use for inference.	None
device	Optional[str]	Device to use for inference ('cuda:0', 'cpu', etc.).	None
band_selection	Optional[List[int]]	List of Sentinel-2 band indices to use (None = adapt based on model)	None
use_ndvi	bool	Whether to calculate and include NDVI as an additional channel	False

Source code in geoai/extract.py

def __init__(
    self,
    model_path: str = "field_boundary_detector.pth",
    repo_id: Optional[str] = None,
    model: Optional[Any] = None,
    device: Optional[str] = None,
    band_selection: Optional[List[int]] = None,
    use_ndvi: bool = False,
) -> None:
    """
    Initialize the field boundary delineator.

    Args:
        model_path: Path to the .pth model file.
        repo_id: Repo ID for loading models from the Hub.
        model: Custom model to use for inference.
        device: Device to use for inference ('cuda:0', 'cpu', etc.).
        band_selection: List of Sentinel-2 band indices to use (None = adapt based on model)
        use_ndvi: Whether to calculate and include NDVI as an additional channel
    """
    # Save parameters before calling parent constructor
    self.custom_band_selection = band_selection
    self.use_ndvi = use_ndvi

    # Set device (copied from parent init to ensure it's set before initialize_model)
    if device is None:
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    else:
        self.device = torch.device(device)

    # Initialize model differently for multi-spectral input
    model = self.initialize_sentinel2_model(model)

    # Call parent but with our custom model
    super().__init__(
        model_path=model_path, repo_id=repo_id, model=model, device=device
    )

    # Default Sentinel-2 band statistics (can be overridden with actual stats)
    # Band order: [B1, B2, B3, B4, B5, B6, B7, B8, B8A, B9, B10, B11, B12]
    self.sentinel_band_stats = {
        "means": [
            0.0975,
            0.0476,
            0.0598,
            0.0697,
            0.1077,
            0.1859,
            0.2378,
            0.2061,
            0.2598,
            0.4120,
            0.1956,
            0.1410,
        ],
        "stds": [
            0.0551,
            0.0290,
            0.0298,
            0.0479,
            0.0506,
            0.0505,
            0.0747,
            0.0642,
            0.0782,
            0.1187,
            0.0651,
            0.0679,
        ],
    }

    # Set default band selection (RGB - typically B4, B3, B2 for Sentinel-2)
    self.band_selection = (
        self.custom_band_selection
        if self.custom_band_selection is not None
        else [3, 2, 1]
    )  # R, G, B bands

    # Customize parameters for field delineation
    self.confidence_threshold = 0.5  # Default confidence threshold
    self.overlap = 0.5  # Higher overlap for field boundary detection
    self.min_object_area = 1000  # Minimum area in pixels for field detection
    self.simplify_tolerance = 2.0  # Higher tolerance for field boundaries

initialize_sentinel2_model(model=None)

Initialize a Mask R-CNN model with a modified first layer to accept Sentinel-2 data.

Parameters:

Name	Type	Description	Default
model	Optional[Any]	Pre-initialized model (optional)	None

Returns:

Type	Description
Any	Modified model with appropriate input channels

Source code in geoai/extract.py

def initialize_sentinel2_model(self, model: Optional[Any] = None) -> Any:
    """
    Initialize a Mask R-CNN model with a modified first layer to accept Sentinel-2 data.

    Args:
        model: Pre-initialized model (optional)

    Returns:
        Modified model with appropriate input channels
    """
    import torchvision
    from torchvision.models.detection import maskrcnn_resnet50_fpn
    from torchvision.models.detection.backbone_utils import resnet_fpn_backbone

    if model is not None:
        return model

    # Determine number of input channels based on band selection and NDVI
    num_input_channels = (
        len(self.custom_band_selection)
        if self.custom_band_selection is not None
        else 3
    )
    if self.use_ndvi:
        num_input_channels += 1

    print(f"Initializing Mask R-CNN model with {num_input_channels} input channels")

    # Create a ResNet50 backbone with modified input channels
    backbone = resnet_fpn_backbone("resnet50", weights=None)

    # Replace the first conv layer to accept multi-spectral input
    original_conv = backbone.body.conv1
    backbone.body.conv1 = torch.nn.Conv2d(
        num_input_channels,
        original_conv.out_channels,
        kernel_size=original_conv.kernel_size,
        stride=original_conv.stride,
        padding=original_conv.padding,
        bias=original_conv.bias is not None,
    )

    # Create Mask R-CNN with our modified backbone
    model = maskrcnn_resnet50_fpn(
        backbone=backbone,
        num_classes=2,  # Background + field
        image_mean=[0.485] * num_input_channels,  # Extend mean to all channels
        image_std=[0.229] * num_input_channels,  # Extend std to all channels
    )

    model.to(self.device)
    return model

preprocess_sentinel_bands(image_data, band_selection=None, use_ndvi=None)

Preprocess Sentinel-2 band data for model input.

Parameters:

Name	Type	Description	Default
image_data	ndarray	Raw Sentinel-2 image data as numpy array [bands, height, width]	required
band_selection	Optional[List[int]]	List of band indices to use (overrides instance default if provided)	None
use_ndvi	Optional[bool]	Whether to include NDVI (overrides instance default if provided)	None

Returns:

Type	Description
Tensor	Processed tensor ready for model input

Source code in geoai/extract.py

def preprocess_sentinel_bands(
    self,
    image_data: np.ndarray,
    band_selection: Optional[List[int]] = None,
    use_ndvi: Optional[bool] = None,
) -> torch.Tensor:
    """
    Preprocess Sentinel-2 band data for model input.

    Args:
        image_data: Raw Sentinel-2 image data as numpy array [bands, height, width]
        band_selection: List of band indices to use (overrides instance default if provided)
        use_ndvi: Whether to include NDVI (overrides instance default if provided)

    Returns:
        Processed tensor ready for model input
    """
    # Use instance defaults if not specified
    band_selection = (
        band_selection if band_selection is not None else self.band_selection
    )
    use_ndvi = use_ndvi if use_ndvi is not None else self.use_ndvi

    # Select bands
    selected_bands = image_data[band_selection]

    # Calculate NDVI if requested (using B8 and B4 which are indices 7 and 3)
    if (
        use_ndvi
        and 7 in range(image_data.shape[0])
        and 3 in range(image_data.shape[0])
    ):
        nir = image_data[7].astype(np.float32)  # B8 (NIR)
        red = image_data[3].astype(np.float32)  # B4 (Red)

        # Avoid division by zero
        denominator = nir + red
        ndvi = np.zeros_like(nir)
        valid_mask = denominator > 0
        ndvi[valid_mask] = (nir[valid_mask] - red[valid_mask]) / denominator[
            valid_mask
        ]

        # Rescale NDVI from [-1, 1] to [0, 1]
        ndvi = (ndvi + 1) / 2

        # Add NDVI as an additional channel
        selected_bands = np.vstack([selected_bands, ndvi[np.newaxis, :, :]])

    # Convert to tensor
    image_tensor = torch.from_numpy(selected_bands).float()

    # Normalize using band statistics
    for i, band_idx in enumerate(band_selection):
        # Make sure band_idx is within range of our statistics
        if band_idx < len(self.sentinel_band_stats["means"]):
            mean = self.sentinel_band_stats["means"][band_idx]
            std = self.sentinel_band_stats["stds"][band_idx]
            image_tensor[i] = (image_tensor[i] - mean) / std

    # If NDVI was added, normalize it too (last channel)
    if use_ndvi:
        # NDVI is already roughly in [0,1] range, just standardize it slightly
        image_tensor[-1] = (image_tensor[-1] - 0.5) / 0.5

    return image_tensor

process_sentinel_raster(raster_path, output_path=None, batch_size=4, band_selection=None, use_ndvi=None, filter_edges=True, edge_buffer=20, **kwargs)

Process a Sentinel-2 raster to extract field boundaries.

Parameters:

Name	Type	Description	Default
raster_path	str	Path to Sentinel-2 raster file	required
output_path	Optional[str]	Path to output GeoJSON or Parquet file (optional)	None
batch_size	int	Batch size for processing	4
band_selection	Optional[List[int]]	List of bands to use (None = use instance default)	None
use_ndvi	Optional[bool]	Whether to include NDVI (None = use instance default)	None
filter_edges	bool	Whether to filter out objects at the edges of the image	True
edge_buffer	int	Size of edge buffer in pixels to filter out objects	20
**kwargs	Any	Additional parameters for processing	{}

Returns:

Type	Description
GeoDataFrame	GeoDataFrame with field boundaries

Source code in geoai/extract.py

def process_sentinel_raster(
    self,
    raster_path: str,
    output_path: Optional[str] = None,
    batch_size: int = 4,
    band_selection: Optional[List[int]] = None,
    use_ndvi: Optional[bool] = None,
    filter_edges: bool = True,
    edge_buffer: int = 20,
    **kwargs: Any,
) -> gpd.GeoDataFrame:
    """
    Process a Sentinel-2 raster to extract field boundaries.

    Args:
        raster_path: Path to Sentinel-2 raster file
        output_path: Path to output GeoJSON or Parquet file (optional)
        batch_size: Batch size for processing
        band_selection: List of bands to use (None = use instance default)
        use_ndvi: Whether to include NDVI (None = use instance default)
        filter_edges: Whether to filter out objects at the edges of the image
        edge_buffer: Size of edge buffer in pixels to filter out objects
        **kwargs: Additional parameters for processing

    Returns:
        GeoDataFrame with field boundaries
    """
    # Use instance defaults if not specified
    band_selection = (
        band_selection if band_selection is not None else self.band_selection
    )
    use_ndvi = use_ndvi if use_ndvi is not None else self.use_ndvi

    # Get parameters from kwargs or use instance defaults
    confidence_threshold = kwargs.get(
        "confidence_threshold", self.confidence_threshold
    )
    overlap = kwargs.get("overlap", self.overlap)
    chip_size = kwargs.get("chip_size", self.chip_size)
    nms_iou_threshold = kwargs.get("nms_iou_threshold", self.nms_iou_threshold)
    mask_threshold = kwargs.get("mask_threshold", self.mask_threshold)
    min_object_area = kwargs.get("min_object_area", self.min_object_area)
    simplify_tolerance = kwargs.get("simplify_tolerance", self.simplify_tolerance)

    # Update band statistics if not already done
    if kwargs.get("update_stats", True):
        self.update_band_stats(raster_path, band_selection)

    print(f"Processing with parameters:")
    print(f"- Using bands: {band_selection}")
    print(f"- Include NDVI: {use_ndvi}")
    print(f"- Confidence threshold: {confidence_threshold}")
    print(f"- Tile overlap: {overlap}")
    print(f"- Chip size: {chip_size}")
    print(f"- Filter edge objects: {filter_edges}")

    # Create a custom Sentinel-2 dataset class
    class Sentinel2Dataset(torch.utils.data.Dataset):
        def __init__(
            self,
            raster_path: str,
            chip_size: Tuple[int, int],
            stride_x: int,
            stride_y: int,
            band_selection: List[int],
            use_ndvi: bool,
            field_delineator: Any,
        ) -> None:
            self.raster_path = raster_path
            self.chip_size = chip_size
            self.stride_x = stride_x
            self.stride_y = stride_y
            self.band_selection = band_selection
            self.use_ndvi = use_ndvi
            self.field_delineator = field_delineator

            with rasterio.open(self.raster_path) as src:
                self.height = src.height
                self.width = src.width
                self.count = src.count
                self.crs = src.crs
                self.transform = src.transform

                # Calculate row_starts and col_starts
                self.row_starts = []
                self.col_starts = []

                # Normal row starts using stride
                for r in range((self.height - 1) // self.stride_y):
                    self.row_starts.append(r * self.stride_y)

                # Add a special last row that ensures we reach the bottom edge
                if self.height > self.chip_size[0]:
                    self.row_starts.append(max(0, self.height - self.chip_size[0]))
                else:
                    # If the image is smaller than chip size, just start at 0
                    if not self.row_starts:
                        self.row_starts.append(0)

                # Normal column starts using stride
                for c in range((self.width - 1) // self.stride_x):
                    self.col_starts.append(c * self.stride_x)

                # Add a special last column that ensures we reach the right edge
                if self.width > self.chip_size[1]:
                    self.col_starts.append(max(0, self.width - self.chip_size[1]))
                else:
                    # If the image is smaller than chip size, just start at 0
                    if not self.col_starts:
                        self.col_starts.append(0)

            # Calculate number of tiles
            self.rows = len(self.row_starts)
            self.cols = len(self.col_starts)

            print(
                f"Dataset initialized with {self.rows} rows and {self.cols} columns of chips"
            )
            print(f"Image dimensions: {self.width} x {self.height} pixels")
            print(f"Chip size: {self.chip_size[1]} x {self.chip_size[0]} pixels")

        def __len__(self) -> int:
            return self.rows * self.cols

        def __getitem__(self, idx: int) -> Dict[str, Any]:
            # Convert flat index to grid position
            row = idx // self.cols
            col = idx % self.cols

            # Get pre-calculated starting positions
            j = self.row_starts[row]
            i = self.col_starts[col]

            # Read window from raster
            with rasterio.open(self.raster_path) as src:
                # Make sure we don't read outside the image
                width = min(self.chip_size[1], self.width - i)
                height = min(self.chip_size[0], self.height - j)

                window = Window(i, j, width, height)

                # Read all bands
                image = src.read(window=window)

                # Handle partial windows at edges by padding
                if (
                    image.shape[1] != self.chip_size[0]
                    or image.shape[2] != self.chip_size[1]
                ):
                    temp = np.zeros(
                        (image.shape[0], self.chip_size[0], self.chip_size[1]),
                        dtype=image.dtype,
                    )
                    temp[:, : image.shape[1], : image.shape[2]] = image
                    image = temp

            # Preprocess bands for the model
            image_tensor = self.field_delineator.preprocess_sentinel_bands(
                image, self.band_selection, self.use_ndvi
            )

            # Get geographic bounds for the window
            with rasterio.open(self.raster_path) as src:
                window_transform = src.window_transform(window)
                minx, miny = window_transform * (0, height)
                maxx, maxy = window_transform * (width, 0)
                bbox = [minx, miny, maxx, maxy]

            return {
                "image": image_tensor,
                "bbox": bbox,
                "coords": torch.tensor([i, j], dtype=torch.long),
                "window_size": torch.tensor([width, height], dtype=torch.long),
            }

    # Calculate stride based on overlap
    stride_x = int(chip_size[1] * (1 - overlap))
    stride_y = int(chip_size[0] * (1 - overlap))

    # Create dataset
    dataset = Sentinel2Dataset(
        raster_path=raster_path,
        chip_size=chip_size,
        stride_x=stride_x,
        stride_y=stride_y,
        band_selection=band_selection,
        use_ndvi=use_ndvi,
        field_delineator=self,
    )

    # Define custom collate function
    def custom_collate(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
        elem = batch[0]
        if isinstance(elem, dict):
            result = {}
            for key in elem:
                if key == "bbox":
                    # Don't collate bbox objects, keep as list
                    result[key] = [d[key] for d in batch]
                else:
                    # For tensors and other collatable types
                    try:
                        result[key] = (
                            torch.utils.data._utils.collate.default_collate(
                                [d[key] for d in batch]
                            )
                        )
                    except TypeError:
                        # Fall back to list for non-collatable types
                        result[key] = [d[key] for d in batch]
            return result
        else:
            # Default collate for non-dict types
            return torch.utils.data._utils.collate.default_collate(batch)

    # Create dataloader
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=0,
        collate_fn=custom_collate,
    )

    # Process batches (call the parent class's process_raster method)
    # We'll adapt the process_raster method to work with our Sentinel2Dataset
    results = super().process_raster(
        raster_path=raster_path,
        output_path=output_path,
        batch_size=batch_size,
        filter_edges=filter_edges,
        edge_buffer=edge_buffer,
        confidence_threshold=confidence_threshold,
        overlap=overlap,
        chip_size=chip_size,
        nms_iou_threshold=nms_iou_threshold,
        mask_threshold=mask_threshold,
        min_object_area=min_object_area,
        simplify_tolerance=simplify_tolerance,
    )

    return results

update_band_stats(raster_path, band_selection=None, sample_size=1000)

Update band statistics from the input Sentinel-2 raster.

Parameters:

Name	Type	Description	Default
raster_path	str	Path to the Sentinel-2 raster file	required
band_selection	Optional[List[int]]	Specific bands to update (None = update all available)	None
sample_size	int	Number of random pixels to sample for statistics calculation	1000

Returns:

Type	Description
Dict[str, List[float]]	Updated band statistics dictionary

Source code in geoai/extract.py

def update_band_stats(
    self,
    raster_path: str,
    band_selection: Optional[List[int]] = None,
    sample_size: int = 1000,
) -> Dict[str, List[float]]:
    """
    Update band statistics from the input Sentinel-2 raster.

    Args:
        raster_path: Path to the Sentinel-2 raster file
        band_selection: Specific bands to update (None = update all available)
        sample_size: Number of random pixels to sample for statistics calculation

    Returns:
        Updated band statistics dictionary
    """
    with rasterio.open(raster_path) as src:
        # Check if this is likely a Sentinel-2 product
        band_count = src.count
        if band_count < 3:
            print(
                f"Warning: Raster has only {band_count} bands, may not be Sentinel-2 data"
            )

        # Get dimensions
        height, width = src.height, src.width

        # Determine which bands to analyze
        if band_selection is None:
            band_selection = list(range(1, band_count + 1))  # 1-indexed

        # Initialize arrays for band statistics
        means = []
        stds = []

        # Sample random pixels
        np.random.seed(42)  # For reproducibility
        sample_rows = np.random.randint(0, height, sample_size)
        sample_cols = np.random.randint(0, width, sample_size)

        # Calculate statistics for each band
        for band in band_selection:
            # Read band data
            band_data = src.read(band)

            # Sample values
            sample_values = band_data[sample_rows, sample_cols]

            # Remove invalid values (e.g., nodata)
            valid_samples = sample_values[np.isfinite(sample_values)]

            # Calculate statistics
            mean = float(np.mean(valid_samples))
            std = float(np.std(valid_samples))

            # Store results
            means.append(mean)
            stds.append(std)

            print(f"Band {band}: mean={mean:.4f}, std={std:.4f}")

        # Update instance variables
        self.sentinel_band_stats = {"means": means, "stds": stds}

        return self.sentinel_band_stats

BoundingBox dataclass

Represents a bounding box with coordinates.

Source code in geoai/segment.py

@dataclass
class BoundingBox:
    """Represents a bounding box with coordinates."""

    xmin: int
    ymin: int
    xmax: int
    ymax: int

    @property
    def xyxy(self) -> List[float]:
        return [self.xmin, self.ymin, self.xmax, self.ymax]

BuildingFootprintExtractor

Bases: ObjectDetector

Building footprint extraction using a pre-trained Mask R-CNN model.

This class extends the ObjectDetector class with additional methods for building footprint extraction."

Source code in geoai/extract.py

class BuildingFootprintExtractor(ObjectDetector):
    """
    Building footprint extraction using a pre-trained Mask R-CNN model.

    This class extends the
    `ObjectDetector` class with additional methods for building footprint extraction."
    """

    def __init__(
        self,
        model_path: str = "building_footprints_usa.pth",
        repo_id: Optional[str] = None,
        model: Optional[Any] = None,
        device: Optional[str] = None,
    ) -> None:
        """
        Initialize the object extractor.

        Args:
            model_path: Path to the .pth model file.
            repo_id: Repo ID for loading models from the Hub.
            model: Custom model to use for inference.
            device: Device to use for inference ('cuda:0', 'cpu', etc.).
        """
        super().__init__(
            model_path=model_path, repo_id=repo_id, model=model, device=device
        )

    def regularize_buildings(
        self,
        gdf: gpd.GeoDataFrame,
        min_area: int = 10,
        angle_threshold: int = 15,
        orthogonality_threshold: float = 0.3,
        rectangularity_threshold: float = 0.7,
    ) -> gpd.GeoDataFrame:
        """
        Regularize building footprints to enforce right angles and rectangular shapes.

        Args:
            gdf: GeoDataFrame with building footprints
            min_area: Minimum area in square units to keep a building
            angle_threshold: Maximum deviation from 90 degrees to consider an angle as orthogonal (degrees)
            orthogonality_threshold: Percentage of angles that must be orthogonal for a building to be regularized
            rectangularity_threshold: Minimum area ratio to building's oriented bounding box for rectangular simplification

        Returns:
            GeoDataFrame with regularized building footprints
        """
        return self.regularize_objects(
            gdf,
            min_area=min_area,
            angle_threshold=angle_threshold,
            orthogonality_threshold=orthogonality_threshold,
            rectangularity_threshold=rectangularity_threshold,
        )

__init__(model_path='building_footprints_usa.pth', repo_id=None, model=None, device=None)

Initialize the object extractor.

Parameters:

Name	Type	Description	Default
model_path	str	Path to the .pth model file.	'building_footprints_usa.pth'
repo_id	Optional[str]	Repo ID for loading models from the Hub.	None
model	Optional[Any]	Custom model to use for inference.	None
device	Optional[str]	Device to use for inference ('cuda:0', 'cpu', etc.).	None

Source code in geoai/extract.py

def __init__(
    self,
    model_path: str = "building_footprints_usa.pth",
    repo_id: Optional[str] = None,
    model: Optional[Any] = None,
    device: Optional[str] = None,
) -> None:
    """
    Initialize the object extractor.

    Args:
        model_path: Path to the .pth model file.
        repo_id: Repo ID for loading models from the Hub.
        model: Custom model to use for inference.
        device: Device to use for inference ('cuda:0', 'cpu', etc.).
    """
    super().__init__(
        model_path=model_path, repo_id=repo_id, model=model, device=device
    )

regularize_buildings(gdf, min_area=10, angle_threshold=15, orthogonality_threshold=0.3, rectangularity_threshold=0.7)

Regularize building footprints to enforce right angles and rectangular shapes.

Parameters:

Name	Type	Description	Default
gdf	GeoDataFrame	GeoDataFrame with building footprints	required
min_area	int	Minimum area in square units to keep a building	10
angle_threshold	int	Maximum deviation from 90 degrees to consider an angle as orthogonal (degrees)	15
orthogonality_threshold	float	Percentage of angles that must be orthogonal for a building to be regularized	0.3
rectangularity_threshold	float	Minimum area ratio to building's oriented bounding box for rectangular simplification	0.7

Returns:

Type	Description
GeoDataFrame	GeoDataFrame with regularized building footprints

Source code in geoai/extract.py

def regularize_buildings(
    self,
    gdf: gpd.GeoDataFrame,
    min_area: int = 10,
    angle_threshold: int = 15,
    orthogonality_threshold: float = 0.3,
    rectangularity_threshold: float = 0.7,
) -> gpd.GeoDataFrame:
    """
    Regularize building footprints to enforce right angles and rectangular shapes.

    Args:
        gdf: GeoDataFrame with building footprints
        min_area: Minimum area in square units to keep a building
        angle_threshold: Maximum deviation from 90 degrees to consider an angle as orthogonal (degrees)
        orthogonality_threshold: Percentage of angles that must be orthogonal for a building to be regularized
        rectangularity_threshold: Minimum area ratio to building's oriented bounding box for rectangular simplification

    Returns:
        GeoDataFrame with regularized building footprints
    """
    return self.regularize_objects(
        gdf,
        min_area=min_area,
        angle_threshold=angle_threshold,
        orthogonality_threshold=orthogonality_threshold,
        rectangularity_threshold=rectangularity_threshold,
    )

CLIPSegmentation

A class for segmenting high-resolution satellite imagery using text prompts with CLIP-based models.

This segmenter utilizes the CLIP-Seg model to perform semantic segmentation based on text prompts. It can process large GeoTIFF files by tiling them and handles proper georeferencing in the output.

Parameters:

Name	Type	Description	Default
model_name	str	Name of the CLIP-Seg model to use. Defaults to "CIDAS/clipseg-rd64-refined".	'CIDAS/clipseg-rd64-refined'
device	str	Device to run the model on ('cuda', 'mps', 'cpu'). If None, auto-selects the best available device (CUDA > MPS > CPU) via get_device().	None
tile_size	int	Size of tiles to process the image in chunks. Defaults to 352.	512
overlap	int	Overlap between tiles to avoid edge artifacts. Defaults to 16.	32

Attributes:

Name	Type	Description
processor	CLIPSegProcessor	The processor for the CLIP-Seg model.
model	CLIPSegForImageSegmentation	The CLIP-Seg model for segmentation.
device	str	The device being used ('cuda', 'mps', or 'cpu').
tile_size	int	Size of tiles for processing.
overlap	int	Overlap between tiles.

Source code in geoai/segment.py

class CLIPSegmentation:
    """
    A class for segmenting high-resolution satellite imagery using text prompts with CLIP-based models.

    This segmenter utilizes the CLIP-Seg model to perform semantic segmentation based on text prompts.
    It can process large GeoTIFF files by tiling them and handles proper georeferencing in the output.

    Args:
        model_name (str): Name of the CLIP-Seg model to use. Defaults to "CIDAS/clipseg-rd64-refined".
        device (str): Device to run the model on ('cuda', 'mps', 'cpu').
            If None, auto-selects the best available device
            (CUDA > MPS > CPU) via ``get_device()``.
        tile_size (int): Size of tiles to process the image in chunks. Defaults to 352.
        overlap (int): Overlap between tiles to avoid edge artifacts. Defaults to 16.

    Attributes:
        processor (CLIPSegProcessor): The processor for the CLIP-Seg model.
        model (CLIPSegForImageSegmentation): The CLIP-Seg model for segmentation.
        device (str): The device being used ('cuda', 'mps', or 'cpu').
        tile_size (int): Size of tiles for processing.
        overlap (int): Overlap between tiles.
    """

    def __init__(
        self,
        model_name: str = "CIDAS/clipseg-rd64-refined",
        device: Optional[str] = None,
        tile_size: int = 512,
        overlap: int = 32,
    ) -> None:
        """
        Initialize the ImageSegmenter with the specified model and settings.

        Args:
            model_name (str): Name of the CLIP-Seg model to use. Defaults to "CIDAS/clipseg-rd64-refined".
            device (str): Device to run the model on ('cuda', 'mps', 'cpu').
                If None, auto-selects the best available device
                (CUDA > MPS > CPU) via ``get_device()``.
            tile_size (int): Size of tiles to process the image in chunks. Defaults to 512.
            overlap (int): Overlap between tiles to avoid edge artifacts. Defaults to 32.
        """
        self.tile_size = tile_size
        self.overlap = overlap

        # Lazy import to avoid pulling in geoai.utils (which eagerly
        # imports leafmap) at module-collection time.
        from .utils.device import get_device

        self.device = str(device or get_device())

        # Load model and processor
        self.processor = CLIPSegProcessor.from_pretrained(model_name)
        self.model = CLIPSegForImageSegmentation.from_pretrained(model_name).to(
            self.device
        )

        logger.info("CLIPSeg model loaded on %s", self.device)

    def segment_image(
        self,
        input_path: str,
        output_path: str,
        text_prompt: str,
        threshold: float = 0.5,
        smoothing_sigma: float = 1.0,
    ) -> str:
        """
        Segment a GeoTIFF image using the provided text prompt.

        The function processes the image in tiles and saves the result as a GeoTIFF with two bands:
        - Band 1: Binary segmentation mask (0 or 1)
        - Band 2: Probability scores (0.0 to 1.0)

        Args:
            input_path (str): Path to the input GeoTIFF file.
            output_path (str): Path where the output GeoTIFF will be saved.
            text_prompt (str): Text description of what to segment (e.g., "water", "buildings").
            threshold (float): Threshold for binary segmentation (0.0 to 1.0). Defaults to 0.5.
            smoothing_sigma (float): Sigma value for Gaussian smoothing to reduce blockiness. Defaults to 1.0.

        Returns:
            str: Path to the saved output file.
        """
        # Open the input GeoTIFF
        with rasterio.open(input_path) as src:
            # Get metadata
            meta = src.meta
            height = src.height
            width = src.width

            # Create output metadata
            out_meta = meta.copy()
            out_meta.update({"count": 2, "dtype": "float32", "nodata": None})

            # Create arrays for results
            segmentation = np.zeros((height, width), dtype=np.float32)
            probabilities = np.zeros((height, width), dtype=np.float32)

            # Calculate effective tile size (accounting for overlap)
            effective_tile_size = self.tile_size - 2 * self.overlap

            # Calculate number of tiles
            n_tiles_x = max(1, int(np.ceil(width / effective_tile_size)))
            n_tiles_y = max(1, int(np.ceil(height / effective_tile_size)))
            total_tiles = n_tiles_x * n_tiles_y

            # Process tiles with tqdm progress bar
            with tqdm(total=total_tiles, desc="Processing tiles") as pbar:
                # Iterate through tiles
                for y in range(n_tiles_y):
                    for x in range(n_tiles_x):
                        # Calculate tile coordinates with overlap
                        x_start = max(0, x * effective_tile_size - self.overlap)
                        y_start = max(0, y * effective_tile_size - self.overlap)
                        x_end = min(width, (x + 1) * effective_tile_size + self.overlap)
                        y_end = min(
                            height, (y + 1) * effective_tile_size + self.overlap
                        )

                        tile_width = x_end - x_start
                        tile_height = y_end - y_start

                        # Read the tile
                        window = Window(x_start, y_start, tile_width, tile_height)
                        tile_data = src.read(window=window)

                        # Process the tile
                        try:
                            # Convert to RGB format
                            bands = tile_data.shape[0]
                            if bands >= 3:
                                rgb_tile = tile_data[:3].transpose(1, 2, 0)
                            elif bands == 1:
                                rgb_tile = np.repeat(
                                    tile_data[0][:, :, np.newaxis], 3, axis=2
                                )
                            else:
                                # 2-band: replicate first band into 3 channels
                                rgb_tile = np.repeat(
                                    tile_data[0][:, :, np.newaxis], 3, axis=2
                                )

                            # Normalize to uint8 if needed
                            if rgb_tile.dtype == np.uint8:
                                # Already uint8 — skip normalization
                                if rgb_tile.max() == 0:
                                    pbar.update(1)
                                    continue
                            else:
                                # Float / other dtypes: per-channel normalization
                                rgb_tile = rgb_tile.astype(np.float64)
                                result_tile = np.empty_like(rgb_tile)
                                all_zero = True
                                for ch in range(rgb_tile.shape[2]):
                                    cmin = rgb_tile[:, :, ch].min()
                                    cmax = rgb_tile[:, :, ch].max()
                                    if cmax > cmin:
                                        result_tile[:, :, ch] = (
                                            (rgb_tile[:, :, ch] - cmin)
                                            / (cmax - cmin)
                                            * 255.0
                                        )
                                        all_zero = False
                                    elif cmax > 0:
                                        result_tile[:, :, ch] = 128.0
                                        all_zero = False
                                    else:
                                        result_tile[:, :, ch] = 0.0
                                if all_zero:
                                    pbar.update(1)
                                    continue
                                rgb_tile = result_tile.astype(np.uint8)

                            # Convert to PIL Image
                            pil_image = Image.fromarray(rgb_tile)

                            # Resize if needed to match model's requirements
                            if (
                                pil_image.width > self.tile_size
                                or pil_image.height > self.tile_size
                            ):
                                # Keep aspect ratio - use LANCZOS resampling instead of deprecated constant
                                pil_image.thumbnail(
                                    (self.tile_size, self.tile_size),
                                    Image.Resampling.LANCZOS,
                                )

                            # Process with CLIP-Seg
                            inputs = self.processor(
                                text=text_prompt, images=pil_image, return_tensors="pt"
                            ).to(self.device)

                            # Forward pass
                            with torch.no_grad():
                                outputs = self.model(**inputs)

                            # Get logits and resize to original tile size
                            logits = outputs.logits[0]

                            # Convert logits to probabilities with sigmoid
                            probs = torch.sigmoid(logits).cpu().numpy()

                            # Resize back to original tile size if needed
                            if probs.shape != (tile_height, tile_width):
                                # Use bicubic interpolation for smoother results
                                probs_resized = np.array(
                                    Image.fromarray(probs).resize(
                                        (tile_width, tile_height),
                                        Image.Resampling.BICUBIC,
                                    )
                                )
                            else:
                                probs_resized = probs

                            # Apply gaussian blur to reduce blockiness
                            try:
                                from scipy.ndimage import gaussian_filter

                                probs_resized = gaussian_filter(
                                    probs_resized, sigma=smoothing_sigma
                                )
                            except ImportError:
                                pass  # Continue without smoothing if scipy is not available

                            # Store results in the full arrays
                            # Only store the non-overlapping part (except at edges)
                            valid_x_start = self.overlap if x > 0 else 0
                            valid_y_start = self.overlap if y > 0 else 0
                            valid_x_end = (
                                tile_width - self.overlap
                                if x < n_tiles_x - 1
                                else tile_width
                            )
                            valid_y_end = (
                                tile_height - self.overlap
                                if y < n_tiles_y - 1
                                else tile_height
                            )

                            dest_x_start = x_start + valid_x_start
                            dest_y_start = y_start + valid_y_start
                            dest_x_end = x_start + valid_x_end
                            dest_y_end = y_start + valid_y_end

                            # Store probabilities
                            probabilities[
                                dest_y_start:dest_y_end, dest_x_start:dest_x_end
                            ] = probs_resized[
                                valid_y_start:valid_y_end, valid_x_start:valid_x_end
                            ]

                        except Exception as e:
                            logger.warning(
                                "Error processing tile at (%d, %d): %s", x, y, str(e)
                            )
                            # Continue with next tile

                        # Update progress bar
                        pbar.update(1)

            # Create binary segmentation from probabilities
            segmentation = (probabilities >= threshold).astype(np.float32)

            # Write the output GeoTIFF
            with rasterio.open(output_path, "w", **out_meta) as dst:
                dst.write(segmentation, 1)
                dst.write(probabilities, 2)

                # Add descriptions to bands
                dst.set_band_description(1, "Binary Segmentation")
                dst.set_band_description(2, "Probability Scores")

            logger.info("Segmentation saved to %s", output_path)
            return output_path

    def segment_image_batch(
        self,
        input_paths: List[str],
        output_dir: str,
        text_prompt: str,
        threshold: float = 0.5,
        smoothing_sigma: float = 1.0,
        suffix: str = "_segmented",
    ) -> List[str]:
        """
        Segment multiple GeoTIFF images using the provided text prompt.

        Args:
            input_paths (list): List of paths to input GeoTIFF files.
            output_dir (str): Directory where output GeoTIFFs will be saved.
            text_prompt (str): Text description of what to segment.
            threshold (float): Threshold for binary segmentation. Defaults to 0.5.
            smoothing_sigma (float): Sigma value for Gaussian smoothing to reduce blockiness. Defaults to 1.0.
            suffix (str): Suffix to add to output filenames. Defaults to "_segmented".

        Returns:
            list: Paths to all saved output files.
        """
        # Create output directory if it doesn't exist
        os.makedirs(output_dir, exist_ok=True)

        output_paths = []

        # Process each input file
        for input_path in tqdm(input_paths, desc="Processing files"):
            # Generate output path
            filename = os.path.basename(input_path)
            base_name, ext = os.path.splitext(filename)
            output_path = os.path.join(output_dir, f"{base_name}{suffix}{ext}")

            # Segment the image
            result_path = self.segment_image(
                input_path, output_path, text_prompt, threshold, smoothing_sigma
            )
            output_paths.append(result_path)

        return output_paths

__init__(model_name='CIDAS/clipseg-rd64-refined', device=None, tile_size=512, overlap=32)

Initialize the ImageSegmenter with the specified model and settings.

Parameters:

Name	Type	Description	Default
model_name	str	Name of the CLIP-Seg model to use. Defaults to "CIDAS/clipseg-rd64-refined".	'CIDAS/clipseg-rd64-refined'
device	str	Device to run the model on ('cuda', 'mps', 'cpu'). If None, auto-selects the best available device (CUDA > MPS > CPU) via get_device().	None
tile_size	int	Size of tiles to process the image in chunks. Defaults to 512.	512
overlap	int	Overlap between tiles to avoid edge artifacts. Defaults to 32.	32

Source code in geoai/segment.py

def __init__(
    self,
    model_name: str = "CIDAS/clipseg-rd64-refined",
    device: Optional[str] = None,
    tile_size: int = 512,
    overlap: int = 32,
) -> None:
    """
    Initialize the ImageSegmenter with the specified model and settings.

    Args:
        model_name (str): Name of the CLIP-Seg model to use. Defaults to "CIDAS/clipseg-rd64-refined".
        device (str): Device to run the model on ('cuda', 'mps', 'cpu').
            If None, auto-selects the best available device
            (CUDA > MPS > CPU) via ``get_device()``.
        tile_size (int): Size of tiles to process the image in chunks. Defaults to 512.
        overlap (int): Overlap between tiles to avoid edge artifacts. Defaults to 32.
    """
    self.tile_size = tile_size
    self.overlap = overlap

    # Lazy import to avoid pulling in geoai.utils (which eagerly
    # imports leafmap) at module-collection time.
    from .utils.device import get_device

    self.device = str(device or get_device())

    # Load model and processor
    self.processor = CLIPSegProcessor.from_pretrained(model_name)
    self.model = CLIPSegForImageSegmentation.from_pretrained(model_name).to(
        self.device
    )

    logger.info("CLIPSeg model loaded on %s", self.device)

segment_image(input_path, output_path, text_prompt, threshold=0.5, smoothing_sigma=1.0)

Segment a GeoTIFF image using the provided text prompt.

The function processes the image in tiles and saves the result as a GeoTIFF with two bands: - Band 1: Binary segmentation mask (0 or 1) - Band 2: Probability scores (0.0 to 1.0)

Parameters:

Name	Type	Description	Default
input_path	str	Path to the input GeoTIFF file.	required
output_path	str	Path where the output GeoTIFF will be saved.	required
text_prompt	str	Text description of what to segment (e.g., "water", "buildings").	required
threshold	float	Threshold for binary segmentation (0.0 to 1.0). Defaults to 0.5.	0.5
smoothing_sigma	float	Sigma value for Gaussian smoothing to reduce blockiness. Defaults to 1.0.	1.0

Returns:

Name	Type	Description
str	str	Path to the saved output file.

Source code in geoai/segment.py

def segment_image(
    self,
    input_path: str,
    output_path: str,
    text_prompt: str,
    threshold: float = 0.5,
    smoothing_sigma: float = 1.0,
) -> str:
    """
    Segment a GeoTIFF image using the provided text prompt.

    The function processes the image in tiles and saves the result as a GeoTIFF with two bands:
    - Band 1: Binary segmentation mask (0 or 1)
    - Band 2: Probability scores (0.0 to 1.0)

    Args:
        input_path (str): Path to the input GeoTIFF file.
        output_path (str): Path where the output GeoTIFF will be saved.
        text_prompt (str): Text description of what to segment (e.g., "water", "buildings").
        threshold (float): Threshold for binary segmentation (0.0 to 1.0). Defaults to 0.5.
        smoothing_sigma (float): Sigma value for Gaussian smoothing to reduce blockiness. Defaults to 1.0.

    Returns:
        str: Path to the saved output file.
    """
    # Open the input GeoTIFF
    with rasterio.open(input_path) as src:
        # Get metadata
        meta = src.meta
        height = src.height
        width = src.width

        # Create output metadata
        out_meta = meta.copy()
        out_meta.update({"count": 2, "dtype": "float32", "nodata": None})

        # Create arrays for results
        segmentation = np.zeros((height, width), dtype=np.float32)
        probabilities = np.zeros((height, width), dtype=np.float32)

        # Calculate effective tile size (accounting for overlap)
        effective_tile_size = self.tile_size - 2 * self.overlap

        # Calculate number of tiles
        n_tiles_x = max(1, int(np.ceil(width / effective_tile_size)))
        n_tiles_y = max(1, int(np.ceil(height / effective_tile_size)))
        total_tiles = n_tiles_x * n_tiles_y

        # Process tiles with tqdm progress bar
        with tqdm(total=total_tiles, desc="Processing tiles") as pbar:
            # Iterate through tiles
            for y in range(n_tiles_y):
                for x in range(n_tiles_x):
                    # Calculate tile coordinates with overlap
                    x_start = max(0, x * effective_tile_size - self.overlap)
                    y_start = max(0, y * effective_tile_size - self.overlap)
                    x_end = min(width, (x + 1) * effective_tile_size + self.overlap)
                    y_end = min(
                        height, (y + 1) * effective_tile_size + self.overlap
                    )

                    tile_width = x_end - x_start
                    tile_height = y_end - y_start

                    # Read the tile
                    window = Window(x_start, y_start, tile_width, tile_height)
                    tile_data = src.read(window=window)

                    # Process the tile
                    try:
                        # Convert to RGB format
                        bands = tile_data.shape[0]
                        if bands >= 3:
                            rgb_tile = tile_data[:3].transpose(1, 2, 0)
                        elif bands == 1:
                            rgb_tile = np.repeat(
                                tile_data[0][:, :, np.newaxis], 3, axis=2
                            )
                        else:
                            # 2-band: replicate first band into 3 channels
                            rgb_tile = np.repeat(
                                tile_data[0][:, :, np.newaxis], 3, axis=2
                            )

                        # Normalize to uint8 if needed
                        if rgb_tile.dtype == np.uint8:
                            # Already uint8 — skip normalization
                            if rgb_tile.max() == 0:
                                pbar.update(1)
                                continue
                        else:
                            # Float / other dtypes: per-channel normalization
                            rgb_tile = rgb_tile.astype(np.float64)
                            result_tile = np.empty_like(rgb_tile)
                            all_zero = True
                            for ch in range(rgb_tile.shape[2]):
                                cmin = rgb_tile[:, :, ch].min()
                                cmax = rgb_tile[:, :, ch].max()
                                if cmax > cmin:
                                    result_tile[:, :, ch] = (
                                        (rgb_tile[:, :, ch] - cmin)
                                        / (cmax - cmin)
                                        * 255.0
                                    )
                                    all_zero = False
                                elif cmax > 0:
                                    result_tile[:, :, ch] = 128.0
                                    all_zero = False
                                else:
                                    result_tile[:, :, ch] = 0.0
                            if all_zero:
                                pbar.update(1)
                                continue
                            rgb_tile = result_tile.astype(np.uint8)

                        # Convert to PIL Image
                        pil_image = Image.fromarray(rgb_tile)

                        # Resize if needed to match model's requirements
                        if (
                            pil_image.width > self.tile_size
                            or pil_image.height > self.tile_size
                        ):
                            # Keep aspect ratio - use LANCZOS resampling instead of deprecated constant
                            pil_image.thumbnail(
                                (self.tile_size, self.tile_size),
                                Image.Resampling.LANCZOS,
                            )

                        # Process with CLIP-Seg
                        inputs = self.processor(
                            text=text_prompt, images=pil_image, return_tensors="pt"
                        ).to(self.device)

                        # Forward pass
                        with torch.no_grad():
                            outputs = self.model(**inputs)

                        # Get logits and resize to original tile size
                        logits = outputs.logits[0]

                        # Convert logits to probabilities with sigmoid
                        probs = torch.sigmoid(logits).cpu().numpy()

                        # Resize back to original tile size if needed
                        if probs.shape != (tile_height, tile_width):
                            # Use bicubic interpolation for smoother results
                            probs_resized = np.array(
                                Image.fromarray(probs).resize(
                                    (tile_width, tile_height),
                                    Image.Resampling.BICUBIC,
                                )
                            )
                        else:
                            probs_resized = probs

                        # Apply gaussian blur to reduce blockiness
                        try:
                            from scipy.ndimage import gaussian_filter

                            probs_resized = gaussian_filter(
                                probs_resized, sigma=smoothing_sigma
                            )
                        except ImportError:
                            pass  # Continue without smoothing if scipy is not available

                        # Store results in the full arrays
                        # Only store the non-overlapping part (except at edges)
                        valid_x_start = self.overlap if x > 0 else 0
                        valid_y_start = self.overlap if y > 0 else 0
                        valid_x_end = (
                            tile_width - self.overlap
                            if x < n_tiles_x - 1
                            else tile_width
                        )
                        valid_y_end = (
                            tile_height - self.overlap
                            if y < n_tiles_y - 1
                            else tile_height
                        )

                        dest_x_start = x_start + valid_x_start
                        dest_y_start = y_start + valid_y_start
                        dest_x_end = x_start + valid_x_end
                        dest_y_end = y_start + valid_y_end

                        # Store probabilities
                        probabilities[
                            dest_y_start:dest_y_end, dest_x_start:dest_x_end
                        ] = probs_resized[
                            valid_y_start:valid_y_end, valid_x_start:valid_x_end
                        ]

                    except Exception as e:
                        logger.warning(
                            "Error processing tile at (%d, %d): %s", x, y, str(e)
                        )
                        # Continue with next tile

                    # Update progress bar
                    pbar.update(1)

        # Create binary segmentation from probabilities
        segmentation = (probabilities >= threshold).astype(np.float32)

        # Write the output GeoTIFF
        with rasterio.open(output_path, "w", **out_meta) as dst:
            dst.write(segmentation, 1)
            dst.write(probabilities, 2)

            # Add descriptions to bands
            dst.set_band_description(1, "Binary Segmentation")
            dst.set_band_description(2, "Probability Scores")

        logger.info("Segmentation saved to %s", output_path)
        return output_path

segment_image_batch(input_paths, output_dir, text_prompt, threshold=0.5, smoothing_sigma=1.0, suffix='_segmented')

Segment multiple GeoTIFF images using the provided text prompt.

Parameters:

Name	Type	Description	Default
input_paths	list	List of paths to input GeoTIFF files.	required
output_dir	str	Directory where output GeoTIFFs will be saved.	required
text_prompt	str	Text description of what to segment.	required
threshold	float	Threshold for binary segmentation. Defaults to 0.5.	0.5
smoothing_sigma	float	Sigma value for Gaussian smoothing to reduce blockiness. Defaults to 1.0.	1.0
suffix	str	Suffix to add to output filenames. Defaults to "_segmented".	'_segmented'

Returns:

Name	Type	Description
list	List[str]	Paths to all saved output files.

Source code in geoai/segment.py

def segment_image_batch(
    self,
    input_paths: List[str],
    output_dir: str,
    text_prompt: str,
    threshold: float = 0.5,
    smoothing_sigma: float = 1.0,
    suffix: str = "_segmented",
) -> List[str]:
    """
    Segment multiple GeoTIFF images using the provided text prompt.

    Args:
        input_paths (list): List of paths to input GeoTIFF files.
        output_dir (str): Directory where output GeoTIFFs will be saved.
        text_prompt (str): Text description of what to segment.
        threshold (float): Threshold for binary segmentation. Defaults to 0.5.
        smoothing_sigma (float): Sigma value for Gaussian smoothing to reduce blockiness. Defaults to 1.0.
        suffix (str): Suffix to add to output filenames. Defaults to "_segmented".

    Returns:
        list: Paths to all saved output files.
    """
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)

    output_paths = []

    # Process each input file
    for input_path in tqdm(input_paths, desc="Processing files"):
        # Generate output path
        filename = os.path.basename(input_path)
        base_name, ext = os.path.splitext(filename)
        output_path = os.path.join(output_dir, f"{base_name}{suffix}{ext}")

        # Segment the image
        result_path = self.segment_image(
            input_path, output_path, text_prompt, threshold, smoothing_sigma
        )
        output_paths.append(result_path)

    return output_paths

CarDetector

Bases: ObjectDetector

Car detection using a pre-trained Mask R-CNN model.

This class extends the ObjectDetector class with additional methods for car detection.

Source code in geoai/extract.py

class CarDetector(ObjectDetector):
    """
    Car detection using a pre-trained Mask R-CNN model.

    This class extends the `ObjectDetector` class with additional methods for car detection.
    """

    def __init__(
        self,
        model_path: str = "car_detection_usa.pth",
        repo_id: Optional[str] = None,
        model: Optional[Any] = None,
        device: Optional[str] = None,
    ) -> None:
        """
        Initialize the object extractor.

        Args:
            model_path: Path to the .pth model file.
            repo_id: Repo ID for loading models from the Hub.
            model: Custom model to use for inference.
            device: Device to use for inference ('cuda:0', 'cpu', etc.).
        """
        super().__init__(
            model_path=model_path, repo_id=repo_id, model=model, device=device
        )

__init__(model_path='car_detection_usa.pth', repo_id=None, model=None, device=None)

Initialize the object extractor.

Parameters:

Name	Type	Description	Default
model_path	str	Path to the .pth model file.	'car_detection_usa.pth'
repo_id	Optional[str]	Repo ID for loading models from the Hub.	None
model	Optional[Any]	Custom model to use for inference.	None
device	Optional[str]	Device to use for inference ('cuda:0', 'cpu', etc.).	None

Source code in geoai/extract.py

def __init__(
    self,
    model_path: str = "car_detection_usa.pth",
    repo_id: Optional[str] = None,
    model: Optional[Any] = None,
    device: Optional[str] = None,
) -> None:
    """
    Initialize the object extractor.

    Args:
        model_path: Path to the .pth model file.
        repo_id: Repo ID for loading models from the Hub.
        model: Custom model to use for inference.
        device: Device to use for inference ('cuda:0', 'cpu', etc.).
    """
    super().__init__(
        model_path=model_path, repo_id=repo_id, model=model, device=device
    )

CustomDataset

Bases: NonGeoDataset

A TorchGeo dataset for object extraction with overlapping tiles support.

This dataset class creates overlapping image tiles for object detection, ensuring complete coverage of the input raster including right and bottom edges. It inherits from NonGeoDataset to avoid spatial indexing issues.

Attributes:

Name	Type	Description
raster_path		Path to the input raster file.
chip_size		Size of image chips to extract (height, width).
overlap		Amount of overlap between adjacent tiles (0.0-1.0).
transforms		Transforms to apply to the image.
verbose		Whether to print detailed processing information.
stride_x		Horizontal stride between tiles based on overlap.
stride_y		Vertical stride between tiles based on overlap.
row_starts		Starting Y positions for each row of tiles.
col_starts		Starting X positions for each column of tiles.
crs		Coordinate reference system of the raster.
transform		Affine transform of the raster.
height		Height of the raster in pixels.
width		Width of the raster in pixels.
count		Number of bands in the raster.
bounds		Geographic bounds of the raster (west, south, east, north).
roi		Shapely box representing the region of interest.
rows		Number of rows of tiles.
cols		Number of columns of tiles.
raster_stats		Statistics of the raster.

Source code in geoai/extract.py

class CustomDataset(NonGeoDataset):
    """
    A TorchGeo dataset for object extraction with overlapping tiles support.

    This dataset class creates overlapping image tiles for object detection,
    ensuring complete coverage of the input raster including right and bottom edges.
    It inherits from NonGeoDataset to avoid spatial indexing issues.

    Attributes:
        raster_path: Path to the input raster file.
        chip_size: Size of image chips to extract (height, width).
        overlap: Amount of overlap between adjacent tiles (0.0-1.0).
        transforms: Transforms to apply to the image.
        verbose: Whether to print detailed processing information.
        stride_x: Horizontal stride between tiles based on overlap.
        stride_y: Vertical stride between tiles based on overlap.
        row_starts: Starting Y positions for each row of tiles.
        col_starts: Starting X positions for each column of tiles.
        crs: Coordinate reference system of the raster.
        transform: Affine transform of the raster.
        height: Height of the raster in pixels.
        width: Width of the raster in pixels.
        count: Number of bands in the raster.
        bounds: Geographic bounds of the raster (west, south, east, north).
        roi: Shapely box representing the region of interest.
        rows: Number of rows of tiles.
        cols: Number of columns of tiles.
        raster_stats: Statistics of the raster.
    """

    def __init__(
        self,
        raster_path: str,
        chip_size: Tuple[int, int] = (512, 512),
        overlap: float = 0.5,
        transforms: Optional[Any] = None,
        band_indexes: Optional[List[int]] = None,
        verbose: bool = False,
    ) -> None:
        """
        Initialize the dataset with overlapping tiles.

        Args:
            raster_path: Path to the input raster file.
            chip_size: Size of image chips to extract (height, width). Default is (512, 512).
            overlap: Amount of overlap between adjacent tiles (0.0-1.0). Default is 0.5 (50%).
            transforms: Transforms to apply to the image. Default is None.
            band_indexes: List of band indexes to use. Default is None (use all bands).
            verbose: Whether to print detailed processing information. Default is False.

        Raises:
            ValueError: If overlap is too high resulting in non-positive stride.
        """
        super().__init__()

        # Initialize parameters
        self.raster_path = raster_path
        self.chip_size = chip_size
        self.overlap = overlap
        self.transforms = transforms
        self.band_indexes = band_indexes
        self.verbose = verbose
        self.warned_about_bands = False

        # Calculate stride based on overlap
        self.stride_x = int(chip_size[1] * (1 - overlap))
        self.stride_y = int(chip_size[0] * (1 - overlap))

        if self.stride_x <= 0 or self.stride_y <= 0:
            raise ValueError(
                f"Overlap {overlap} is too high, resulting in non-positive stride"
            )

        with rasterio.open(self.raster_path) as src:
            self.crs = src.crs
            self.transform = src.transform
            self.height = src.height
            self.width = src.width
            self.count = src.count

            # Define the bounds of the dataset
            west, south, east, north = src.bounds
            self.bounds = (west, south, east, north)
            self.roi = box(*self.bounds)

            # Calculate starting positions for each tile
            self.row_starts = []
            self.col_starts = []

            # Normal row starts using stride
            for r in range((self.height - 1) // self.stride_y):
                self.row_starts.append(r * self.stride_y)

            # Add a special last row that ensures we reach the bottom edge
            if self.height > self.chip_size[0]:
                self.row_starts.append(max(0, self.height - self.chip_size[0]))
            else:
                # If the image is smaller than chip size, just start at 0
                if not self.row_starts:
                    self.row_starts.append(0)

            # Normal column starts using stride
            for c in range((self.width - 1) // self.stride_x):
                self.col_starts.append(c * self.stride_x)

            # Add a special last column that ensures we reach the right edge
            if self.width > self.chip_size[1]:
                self.col_starts.append(max(0, self.width - self.chip_size[1]))
            else:
                # If the image is smaller than chip size, just start at 0
                if not self.col_starts:
                    self.col_starts.append(0)

            # Update rows and cols based on actual starting positions
            self.rows = len(self.row_starts)
            self.cols = len(self.col_starts)

            print(
                f"Dataset initialized with {self.rows} rows and {self.cols} columns of chips"
            )
            print(f"Image dimensions: {self.width} x {self.height} pixels")
            print(f"Chip size: {self.chip_size[1]} x {self.chip_size[0]} pixels")
            print(
                f"Overlap: {overlap*100}% (stride_x={self.stride_x}, stride_y={self.stride_y})"
            )
            if src.crs:
                print(f"CRS: {src.crs}")

        # Get raster stats
        self.raster_stats = get_raster_stats(raster_path, divide_by=255)

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        """
        Get an image chip from the dataset by index.

        Retrieves an image tile with the specified overlap pattern, ensuring
        proper coverage of the entire raster including edges.

        Args:
            idx: Index of the chip to retrieve.

        Returns:
            dict: Dictionary containing:
                - image: Image tensor.
                - bbox: Geographic bounding box for the window.
                - coords: Pixel coordinates as tensor [i, j].
                - window_size: Window size as tensor [width, height].
        """
        # Convert flat index to grid position
        row = idx // self.cols
        col = idx % self.cols

        # Get pre-calculated starting positions
        j = self.row_starts[row]
        i = self.col_starts[col]

        # Read window from raster
        with rasterio.open(self.raster_path) as src:
            # Make sure we don't read outside the image
            width = min(self.chip_size[1], self.width - i)
            height = min(self.chip_size[0], self.height - j)

            window = Window(i, j, width, height)
            image = src.read(window=window)

            # Handle RGBA or multispectral images - keep only first 3 bands
            if image.shape[0] > 3:
                if not self.warned_about_bands and self.verbose:
                    print(f"Image has {image.shape[0]} bands, using first 3 bands only")
                    self.warned_about_bands = True
                if self.band_indexes is not None:
                    image = image[self.band_indexes]
                else:
                    image = image[:3]
            elif image.shape[0] < 3:
                # If image has fewer than 3 bands, duplicate the last band to make 3
                if not self.warned_about_bands and self.verbose:
                    print(
                        f"Image has {image.shape[0]} bands, duplicating bands to make 3"
                    )
                    self.warned_about_bands = True
                temp = np.zeros((3, image.shape[1], image.shape[2]), dtype=image.dtype)
                for c in range(3):
                    temp[c] = image[min(c, image.shape[0] - 1)]
                image = temp

            # Handle partial windows at edges by padding
            if (
                image.shape[1] != self.chip_size[0]
                or image.shape[2] != self.chip_size[1]
            ):
                temp = np.zeros(
                    (image.shape[0], self.chip_size[0], self.chip_size[1]),
                    dtype=image.dtype,
                )
                temp[:, : image.shape[1], : image.shape[2]] = image
                image = temp

        # Convert to format expected by model (C,H,W)
        image = torch.from_numpy(image).float()

        # Normalize to [0, 1]
        if image.max() > 1:
            image = image / 255.0

        # Apply transforms if any
        if self.transforms is not None:
            image = self.transforms(image)

        # Create geographic bounding box for the window
        minx, miny = self.transform * (i, j + height)
        maxx, maxy = self.transform * (i + width, j)
        bbox = box(minx, miny, maxx, maxy)

        return {
            "image": image,
            "bbox": bbox,
            "coords": torch.tensor([i, j], dtype=torch.long),  # Consistent format
            "window_size": torch.tensor(
                [width, height], dtype=torch.long
            ),  # Consistent format
        }

    def __len__(self) -> int:
        """
        Return the number of samples in the dataset.

        Returns:
            int: Total number of tiles in the dataset.
        """
        return self.rows * self.cols

__getitem__(idx)

Get an image chip from the dataset by index.

Retrieves an image tile with the specified overlap pattern, ensuring proper coverage of the entire raster including edges.

Parameters:

Name	Type	Description	Default
idx	int	Index of the chip to retrieve.	required

Returns:

Name	Type	Description
dict	Dict[str, Any]	Dictionary containing: - image: Image tensor. - bbox: Geographic bounding box for the window. - coords: Pixel coordinates as tensor [i, j]. - window_size: Window size as tensor [width, height].

Source code in geoai/extract.py

def __getitem__(self, idx: int) -> Dict[str, Any]:
    """
    Get an image chip from the dataset by index.

    Retrieves an image tile with the specified overlap pattern, ensuring
    proper coverage of the entire raster including edges.

    Args:
        idx: Index of the chip to retrieve.

    Returns:
        dict: Dictionary containing:
            - image: Image tensor.
            - bbox: Geographic bounding box for the window.
            - coords: Pixel coordinates as tensor [i, j].
            - window_size: Window size as tensor [width, height].
    """
    # Convert flat index to grid position
    row = idx // self.cols
    col = idx % self.cols

    # Get pre-calculated starting positions
    j = self.row_starts[row]
    i = self.col_starts[col]

    # Read window from raster
    with rasterio.open(self.raster_path) as src:
        # Make sure we don't read outside the image
        width = min(self.chip_size[1], self.width - i)
        height = min(self.chip_size[0], self.height - j)

        window = Window(i, j, width, height)
        image = src.read(window=window)

        # Handle RGBA or multispectral images - keep only first 3 bands
        if image.shape[0] > 3:
            if not self.warned_about_bands and self.verbose:
                print(f"Image has {image.shape[0]} bands, using first 3 bands only")
                self.warned_about_bands = True
            if self.band_indexes is not None:
                image = image[self.band_indexes]
            else:
                image = image[:3]
        elif image.shape[0] < 3:
            # If image has fewer than 3 bands, duplicate the last band to make 3
            if not self.warned_about_bands and self.verbose:
                print(
                    f"Image has {image.shape[0]} bands, duplicating bands to make 3"
                )
                self.warned_about_bands = True
            temp = np.zeros((3, image.shape[1], image.shape[2]), dtype=image.dtype)
            for c in range(3):
                temp[c] = image[min(c, image.shape[0] - 1)]
            image = temp

        # Handle partial windows at edges by padding
        if (
            image.shape[1] != self.chip_size[0]
            or image.shape[2] != self.chip_size[1]
        ):
            temp = np.zeros(
                (image.shape[0], self.chip_size[0], self.chip_size[1]),
                dtype=image.dtype,
            )
            temp[:, : image.shape[1], : image.shape[2]] = image
            image = temp

    # Convert to format expected by model (C,H,W)
    image = torch.from_numpy(image).float()

    # Normalize to [0, 1]
    if image.max() > 1:
        image = image / 255.0

    # Apply transforms if any
    if self.transforms is not None:
        image = self.transforms(image)

    # Create geographic bounding box for the window
    minx, miny = self.transform * (i, j + height)
    maxx, maxy = self.transform * (i + width, j)
    bbox = box(minx, miny, maxx, maxy)

    return {
        "image": image,
        "bbox": bbox,
        "coords": torch.tensor([i, j], dtype=torch.long),  # Consistent format
        "window_size": torch.tensor(
            [width, height], dtype=torch.long
        ),  # Consistent format
    }

__init__(raster_path, chip_size=(512, 512), overlap=0.5, transforms=None, band_indexes=None, verbose=False)

Initialize the dataset with overlapping tiles.

Parameters:

Name	Type	Description	Default
raster_path	str	Path to the input raster file.	required
chip_size	Tuple[int, int]	Size of image chips to extract (height, width). Default is (512, 512).	(512, 512)
overlap	float	Amount of overlap between adjacent tiles (0.0-1.0). Default is 0.5 (50%).	0.5
transforms	Optional[Any]	Transforms to apply to the image. Default is None.	None
band_indexes	Optional[List[int]]	List of band indexes to use. Default is None (use all bands).	None
verbose	bool	Whether to print detailed processing information. Default is False.	False

Raises:

Type	Description
ValueError	If overlap is too high resulting in non-positive stride.

Source code in geoai/extract.py

def __init__(
    self,
    raster_path: str,
    chip_size: Tuple[int, int] = (512, 512),
    overlap: float = 0.5,
    transforms: Optional[Any] = None,
    band_indexes: Optional[List[int]] = None,
    verbose: bool = False,
) -> None:
    """
    Initialize the dataset with overlapping tiles.

    Args:
        raster_path: Path to the input raster file.
        chip_size: Size of image chips to extract (height, width). Default is (512, 512).
        overlap: Amount of overlap between adjacent tiles (0.0-1.0). Default is 0.5 (50%).
        transforms: Transforms to apply to the image. Default is None.
        band_indexes: List of band indexes to use. Default is None (use all bands).
        verbose: Whether to print detailed processing information. Default is False.

    Raises:
        ValueError: If overlap is too high resulting in non-positive stride.
    """
    super().__init__()

    # Initialize parameters
    self.raster_path = raster_path
    self.chip_size = chip_size
    self.overlap = overlap
    self.transforms = transforms
    self.band_indexes = band_indexes
    self.verbose = verbose
    self.warned_about_bands = False

    # Calculate stride based on overlap
    self.stride_x = int(chip_size[1] * (1 - overlap))
    self.stride_y = int(chip_size[0] * (1 - overlap))

    if self.stride_x <= 0 or self.stride_y <= 0:
        raise ValueError(
            f"Overlap {overlap} is too high, resulting in non-positive stride"
        )

    with rasterio.open(self.raster_path) as src:
        self.crs = src.crs
        self.transform = src.transform
        self.height = src.height
        self.width = src.width
        self.count = src.count

        # Define the bounds of the dataset
        west, south, east, north = src.bounds
        self.bounds = (west, south, east, north)
        self.roi = box(*self.bounds)

        # Calculate starting positions for each tile
        self.row_starts = []
        self.col_starts = []

        # Normal row starts using stride
        for r in range((self.height - 1) // self.stride_y):
            self.row_starts.append(r * self.stride_y)

        # Add a special last row that ensures we reach the bottom edge
        if self.height > self.chip_size[0]:
            self.row_starts.append(max(0, self.height - self.chip_size[0]))
        else:
            # If the image is smaller than chip size, just start at 0
            if not self.row_starts:
                self.row_starts.append(0)

        # Normal column starts using stride
        for c in range((self.width - 1) // self.stride_x):
            self.col_starts.append(c * self.stride_x)

        # Add a special last column that ensures we reach the right edge
        if self.width > self.chip_size[1]:
            self.col_starts.append(max(0, self.width - self.chip_size[1]))
        else:
            # If the image is smaller than chip size, just start at 0
            if not self.col_starts:
                self.col_starts.append(0)

        # Update rows and cols based on actual starting positions
        self.rows = len(self.row_starts)
        self.cols = len(self.col_starts)

        print(
            f"Dataset initialized with {self.rows} rows and {self.cols} columns of chips"
        )
        print(f"Image dimensions: {self.width} x {self.height} pixels")
        print(f"Chip size: {self.chip_size[1]} x {self.chip_size[0]} pixels")
        print(
            f"Overlap: {overlap*100}% (stride_x={self.stride_x}, stride_y={self.stride_y})"
        )
        if src.crs:
            print(f"CRS: {src.crs}")

    # Get raster stats
    self.raster_stats = get_raster_stats(raster_path, divide_by=255)

__len__()

Return the number of samples in the dataset.

Returns:

Name	Type	Description
int	int	Total number of tiles in the dataset.

Source code in geoai/extract.py

def __len__(self) -> int:
    """
    Return the number of samples in the dataset.

    Returns:
        int: Total number of tiles in the dataset.
    """
    return self.rows * self.cols

DetectionResult dataclass

Represents a detection result with score, label, bounding box, and optional mask.

Source code in geoai/segment.py

@dataclass
class DetectionResult:
    """Represents a detection result with score, label, bounding box, and optional mask."""

    score: float
    label: str
    box: BoundingBox
    mask: Optional[np.array] = None

    @classmethod
    def from_dict(cls, detection_dict: Dict) -> "DetectionResult":
        return cls(
            score=detection_dict["score"],
            label=detection_dict["label"],
            box=BoundingBox(
                xmin=detection_dict["box"]["xmin"],
                ymin=detection_dict["box"]["ymin"],
                xmax=detection_dict["box"]["xmax"],
                ymax=detection_dict["box"]["ymax"],
            ),
        )

FocalLoss

Bases: Module

Focal Loss for addressing class imbalance in segmentation.

Reference: Lin, T. Y., Goyal, P., Girshick, R., He, K., & Dollár, P. (2017). Focal loss for dense object detection. ICCV.

Parameters:

Name	Type	Description	Default
alpha		Weighting factor in range (0,1) to balance positive/negative examples	1.0
gamma		Exponent of the modulating factor (1 - p_t)^gamma	2.0
ignore_index		Specifies a target value that is ignored	-100
reduction		Specifies the reduction to apply to the output	'mean'
weight		Manual rescaling weight given to each class	None

Source code in geoai/landcover_train.py

class FocalLoss(nn.Module):
    """
    Focal Loss for addressing class imbalance in segmentation.

    Reference: Lin, T. Y., Goyal, P., Girshick, R., He, K., & Dollár, P. (2017).
    Focal loss for dense object detection. ICCV.

    Args:
        alpha: Weighting factor in range (0,1) to balance positive/negative examples
        gamma: Exponent of the modulating factor (1 - p_t)^gamma
        ignore_index: Specifies a target value that is ignored
        reduction: Specifies the reduction to apply to the output
        weight: Manual rescaling weight given to each class
    """

    def __init__(
        self, alpha=1.0, gamma=2.0, ignore_index=-100, reduction="mean", weight=None
    ):
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.ignore_index = ignore_index
        self.reduction = reduction
        self.weight = weight

    def forward(self, inputs, targets):
        """
        Forward pass of focal loss.

        Args:
            inputs: Predictions (N, C, H, W) where C = number of classes
            targets: Ground truth (N, H, W) with class indices

        Returns:
            Loss value
        """
        # Get class probabilities
        ce_loss = F.cross_entropy(
            inputs,
            targets,
            weight=self.weight,
            ignore_index=self.ignore_index,
            reduction="none",
        )

        # Get probability of true class
        p_t = torch.exp(-ce_loss)

        # Calculate focal loss
        focal_loss = self.alpha * (1 - p_t) ** self.gamma * ce_loss

        # Apply reduction
        if self.reduction == "mean":
            return focal_loss.mean()
        elif self.reduction == "sum":
            return focal_loss.sum()
        else:
            return focal_loss

forward(inputs, targets)

Forward pass of focal loss.

Parameters:

Name	Type	Description	Default
inputs		Predictions (N, C, H, W) where C = number of classes	required
targets		Ground truth (N, H, W) with class indices	required

Returns:

Type	Description
	Loss value

Source code in geoai/landcover_train.py

def forward(self, inputs, targets):
    """
    Forward pass of focal loss.

    Args:
        inputs: Predictions (N, C, H, W) where C = number of classes
        targets: Ground truth (N, H, W) with class indices

    Returns:
        Loss value
    """
    # Get class probabilities
    ce_loss = F.cross_entropy(
        inputs,
        targets,
        weight=self.weight,
        ignore_index=self.ignore_index,
        reduction="none",
    )

    # Get probability of true class
    p_t = torch.exp(-ce_loss)

    # Calculate focal loss
    focal_loss = self.alpha * (1 - p_t) ** self.gamma * ce_loss

    # Apply reduction
    if self.reduction == "mean":
        return focal_loss.mean()
    elif self.reduction == "sum":
        return focal_loss.sum()
    else:
        return focal_loss

GroundedSAM

A class for segmenting remote sensing imagery using text prompts with Grounding DINO + SAM.

This class combines Grounding DINO for object detection and Segment Anything Model (SAM) for precise segmentation based on text prompts. It can process large GeoTIFF files by tiling them and handles proper georeferencing in the outputs.

Parameters:

Name	Type	Description	Default
detector_id	str	Hugging Face model ID for Grounding DINO. Defaults to "IDEA-Research/grounding-dino-tiny".	'IDEA-Research/grounding-dino-tiny'
segmenter_id	str	Hugging Face model ID for SAM. Defaults to "facebook/sam-vit-base".	'facebook/sam-vit-base'
device	str	Device to run the models on ('cuda', 'cpu'). If None, will use CUDA if available.	None
tile_size	int	Size of tiles to process the image in chunks. Defaults to 1024.	1024
overlap	int	Overlap between tiles to avoid edge artifacts. Defaults to 128.	128
threshold	float	Detection threshold for Grounding DINO. Defaults to 0.3.	0.3

Attributes:

Name	Type	Description
detector_id	str	The Grounding DINO model ID.
segmenter_id	str	The SAM model ID.
device	str	The device being used ('cuda' or 'cpu').
tile_size	int	Size of tiles for processing.
overlap	int	Overlap between tiles.
threshold	float	Detection threshold.
object_detector	float	The Grounding DINO pipeline.
segmentator	float	The SAM model.
processor	float	The SAM processor.

Source code in geoai/segment.py

class GroundedSAM:
    """
    A class for segmenting remote sensing imagery using text prompts with Grounding DINO + SAM.

    This class combines Grounding DINO for object detection and Segment Anything Model (SAM) for
    precise segmentation based on text prompts. It can process large GeoTIFF files by tiling them
    and handles proper georeferencing in the outputs.

    Args:
        detector_id (str): Hugging Face model ID for Grounding DINO. Defaults to "IDEA-Research/grounding-dino-tiny".
        segmenter_id (str): Hugging Face model ID for SAM. Defaults to "facebook/sam-vit-base".
        device (str): Device to run the models on ('cuda', 'cpu'). If None, will use CUDA if available.
        tile_size (int): Size of tiles to process the image in chunks. Defaults to 1024.
        overlap (int): Overlap between tiles to avoid edge artifacts. Defaults to 128.
        threshold (float): Detection threshold for Grounding DINO. Defaults to 0.3.

    Attributes:
        detector_id (str): The Grounding DINO model ID.
        segmenter_id (str): The SAM model ID.
        device (str): The device being used ('cuda' or 'cpu').
        tile_size (int): Size of tiles for processing.
        overlap (int): Overlap between tiles.
        threshold (float): Detection threshold.
        object_detector: The Grounding DINO pipeline.
        segmentator: The SAM model.
        processor: The SAM processor.
    """

    def __init__(
        self,
        detector_id: str = "IDEA-Research/grounding-dino-tiny",
        segmenter_id: str = "facebook/sam-vit-base",
        device: Optional[str] = None,
        tile_size: int = 1024,
        overlap: int = 128,
        threshold: float = 0.3,
    ) -> None:
        """
        Initialize the GroundedSAM with the specified models and settings.

        Args:
            detector_id (str): Hugging Face model ID for Grounding DINO.
            segmenter_id (str): Hugging Face model ID for SAM.
            device (str): Device to run the models on ('cuda', 'cpu').
            tile_size (int): Size of tiles to process the image in chunks.
            overlap (int): Overlap between tiles to avoid edge artifacts.
            threshold (float): Detection threshold for Grounding DINO.
        """
        self.detector_id = detector_id
        self.segmenter_id = segmenter_id
        self.tile_size = tile_size
        self.overlap = overlap
        self.threshold = threshold

        # Set device
        if device is None:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device

        # Load models
        self._load_models()

        print(f"GroundedSAM initialized on {self.device}")

    def _load_models(self) -> None:
        """Load the Grounding DINO and SAM models."""
        # Load Grounding DINO
        self.object_detector = pipeline(
            model=self.detector_id,
            task="zero-shot-object-detection",
            device=self.device,
        )

        # Load SAM
        self.segmentator = AutoModelForMaskGeneration.from_pretrained(
            self.segmenter_id
        ).to(self.device)
        self.processor = AutoProcessor.from_pretrained(self.segmenter_id)

    def _detect(self, image: Image.Image, labels: List[str]) -> List[DetectionResult]:
        """
        Use Grounding DINO to detect objects in an image.

        Args:
            image (Image.Image): PIL image to detect objects in.
            labels (List[str]): List of text labels to detect.

        Returns:
            List[DetectionResult]: List of detection results.
        """
        # Ensure labels end with periods
        labels = [label if label.endswith(".") else label + "." for label in labels]

        results = self.object_detector(
            image, candidate_labels=labels, threshold=self.threshold
        )
        results = [DetectionResult.from_dict(result) for result in results]

        return results

    def _apply_nms(
        self, detections: List[DetectionResult], iou_threshold: float = 0.5
    ) -> List[DetectionResult]:
        """
        Apply Non-Maximum Suppression to remove overlapping detections.

        Args:
            detections (List[DetectionResult]): List of detection results.
            iou_threshold (float): IoU threshold for NMS.

        Returns:
            List[DetectionResult]: Filtered detection results.
        """
        import cv2  # Lazy import to avoid QGIS opencv conflicts

        if not detections:
            return detections

        # Convert to format for NMS
        boxes = []
        scores = []

        for detection in detections:
            boxes.append(
                [
                    detection.box.xmin,
                    detection.box.ymin,
                    detection.box.xmax,
                    detection.box.ymax,
                ]
            )
            scores.append(detection.score)

        boxes = np.array(boxes, dtype=np.float32)
        scores = np.array(scores, dtype=np.float32)

        # Apply NMS using OpenCV
        indices = cv2.dnn.NMSBoxes(boxes, scores, self.threshold, iou_threshold)

        if len(indices) > 0:
            indices = indices.flatten()
            return [detections[i] for i in indices]
        else:
            return []

    def _get_boxes(self, results: List[DetectionResult]) -> List[List[List[float]]]:
        """Extract bounding boxes from detection results."""
        boxes = []
        for result in results:
            xyxy = result.box.xyxy
            boxes.append(xyxy)
        return [boxes]

    def _refine_masks(
        self, masks: torch.BoolTensor, polygon_refinement: bool = False
    ) -> List[np.ndarray]:
        """Refine masks from SAM output."""
        masks = masks.cpu().float()
        masks = masks.permute(0, 2, 3, 1)
        masks = masks.mean(axis=-1)
        masks = (masks > 0).int()
        masks = masks.numpy().astype(np.uint8)
        masks = list(masks)

        if polygon_refinement:
            for idx, mask in enumerate(masks):
                shape = mask.shape
                polygon = self._mask_to_polygon(mask)
                if polygon:
                    mask = self._polygon_to_mask(polygon, shape)
                    masks[idx] = mask

        return masks

    def _mask_to_polygon(self, mask: np.ndarray) -> List[List[int]]:
        """Convert mask to polygon coordinates."""
        import cv2  # Lazy import to avoid QGIS opencv conflicts

        # Find contours in the binary mask
        contours, _ = cv2.findContours(
            mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )

        if not contours:
            return []

        # Find the contour with the largest area
        largest_contour = max(contours, key=cv2.contourArea)

        # Extract the vertices of the contour
        polygon = largest_contour.reshape(-1, 2).tolist()

        return polygon

    def _polygon_to_mask(
        self, polygon: List[Tuple[int, int]], image_shape: Tuple[int, int]
    ) -> np.ndarray:
        """Convert polygon to mask."""
        import cv2  # Lazy import to avoid QGIS opencv conflicts

        # Create an empty mask
        mask = np.zeros(image_shape, dtype=np.uint8)

        # Convert polygon to an array of points
        pts = np.array(polygon, dtype=np.int32)

        # Fill the polygon with white color (255)
        cv2.fillPoly(mask, [pts], color=(255,))

        return mask

    def _separate_instances(
        self, mask: np.ndarray, min_area: int = 50
    ) -> List[np.ndarray]:
        """
        Separate individual instances from a combined mask using connected components.

        Args:
            mask (np.ndarray): Combined binary mask.
            min_area (int): Minimum area threshold for valid instances.

        Returns:
            List[np.ndarray]: List of individual instance masks.
        """
        import cv2  # Lazy import to avoid QGIS opencv conflicts

        # Find connected components
        num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(
            mask.astype(np.uint8), connectivity=8
        )

        instances = []
        for i in range(1, num_labels):  # Skip background (label 0)
            # Get area of the component
            area = stats[i, cv2.CC_STAT_AREA]

            # Filter by minimum area
            if area >= min_area:
                # Create mask for this instance
                instance_mask = (labels == i).astype(np.uint8) * 255
                instances.append(instance_mask)

        return instances

    def _mask_to_polygons(
        self,
        mask: np.ndarray,
        transform,
        x_offset: int = 0,
        y_offset: int = 0,
        min_area: int = 50,
        simplify_tolerance: float = 1.0,
    ) -> List[Dict]:
        """
        Convert mask to individual polygons with geospatial coordinates.

        Args:
            mask (np.ndarray): Binary mask.
            transform: Rasterio transform object.
            x_offset (int): X offset for tile position.
            y_offset (int): Y offset for tile position.
            min_area (int): Minimum area threshold for valid polygons.
            simplify_tolerance (float): Tolerance for polygon simplification.

        Returns:
            List[Dict]: List of polygon dictionaries with geometry and properties.
        """
        import cv2  # Lazy import to avoid QGIS opencv conflicts

        polygons = []

        # Get individual instances
        instances = self._separate_instances(mask, min_area)

        for instance_mask in instances:
            # Find contours
            contours, _ = cv2.findContours(
                instance_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
            )

            for contour in contours:
                # Filter by minimum area
                area = cv2.contourArea(contour)
                if area < min_area:
                    continue

                # Simplify contour
                epsilon = simplify_tolerance
                simplified_contour = cv2.approxPolyDP(contour, epsilon, True)

                # Convert to pixel coordinates (add offsets)
                pixel_coords = simplified_contour.reshape(-1, 2)
                pixel_coords = pixel_coords + [x_offset, y_offset]

                # Convert to geographic coordinates
                geo_coords = []
                for x, y in pixel_coords:
                    geo_x, geo_y = transform * (x, y)
                    geo_coords.append([geo_x, geo_y])

                # Close the polygon if needed
                if len(geo_coords) > 2:
                    if geo_coords[0] != geo_coords[-1]:
                        geo_coords.append(geo_coords[0])

                    # Create Shapely polygon
                    try:
                        polygon = Polygon(geo_coords)
                        if polygon.is_valid and polygon.area > 0:
                            polygons.append({"geometry": polygon, "area_pixels": area})
                    except Exception as e:
                        print(f"Error creating polygon: {e}")
                        continue

        return polygons

    def _segment(
        self,
        image: Image.Image,
        detection_results: List[DetectionResult],
        polygon_refinement: bool = False,
    ) -> List[DetectionResult]:
        """
        Use SAM to generate masks for detected objects.

        Args:
            image (Image.Image): PIL image.
            detection_results (List[DetectionResult]): Detection results from Grounding DINO.
            polygon_refinement (bool): Whether to refine masks using polygon fitting.

        Returns:
            List[DetectionResult]: Detection results with masks.
        """
        if not detection_results:
            return detection_results

        boxes = self._get_boxes(detection_results)
        inputs = self.processor(
            images=image, input_boxes=boxes, return_tensors="pt"
        ).to(self.device)

        outputs = self.segmentator(**inputs)
        masks = self.processor.post_process_masks(
            masks=outputs.pred_masks,
            original_sizes=inputs.original_sizes,
            reshaped_input_sizes=inputs.reshaped_input_sizes,
        )[0]

        masks = self._refine_masks(masks, polygon_refinement)

        for detection_result, mask in zip(detection_results, masks):
            detection_result.mask = mask

        return detection_results

    def segment_image(
        self,
        input_path: str,
        output_path: str,
        text_prompts: Union[str, List[str]],
        polygon_refinement: bool = False,
        export_boxes: bool = False,
        export_polygons: bool = True,
        smoothing_sigma: float = 1.0,
        nms_threshold: float = 0.5,
        min_polygon_area: int = 50,
        simplify_tolerance: float = 2.0,
    ) -> str:
        """
        Segment a GeoTIFF image using text prompts with improved instance segmentation.

        Args:
            input_path (str): Path to the input GeoTIFF file.
            output_path (str): Path where the output GeoTIFF will be saved.
            text_prompts (Union[str, List[str]]): Text prompt(s) describing what to segment.
            polygon_refinement (bool): Whether to refine masks using polygon fitting.
            export_boxes (bool): Whether to export bounding boxes as a separate vector file.
            export_polygons (bool): Whether to export segmentation polygons as vector file.
            smoothing_sigma (float): Sigma value for Gaussian smoothing to reduce blockiness.
            nms_threshold (float): Non-maximum suppression threshold for removing overlapping detections.
            min_polygon_area (int): Minimum area in pixels for valid polygons.
            simplify_tolerance (float): Tolerance for polygon simplification.

        Returns:
            Dict: Dictionary containing paths to output files.
        """
        if isinstance(text_prompts, str):
            text_prompts = [text_prompts]

        # Open the input GeoTIFF
        with rasterio.open(input_path) as src:
            # Get metadata
            meta = src.meta
            height = src.height
            width = src.width
            transform = src.transform
            crs = src.crs

            # Create output metadata for segmentation masks
            out_meta = meta.copy()
            out_meta.update(
                {"count": len(text_prompts) + 1, "dtype": "uint8", "nodata": 0}
            )

            # Create arrays for results
            all_masks = np.zeros((len(text_prompts), height, width), dtype=np.uint8)
            all_boxes = []
            all_polygons = []

            # Calculate effective tile size (accounting for overlap)
            effective_tile_size = self.tile_size - 2 * self.overlap

            # Calculate number of tiles
            n_tiles_x = max(1, int(np.ceil(width / effective_tile_size)))
            n_tiles_y = max(1, int(np.ceil(height / effective_tile_size)))
            total_tiles = n_tiles_x * n_tiles_y

            print(f"Processing {total_tiles} tiles ({n_tiles_x}x{n_tiles_y})")

            # Process tiles with tqdm progress bar
            with tqdm(total=total_tiles, desc="Processing tiles") as pbar:
                # Iterate through tiles
                for y in range(n_tiles_y):
                    for x in range(n_tiles_x):
                        # Calculate tile coordinates with overlap
                        x_start = max(0, x * effective_tile_size - self.overlap)
                        y_start = max(0, y * effective_tile_size - self.overlap)
                        x_end = min(width, (x + 1) * effective_tile_size + self.overlap)
                        y_end = min(
                            height, (y + 1) * effective_tile_size + self.overlap
                        )

                        tile_width = x_end - x_start
                        tile_height = y_end - y_start

                        # Read the tile
                        window = Window(x_start, y_start, tile_width, tile_height)
                        tile_data = src.read(window=window)

                        # Process the tile
                        try:
                            # Convert to RGB format for processing
                            if tile_data.shape[0] >= 3:
                                # Use first three bands for RGB representation
                                rgb_tile = tile_data[:3].transpose(1, 2, 0)
                            elif tile_data.shape[0] == 1:
                                # Create RGB from grayscale
                                rgb_tile = np.repeat(
                                    tile_data[0][:, :, np.newaxis], 3, axis=2
                                )
                            else:
                                print(
                                    f"Unsupported number of bands: {tile_data.shape[0]}"
                                )
                                continue

                            # Normalize to 0-255 range if needed
                            if rgb_tile.max() > 0:
                                rgb_tile = (
                                    (rgb_tile - rgb_tile.min())
                                    / (rgb_tile.max() - rgb_tile.min())
                                    * 255
                                ).astype(np.uint8)

                            # Convert to PIL Image
                            pil_image = Image.fromarray(rgb_tile)

                            # Detect objects
                            detections = self._detect(pil_image, text_prompts)

                            if detections:
                                # Apply Non-Maximum Suppression to reduce overlapping detections
                                detections = self._apply_nms(detections, nms_threshold)

                                if detections:
                                    # Segment objects
                                    detections = self._segment(
                                        pil_image, detections, polygon_refinement
                                    )

                                    # Process results
                                    for i, prompt in enumerate(text_prompts):
                                        prompt_polygons = []
                                        prompt_mask = np.zeros(
                                            (tile_height, tile_width), dtype=np.uint8
                                        )

                                        for detection in detections:
                                            if (
                                                detection.label.replace(".", "")
                                                .strip()
                                                .lower()
                                                == prompt.lower()
                                            ):
                                                if detection.mask is not None:
                                                    # Apply gaussian blur to reduce blockiness
                                                    try:
                                                        from scipy.ndimage import (
                                                            gaussian_filter,
                                                        )

                                                        smoothed_mask = gaussian_filter(
                                                            detection.mask.astype(
                                                                float
                                                            ),
                                                            sigma=smoothing_sigma,
                                                        )
                                                        detection.mask = (
                                                            smoothed_mask > 0.5
                                                        ).astype(np.uint8)
                                                    except ImportError:
                                                        pass

                                                    # Add to combined mask for this prompt
                                                    prompt_mask = np.maximum(
                                                        prompt_mask, detection.mask
                                                    )

                                                # Store bounding box with geospatial coordinates
                                                if export_boxes:
                                                    bbox = detection.box
                                                    x_geo_min, y_geo_min = transform * (
                                                        x_start + bbox.xmin,
                                                        y_start + bbox.ymin,
                                                    )
                                                    x_geo_max, y_geo_max = transform * (
                                                        x_start + bbox.xmax,
                                                        y_start + bbox.ymax,
                                                    )

                                                    geo_box = {
                                                        "label": detection.label,
                                                        "score": detection.score,
                                                        "prompt": prompt,
                                                        "geometry": box(
                                                            x_geo_min,
                                                            y_geo_max,
                                                            x_geo_max,
                                                            y_geo_min,
                                                        ),
                                                    }
                                                    all_boxes.append(geo_box)

                                        # Convert masks to individual polygons
                                        if export_polygons and np.any(prompt_mask):
                                            tile_polygons = self._mask_to_polygons(
                                                prompt_mask,
                                                transform,
                                                x_start,
                                                y_start,
                                                min_polygon_area,
                                                simplify_tolerance,
                                            )

                                            # Add metadata to polygons
                                            for poly_data in tile_polygons:
                                                poly_data.update(
                                                    {
                                                        "label": prompt,
                                                        "score": max(
                                                            [
                                                                d.score
                                                                for d in detections
                                                                if d.label.replace(
                                                                    ".", ""
                                                                )
                                                                .strip()
                                                                .lower()
                                                                == prompt.lower()
                                                            ],
                                                            default=0.0,
                                                        ),
                                                        "tile_x": x,
                                                        "tile_y": y,
                                                    }
                                                )
                                                all_polygons.append(poly_data)

                                        # Store mask in the global array
                                        valid_x_start = self.overlap if x > 0 else 0
                                        valid_y_start = self.overlap if y > 0 else 0
                                        valid_x_end = (
                                            tile_width - self.overlap
                                            if x < n_tiles_x - 1
                                            else tile_width
                                        )
                                        valid_y_end = (
                                            tile_height - self.overlap
                                            if y < n_tiles_y - 1
                                            else tile_height
                                        )

                                        dest_x_start = x_start + valid_x_start
                                        dest_y_start = y_start + valid_y_start
                                        dest_x_end = x_start + valid_x_end
                                        dest_y_end = y_start + valid_y_end

                                        mask_slice = prompt_mask[
                                            valid_y_start:valid_y_end,
                                            valid_x_start:valid_x_end,
                                        ]
                                        all_masks[
                                            i,
                                            dest_y_start:dest_y_end,
                                            dest_x_start:dest_x_end,
                                        ] = np.maximum(
                                            all_masks[
                                                i,
                                                dest_y_start:dest_y_end,
                                                dest_x_start:dest_x_end,
                                            ],
                                            mask_slice,
                                        )

                        except Exception as e:
                            print(f"Error processing tile at ({x}, {y}): {str(e)}")
                            continue

                        # Update progress bar
                        pbar.update(1)

            # Create combined mask (union of all individual masks)
            combined_mask = np.any(all_masks, axis=0).astype(np.uint8)

            # Write the output GeoTIFF
            with rasterio.open(output_path, "w", **out_meta) as dst:
                # Write combined mask as first band
                dst.write(combined_mask, 1)

                # Write individual masks for each prompt
                for i, mask in enumerate(all_masks):
                    dst.write(mask, i + 2)

                # Add descriptions to bands
                dst.set_band_description(1, "Combined Segmentation")
                for i, prompt in enumerate(text_prompts):
                    dst.set_band_description(i + 2, f"Segmentation: {prompt}")

            result_files = {"segmentation": output_path}

            # Export bounding boxes if requested
            if export_boxes and all_boxes:
                boxes_path = output_path.replace(".tif", "_boxes.geojson")
                gdf = gpd.GeoDataFrame(all_boxes, crs=crs)
                gdf.to_file(boxes_path, driver="GeoJSON")
                result_files["boxes"] = boxes_path
                print(f"Exported {len(all_boxes)} bounding boxes to {boxes_path}")

            # Export instance polygons if requested
            if export_polygons and all_polygons:
                polygons_path = output_path.replace(".tif", "_polygons.geojson")
                gdf = gpd.GeoDataFrame(all_polygons, crs=crs)
                gdf.to_file(polygons_path, driver="GeoJSON")
                result_files["polygons"] = polygons_path
                print(
                    f"Exported {len(all_polygons)} instance polygons to {polygons_path}"
                )

            print(f"Segmentation saved to {output_path}")
            print(
                f"Found {len(all_polygons)} individual building instances"
                if export_polygons
                else ""
            )

            return result_files

__init__(detector_id='IDEA-Research/grounding-dino-tiny', segmenter_id='facebook/sam-vit-base', device=None, tile_size=1024, overlap=128, threshold=0.3)

Initialize the GroundedSAM with the specified models and settings.

Parameters:

Name	Type	Description	Default
detector_id	str	Hugging Face model ID for Grounding DINO.	'IDEA-Research/grounding-dino-tiny'
segmenter_id	str	Hugging Face model ID for SAM.	'facebook/sam-vit-base'
device	str	Device to run the models on ('cuda', 'cpu').	None
tile_size	int	Size of tiles to process the image in chunks.	1024
overlap	int	Overlap between tiles to avoid edge artifacts.	128
threshold	float	Detection threshold for Grounding DINO.	0.3

Source code in geoai/segment.py

def __init__(
    self,
    detector_id: str = "IDEA-Research/grounding-dino-tiny",
    segmenter_id: str = "facebook/sam-vit-base",
    device: Optional[str] = None,
    tile_size: int = 1024,
    overlap: int = 128,
    threshold: float = 0.3,
) -> None:
    """
    Initialize the GroundedSAM with the specified models and settings.

    Args:
        detector_id (str): Hugging Face model ID for Grounding DINO.
        segmenter_id (str): Hugging Face model ID for SAM.
        device (str): Device to run the models on ('cuda', 'cpu').
        tile_size (int): Size of tiles to process the image in chunks.
        overlap (int): Overlap between tiles to avoid edge artifacts.
        threshold (float): Detection threshold for Grounding DINO.
    """
    self.detector_id = detector_id
    self.segmenter_id = segmenter_id
    self.tile_size = tile_size
    self.overlap = overlap
    self.threshold = threshold

    # Set device
    if device is None:
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
    else:
        self.device = device

    # Load models
    self._load_models()

    print(f"GroundedSAM initialized on {self.device}")

segment_image(input_path, output_path, text_prompts, polygon_refinement=False, export_boxes=False, export_polygons=True, smoothing_sigma=1.0, nms_threshold=0.5, min_polygon_area=50, simplify_tolerance=2.0)

Segment a GeoTIFF image using text prompts with improved instance segmentation.

Parameters:

Name	Type	Description	Default
input_path	str	Path to the input GeoTIFF file.	required
output_path	str	Path where the output GeoTIFF will be saved.	required
text_prompts	Union[str, List[str]]	Text prompt(s) describing what to segment.	required
polygon_refinement	bool	Whether to refine masks using polygon fitting.	False
export_boxes	bool	Whether to export bounding boxes as a separate vector file.	False
export_polygons	bool	Whether to export segmentation polygons as vector file.	True
smoothing_sigma	float	Sigma value for Gaussian smoothing to reduce blockiness.	1.0
nms_threshold	float	Non-maximum suppression threshold for removing overlapping detections.	0.5
min_polygon_area	int	Minimum area in pixels for valid polygons.	50
simplify_tolerance	float	Tolerance for polygon simplification.	2.0

Returns:

Name	Type	Description
Dict	str	Dictionary containing paths to output files.

Source code in geoai/segment.py

def segment_image(
    self,
    input_path: str,
    output_path: str,
    text_prompts: Union[str, List[str]],
    polygon_refinement: bool = False,
    export_boxes: bool = False,
    export_polygons: bool = True,
    smoothing_sigma: float = 1.0,
    nms_threshold: float = 0.5,
    min_polygon_area: int = 50,
    simplify_tolerance: float = 2.0,
) -> str:
    """
    Segment a GeoTIFF image using text prompts with improved instance segmentation.

    Args:
        input_path (str): Path to the input GeoTIFF file.
        output_path (str): Path where the output GeoTIFF will be saved.
        text_prompts (Union[str, List[str]]): Text prompt(s) describing what to segment.
        polygon_refinement (bool): Whether to refine masks using polygon fitting.
        export_boxes (bool): Whether to export bounding boxes as a separate vector file.
        export_polygons (bool): Whether to export segmentation polygons as vector file.
        smoothing_sigma (float): Sigma value for Gaussian smoothing to reduce blockiness.
        nms_threshold (float): Non-maximum suppression threshold for removing overlapping detections.
        min_polygon_area (int): Minimum area in pixels for valid polygons.
        simplify_tolerance (float): Tolerance for polygon simplification.

    Returns:
        Dict: Dictionary containing paths to output files.
    """
    if isinstance(text_prompts, str):
        text_prompts = [text_prompts]

    # Open the input GeoTIFF
    with rasterio.open(input_path) as src:
        # Get metadata
        meta = src.meta
        height = src.height
        width = src.width
        transform = src.transform
        crs = src.crs

        # Create output metadata for segmentation masks
        out_meta = meta.copy()
        out_meta.update(
            {"count": len(text_prompts) + 1, "dtype": "uint8", "nodata": 0}
        )

        # Create arrays for results
        all_masks = np.zeros((len(text_prompts), height, width), dtype=np.uint8)
        all_boxes = []
        all_polygons = []

        # Calculate effective tile size (accounting for overlap)
        effective_tile_size = self.tile_size - 2 * self.overlap

        # Calculate number of tiles
        n_tiles_x = max(1, int(np.ceil(width / effective_tile_size)))
        n_tiles_y = max(1, int(np.ceil(height / effective_tile_size)))
        total_tiles = n_tiles_x * n_tiles_y

        print(f"Processing {total_tiles} tiles ({n_tiles_x}x{n_tiles_y})")

        # Process tiles with tqdm progress bar
        with tqdm(total=total_tiles, desc="Processing tiles") as pbar:
            # Iterate through tiles
            for y in range(n_tiles_y):
                for x in range(n_tiles_x):
                    # Calculate tile coordinates with overlap
                    x_start = max(0, x * effective_tile_size - self.overlap)
                    y_start = max(0, y * effective_tile_size - self.overlap)
                    x_end = min(width, (x + 1) * effective_tile_size + self.overlap)
                    y_end = min(
                        height, (y + 1) * effective_tile_size + self.overlap
                    )

                    tile_width = x_end - x_start
                    tile_height = y_end - y_start

                    # Read the tile
                    window = Window(x_start, y_start, tile_width, tile_height)
                    tile_data = src.read(window=window)

                    # Process the tile
                    try:
                        # Convert to RGB format for processing
                        if tile_data.shape[0] >= 3:
                            # Use first three bands for RGB representation
                            rgb_tile = tile_data[:3].transpose(1, 2, 0)
                        elif tile_data.shape[0] == 1:
                            # Create RGB from grayscale
                            rgb_tile = np.repeat(
                                tile_data[0][:, :, np.newaxis], 3, axis=2
                            )
                        else:
                            print(
                                f"Unsupported number of bands: {tile_data.shape[0]}"
                            )
                            continue

                        # Normalize to 0-255 range if needed
                        if rgb_tile.max() > 0:
                            rgb_tile = (
                                (rgb_tile - rgb_tile.min())
                                / (rgb_tile.max() - rgb_tile.min())
                                * 255
                            ).astype(np.uint8)

                        # Convert to PIL Image
                        pil_image = Image.fromarray(rgb_tile)

                        # Detect objects
                        detections = self._detect(pil_image, text_prompts)

                        if detections:
                            # Apply Non-Maximum Suppression to reduce overlapping detections
                            detections = self._apply_nms(detections, nms_threshold)

                            if detections:
                                # Segment objects
                                detections = self._segment(
                                    pil_image, detections, polygon_refinement
                                )

                                # Process results
                                for i, prompt in enumerate(text_prompts):
                                    prompt_polygons = []
                                    prompt_mask = np.zeros(
                                        (tile_height, tile_width), dtype=np.uint8
                                    )

                                    for detection in detections:
                                        if (
                                            detection.label.replace(".", "")
                                            .strip()
                                            .lower()
                                            == prompt.lower()
                                        ):
                                            if detection.mask is not None:
                                                # Apply gaussian blur to reduce blockiness
                                                try:
                                                    from scipy.ndimage import (
                                                        gaussian_filter,
                                                    )

                                                    smoothed_mask = gaussian_filter(
                                                        detection.mask.astype(
                                                            float
                                                        ),
                                                        sigma=smoothing_sigma,
                                                    )
                                                    detection.mask = (
                                                        smoothed_mask > 0.5
                                                    ).astype(np.uint8)
                                                except ImportError:
                                                    pass

                                                # Add to combined mask for this prompt
                                                prompt_mask = np.maximum(
                                                    prompt_mask, detection.mask
                                                )

                                            # Store bounding box with geospatial coordinates
                                            if export_boxes:
                                                bbox = detection.box
                                                x_geo_min, y_geo_min = transform * (
                                                    x_start + bbox.xmin,
                                                    y_start + bbox.ymin,
                                                )
                                                x_geo_max, y_geo_max = transform * (
                                                    x_start + bbox.xmax,
                                                    y_start + bbox.ymax,
                                                )

                                                geo_box = {
                                                    "label": detection.label,
                                                    "score": detection.score,
                                                    "prompt": prompt,
                                                    "geometry": box(
                                                        x_geo_min,
                                                        y_geo_max,
                                                        x_geo_max,
                                                        y_geo_min,
                                                    ),
                                                }
                                                all_boxes.append(geo_box)

                                    # Convert masks to individual polygons
                                    if export_polygons and np.any(prompt_mask):
                                        tile_polygons = self._mask_to_polygons(
                                            prompt_mask,
                                            transform,
                                            x_start,
                                            y_start,
                                            min_polygon_area,
                                            simplify_tolerance,
                                        )

                                        # Add metadata to polygons
                                        for poly_data in tile_polygons:
                                            poly_data.update(
                                                {
                                                    "label": prompt,
                                                    "score": max(
                                                        [
                                                            d.score
                                                            for d in detections
                                                            if d.label.replace(
                                                                ".", ""
                                                            )
                                                            .strip()
                                                            .lower()
                                                            == prompt.lower()
                                                        ],
                                                        default=0.0,
                                                    ),
                                                    "tile_x": x,
                                                    "tile_y": y,
                                                }
                                            )
                                            all_polygons.append(poly_data)

                                    # Store mask in the global array
                                    valid_x_start = self.overlap if x > 0 else 0
                                    valid_y_start = self.overlap if y > 0 else 0
                                    valid_x_end = (
                                        tile_width - self.overlap
                                        if x < n_tiles_x - 1
                                        else tile_width
                                    )
                                    valid_y_end = (
                                        tile_height - self.overlap
                                        if y < n_tiles_y - 1
                                        else tile_height
                                    )

                                    dest_x_start = x_start + valid_x_start
                                    dest_y_start = y_start + valid_y_start
                                    dest_x_end = x_start + valid_x_end
                                    dest_y_end = y_start + valid_y_end

                                    mask_slice = prompt_mask[
                                        valid_y_start:valid_y_end,
                                        valid_x_start:valid_x_end,
                                    ]
                                    all_masks[
                                        i,
                                        dest_y_start:dest_y_end,
                                        dest_x_start:dest_x_end,
                                    ] = np.maximum(
                                        all_masks[
                                            i,
                                            dest_y_start:dest_y_end,
                                            dest_x_start:dest_x_end,
                                        ],
                                        mask_slice,
                                    )

                    except Exception as e:
                        print(f"Error processing tile at ({x}, {y}): {str(e)}")
                        continue

                    # Update progress bar
                    pbar.update(1)

        # Create combined mask (union of all individual masks)
        combined_mask = np.any(all_masks, axis=0).astype(np.uint8)

        # Write the output GeoTIFF
        with rasterio.open(output_path, "w", **out_meta) as dst:
            # Write combined mask as first band
            dst.write(combined_mask, 1)

            # Write individual masks for each prompt
            for i, mask in enumerate(all_masks):
                dst.write(mask, i + 2)

            # Add descriptions to bands
            dst.set_band_description(1, "Combined Segmentation")
            for i, prompt in enumerate(text_prompts):
                dst.set_band_description(i + 2, f"Segmentation: {prompt}")

        result_files = {"segmentation": output_path}

        # Export bounding boxes if requested
        if export_boxes and all_boxes:
            boxes_path = output_path.replace(".tif", "_boxes.geojson")
            gdf = gpd.GeoDataFrame(all_boxes, crs=crs)
            gdf.to_file(boxes_path, driver="GeoJSON")
            result_files["boxes"] = boxes_path
            print(f"Exported {len(all_boxes)} bounding boxes to {boxes_path}")

        # Export instance polygons if requested
        if export_polygons and all_polygons:
            polygons_path = output_path.replace(".tif", "_polygons.geojson")
            gdf = gpd.GeoDataFrame(all_polygons, crs=crs)
            gdf.to_file(polygons_path, driver="GeoJSON")
            result_files["polygons"] = polygons_path
            print(
                f"Exported {len(all_polygons)} instance polygons to {polygons_path}"
            )

        print(f"Segmentation saved to {output_path}")
        print(
            f"Found {len(all_polygons)} individual building instances"
            if export_polygons
            else ""
        )

        return result_files

LandcoverCrossEntropyLoss

Bases: Module

Enhanced CrossEntropyLoss with optional ignore_index and class weights.

This extends the standard CrossEntropyLoss with more flexible ignore_index handling, specifically designed for landcover classification tasks.

Parameters:

Name	Type	Description	Default
weight	Optional[Tensor]	Manual rescaling weight given to each class	None
ignore_index	Union[int, bool]	Specifies a target value that is ignored. - False: No values ignored (standard behavior) - int: Specific class index to ignore (e.g., 0 for background)	False
reduction	str	Specifies the reduction to apply ('mean', 'sum', 'none')	'mean'

Source code in geoai/landcover_train.py

class LandcoverCrossEntropyLoss(nn.Module):
    """
    Enhanced CrossEntropyLoss with optional ignore_index and class weights.

    This extends the standard CrossEntropyLoss with more flexible ignore_index
    handling, specifically designed for landcover classification tasks.

    Args:
        weight: Manual rescaling weight given to each class
        ignore_index: Specifies a target value that is ignored.
            - False: No values ignored (standard behavior)
            - int: Specific class index to ignore (e.g., 0 for background)
        reduction: Specifies the reduction to apply ('mean', 'sum', 'none')
    """

    def __init__(
        self,
        weight: Optional[torch.Tensor] = None,
        ignore_index: Union[int, bool] = False,
        reduction: str = "mean",
    ):
        super().__init__()
        self.weight = weight
        # Convert ignore_index: int stays as-is, False becomes -100 (PyTorch default)
        self.ignore_index = ignore_index if isinstance(ignore_index, int) else -100
        self.reduction = reduction

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """
        Compute cross entropy loss.

        Args:
            input: Predictions (N, C, H, W) where C = number of classes
            target: Ground truth (N, H, W) with class indices

        Returns:
            Loss value
        """
        return F.cross_entropy(
            input,
            target,
            weight=self.weight,
            ignore_index=self.ignore_index,
            reduction=self.reduction,
        )

forward(input, target)

Compute cross entropy loss.

Parameters:

Name	Type	Description	Default
input	Tensor	Predictions (N, C, H, W) where C = number of classes	required
target	Tensor	Ground truth (N, H, W) with class indices	required

Returns:

Type	Description
Tensor	Loss value

Source code in geoai/landcover_train.py

def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """
    Compute cross entropy loss.

    Args:
        input: Predictions (N, C, H, W) where C = number of classes
        target: Ground truth (N, H, W) with class indices

    Returns:
        Loss value
    """
    return F.cross_entropy(
        input,
        target,
        weight=self.weight,
        ignore_index=self.ignore_index,
        reduction=self.reduction,
    )

LeafMap

Bases: Map

A subclass of leafmap.Map for GeoAI applications.

Source code in geoai/geoai.py

class LeafMap(leafmap.Map):
    """A subclass of leafmap.Map for GeoAI applications."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        """Initialize the Map class."""
        super().__init__(*args, **kwargs)

    def add_dinov3_gui(
        self,
        raster: str,
        processor: "DINOv3GeoProcessor",
        features: torch.Tensor,
        **kwargs: Any,
    ) -> None:
        """Add a DINOv3 GUI to the map."""
        return DINOv3GUI(raster, processor, features, host_map=self, **kwargs)

__init__(*args, **kwargs)

Initialize the Map class.

Source code in geoai/geoai.py

def __init__(self, *args: Any, **kwargs: Any) -> None:
    """Initialize the Map class."""
    super().__init__(*args, **kwargs)

add_dinov3_gui(raster, processor, features, **kwargs)

Add a DINOv3 GUI to the map.

Source code in geoai/geoai.py

def add_dinov3_gui(
    self,
    raster: str,
    processor: "DINOv3GeoProcessor",
    features: torch.Tensor,
    **kwargs: Any,
) -> None:
    """Add a DINOv3 GUI to the map."""
    return DINOv3GUI(raster, processor, features, host_map=self, **kwargs)

Map

Bases: Map

A subclass of maplibregl.Map for GeoAI applications.

Source code in geoai/geoai.py

class Map(maplibregl.Map):
    """A subclass of maplibregl.Map for GeoAI applications."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        """Initialize the Map class."""
        super().__init__(*args, **kwargs)

__init__(*args, **kwargs)

Initialize the Map class.

Source code in geoai/geoai.py

def __init__(self, *args: Any, **kwargs: Any) -> None:
    """Initialize the Map class."""
    super().__init__(*args, **kwargs)

ObjectDetector

Object extraction using Mask R-CNN with TorchGeo.

Source code in geoai/extract.py

class ObjectDetector:
    """
    Object extraction using Mask R-CNN with TorchGeo.
    """

    def __init__(
        self,
        model_path: Optional[str] = None,
        repo_id: Optional[str] = None,
        model: Optional[Any] = None,
        num_classes: int = 2,
        device: Optional[str] = None,
    ) -> None:
        """
        Initialize the object extractor.

        Args:
            model_path: Path to the .pth model file.
            repo_id: Hugging Face repository ID for model download.
            model: Pre-initialized model object (optional).
            num_classes: Number of classes for detection (default: 2).
            device: Device to use for inference ('cuda:0', 'cpu', etc.).
        """
        # Set device
        if device is None:
            self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        else:
            self.device = torch.device(device)

        # Default parameters for object detection - these can be overridden in process_raster
        self.chip_size = (512, 512)  # Size of image chips for processing
        self.overlap = 0.25  # Default overlap between tiles
        self.confidence_threshold = 0.5  # Default confidence threshold
        self.nms_iou_threshold = 0.5  # IoU threshold for non-maximum suppression
        self.min_object_area = 100  # Minimum area in pixels to keep an object
        self.max_object_area = None  # Maximum area in pixels to keep an object
        self.mask_threshold = 0.5  # Threshold for mask binarization
        self.simplify_tolerance = 1.0  # Tolerance for polygon simplification

        # Initialize model
        self.model = self.initialize_model(model, num_classes=num_classes)

        # Download model if needed
        if model_path is None or (not os.path.exists(model_path)):
            model_path = self.download_model_from_hf(model_path, repo_id)

        # Load model weights
        self.load_weights(model_path)

        # Set model to evaluation mode
        self.model.eval()

    def download_model_from_hf(
        self, model_path: Optional[str] = None, repo_id: Optional[str] = None
    ) -> str:
        """
        Download the object detection model from Hugging Face.

        Args:
            model_path: Path to the model file.
            repo_id: Hugging Face repository ID.

        Returns:
            Path to the downloaded model file
        """
        try:

            print("Model path not specified, downloading from Hugging Face...")

            # Define the repository ID and model filename
            if repo_id is None:
                repo_id = "giswqs/geoai"

            if model_path is None:
                model_path = "building_footprints_usa.pth"

            # Download the model
            model_path = hf_hub_download(repo_id=repo_id, filename=model_path)
            print(f"Model downloaded to: {model_path}")

            return model_path

        except Exception as e:
            print(f"Error downloading model from Hugging Face: {e}")
            print("Please specify a local model path or ensure internet connectivity.")
            raise

    def initialize_model(self, model: Optional[Any], num_classes: int = 2) -> Any:
        """Initialize a deep learning model for object detection.

        Args:
            model (torch.nn.Module): A pre-initialized model object.
            num_classes (int): Number of classes for detection.

        Returns:
            torch.nn.Module: A deep learning model for object detection.
        """

        if model is None:  # Initialize Mask R-CNN model with ResNet50 backbone.
            # Standard image mean and std for pre-trained models
            image_mean = [0.485, 0.456, 0.406]
            image_std = [0.229, 0.224, 0.225]

            # Create model with explicit normalization parameters
            model = maskrcnn_resnet50_fpn(
                weights=None,
                progress=False,
                num_classes=num_classes,  # Background + object
                weights_backbone=None,
                # These parameters ensure consistent normalization
                image_mean=image_mean,
                image_std=image_std,
            )

        model.to(self.device)
        return model

    def load_weights(self, model_path: str) -> None:
        """
        Load weights from file with error handling for different formats.

        Args:
            model_path: Path to model weights
        """
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"Model file not found: {model_path}")

        try:
            state_dict = torch.load(model_path, map_location=self.device)

            # Handle different state dict formats
            if isinstance(state_dict, dict):
                if "model" in state_dict:
                    state_dict = state_dict["model"]
                elif "state_dict" in state_dict:
                    state_dict = state_dict["state_dict"]

            # Try to load state dict
            try:
                self.model.load_state_dict(state_dict)
                print("Model loaded successfully")
            except Exception as e:
                print(f"Error loading model: {e}")
                print("Attempting to fix state_dict keys...")

                # Try to fix state_dict keys (remove module prefix if needed)
                new_state_dict = {}
                for k, v in state_dict.items():
                    if k.startswith("module."):
                        new_state_dict[k[7:]] = v
                    else:
                        new_state_dict[k] = v

                self.model.load_state_dict(new_state_dict)
                print("Model loaded successfully after key fixing")

        except Exception as e:
            raise RuntimeError(f"Failed to load model: {e}")

    def mask_to_polygons(self, mask: np.ndarray, **kwargs: Any) -> List[Polygon]:
        """
        Convert binary mask to polygon contours using OpenCV.

        Args:
            mask: Binary mask as numpy array
            **kwargs: Optional parameters:
                simplify_tolerance: Tolerance for polygon simplification
                mask_threshold: Threshold for mask binarization
                min_object_area: Minimum area in pixels to keep an object
                max_object_area: Maximum area in pixels to keep an object

        Returns:
            List of polygons as lists of (x, y) coordinates
        """
        import cv2  # Lazy import to avoid QGIS opencv conflicts

        # Get parameters from kwargs or use instance defaults
        simplify_tolerance = kwargs.get("simplify_tolerance", self.simplify_tolerance)
        mask_threshold = kwargs.get("mask_threshold", self.mask_threshold)
        min_object_area = kwargs.get("min_object_area", self.min_object_area)
        max_object_area = kwargs.get("max_object_area", self.max_object_area)

        # Ensure binary mask
        mask = (mask > mask_threshold).astype(np.uint8)

        # Optional: apply morphological operations to improve mask quality
        kernel = np.ones((3, 3), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

        # Find contours
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        # Convert to list of [x, y] coordinates
        polygons = []
        for contour in contours:
            # Filter out too small contours
            if contour.shape[0] < 3 or cv2.contourArea(contour) < min_object_area:
                continue

            # Filter out too large contours
            if (
                max_object_area is not None
                and cv2.contourArea(contour) > max_object_area
            ):
                continue

            # Simplify contour if it has many points
            if contour.shape[0] > 50:
                epsilon = simplify_tolerance * cv2.arcLength(contour, True)
                contour = cv2.approxPolyDP(contour, epsilon, True)

            # Convert to list of [x, y] coordinates
            polygon = contour.reshape(-1, 2).tolist()
            polygons.append(polygon)

        return polygons

    def filter_overlapping_polygons(
        self, gdf: gpd.GeoDataFrame, **kwargs: Any
    ) -> gpd.GeoDataFrame:
        """
        Filter overlapping polygons using non-maximum suppression.

        Args:
            gdf: GeoDataFrame with polygons
            **kwargs: Optional parameters:
                nms_iou_threshold: IoU threshold for filtering

        Returns:
            Filtered GeoDataFrame
        """
        if len(gdf) <= 1:
            return gdf

        # Get parameters from kwargs or use instance defaults
        iou_threshold = kwargs.get("nms_iou_threshold", self.nms_iou_threshold)

        # Sort by confidence
        gdf = gdf.sort_values("confidence", ascending=False)

        # Fix any invalid geometries
        gdf["geometry"] = gdf["geometry"].apply(
            lambda geom: geom.buffer(0) if not geom.is_valid else geom
        )

        keep_indices = []
        polygons = gdf.geometry.values

        for i in range(len(polygons)):
            if i in keep_indices:
                continue

            keep = True
            for j in keep_indices:
                # Skip invalid geometries
                if not polygons[i].is_valid or not polygons[j].is_valid:
                    continue

                # Calculate IoU
                try:
                    intersection = polygons[i].intersection(polygons[j]).area
                    union = polygons[i].area + polygons[j].area - intersection
                    iou = intersection / union if union > 0 else 0

                    if iou > iou_threshold:
                        keep = False
                        break
                except Exception:
                    # Skip on topology exceptions
                    continue

            if keep:
                keep_indices.append(i)

        return gdf.iloc[keep_indices]

    def filter_edge_objects(
        self, gdf: gpd.GeoDataFrame, raster_path: str, edge_buffer: int = 10
    ) -> gpd.GeoDataFrame:
        """
        Filter out object detections that fall in padding/edge areas of the image.

        Args:
            gdf: GeoDataFrame with object detections
            raster_path: Path to the original raster file
            edge_buffer: Buffer in pixels to consider as edge region

        Returns:
            GeoDataFrame with filtered objects
        """
        import rasterio
        from shapely.geometry import box

        # If no objects detected, return empty GeoDataFrame
        if gdf is None or len(gdf) == 0:
            return gdf

        print(f"Objects before filtering: {len(gdf)}")

        with rasterio.open(raster_path) as src:
            # Get raster bounds
            raster_bounds = src.bounds
            raster_width = src.width
            raster_height = src.height

            # Convert edge buffer from pixels to geographic units
            # We need the smallest dimension of a pixel in geographic units
            pixel_width = (raster_bounds[2] - raster_bounds[0]) / raster_width
            pixel_height = (raster_bounds[3] - raster_bounds[1]) / raster_height
            buffer_size = min(pixel_width, pixel_height) * edge_buffer

            # Create a slightly smaller bounding box to exclude edge regions
            inner_bounds = (
                raster_bounds[0] + buffer_size,  # min x (west)
                raster_bounds[1] + buffer_size,  # min y (south)
                raster_bounds[2] - buffer_size,  # max x (east)
                raster_bounds[3] - buffer_size,  # max y (north)
            )

            # Check that inner bounds are valid
            if inner_bounds[0] >= inner_bounds[2] or inner_bounds[1] >= inner_bounds[3]:
                print("Warning: Edge buffer too large, using original bounds")
                inner_box = box(*raster_bounds)
            else:
                inner_box = box(*inner_bounds)

            # Filter out objects that intersect with the edge of the image
            filtered_gdf = gdf[gdf.intersects(inner_box)]

            # Additional check for objects that have >50% of their area outside the valid region
            valid_objects = []
            for idx, row in filtered_gdf.iterrows():
                if row.geometry.intersection(inner_box).area >= 0.5 * row.geometry.area:
                    valid_objects.append(idx)

            filtered_gdf = filtered_gdf.loc[valid_objects]

            print(f"Objects after filtering: {len(filtered_gdf)}")

            return filtered_gdf

    def masks_to_vector(
        self,
        mask_path: str,
        output_path: Optional[str] = None,
        simplify_tolerance: Optional[float] = None,
        mask_threshold: Optional[float] = None,
        min_object_area: Optional[int] = None,
        max_object_area: Optional[int] = None,
        nms_iou_threshold: Optional[float] = None,
        regularize: bool = True,
        angle_threshold: int = 15,
        rectangularity_threshold: float = 0.7,
    ) -> gpd.GeoDataFrame:
        """
        Convert an object mask GeoTIFF to vector polygons and save as GeoJSON.

        Args:
            mask_path: Path to the object masks GeoTIFF
            output_path: Path to save the output GeoJSON or Parquet file (default: mask_path with .geojson extension)
            simplify_tolerance: Tolerance for polygon simplification (default: self.simplify_tolerance)
            mask_threshold: Threshold for mask binarization (default: self.mask_threshold)
            min_object_area: Minimum area in pixels to keep an object (default: self.min_object_area)
            max_object_area: Minimum area in pixels to keep an object (default: self.max_object_area)
            nms_iou_threshold: IoU threshold for non-maximum suppression (default: self.nms_iou_threshold)
            regularize: Whether to regularize objects to right angles (default: True)
            angle_threshold: Maximum deviation from 90 degrees for regularization (default: 15)
            rectangularity_threshold: Threshold for rectangle simplification (default: 0.7)

        Returns:
            GeoDataFrame with objects
        """
        import cv2  # Lazy import to avoid QGIS opencv conflicts

        # Use class defaults if parameters not provided
        simplify_tolerance = (
            simplify_tolerance
            if simplify_tolerance is not None
            else self.simplify_tolerance
        )
        mask_threshold = (
            mask_threshold if mask_threshold is not None else self.mask_threshold
        )
        min_object_area = (
            min_object_area if min_object_area is not None else self.min_object_area
        )
        max_object_area = (
            max_object_area if max_object_area is not None else self.max_object_area
        )
        nms_iou_threshold = (
            nms_iou_threshold
            if nms_iou_threshold is not None
            else self.nms_iou_threshold
        )

        # Set default output path if not provided
        # if output_path is None:
        #     output_path = os.path.splitext(mask_path)[0] + ".geojson"

        print(f"Converting mask to GeoJSON with parameters:")
        print(f"- Mask threshold: {mask_threshold}")
        print(f"- Min object area: {min_object_area}")
        print(f"- Max object area: {max_object_area}")
        print(f"- Simplify tolerance: {simplify_tolerance}")
        print(f"- NMS IoU threshold: {nms_iou_threshold}")
        print(f"- Regularize objects: {regularize}")
        if regularize:
            print(f"- Angle threshold: {angle_threshold}° from 90°")
            print(f"- Rectangularity threshold: {rectangularity_threshold*100}%")

        # Open the mask raster
        with rasterio.open(mask_path) as src:
            # Read the mask data
            mask_data = src.read(1)
            transform = src.transform
            crs = src.crs

            # Print mask statistics
            print(f"Mask dimensions: {mask_data.shape}")
            print(f"Mask value range: {mask_data.min()} to {mask_data.max()}")

            # Prepare for connected component analysis
            # Binarize the mask based on threshold
            binary_mask = (mask_data > (mask_threshold * 255)).astype(np.uint8)

            # Apply morphological operations for better results (optional)
            kernel = np.ones((3, 3), np.uint8)
            binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, kernel)

            # Find connected components
            num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(
                binary_mask, connectivity=8
            )

            print(
                f"Found {num_labels-1} potential objects"
            )  # Subtract 1 for background

            # Create list to store polygons and confidence values
            all_polygons = []
            all_confidences = []

            # Process each component (skip the first one which is background)
            for i in tqdm(range(1, num_labels)):
                # Extract this object
                area = stats[i, cv2.CC_STAT_AREA]

                # Skip if too small
                if area < min_object_area:
                    continue

                # Skip if too large
                if max_object_area is not None and area > max_object_area:
                    continue

                # Create a mask for this object
                object_mask = (labels == i).astype(np.uint8)

                # Find contours
                contours, _ = cv2.findContours(
                    object_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
                )

                # Process each contour
                for contour in contours:
                    # Skip if too few points
                    if contour.shape[0] < 3:
                        continue

                    # Simplify contour if it has many points
                    if contour.shape[0] > 50 and simplify_tolerance > 0:
                        epsilon = simplify_tolerance * cv2.arcLength(contour, True)
                        contour = cv2.approxPolyDP(contour, epsilon, True)

                    # Convert to list of (x, y) coordinates
                    polygon_points = contour.reshape(-1, 2)

                    # Convert pixel coordinates to geographic coordinates
                    geo_points = []
                    for x, y in polygon_points:
                        gx, gy = transform * (x, y)
                        geo_points.append((gx, gy))

                    # Create Shapely polygon
                    if len(geo_points) >= 3:
                        try:
                            shapely_poly = Polygon(geo_points)
                            if shapely_poly.is_valid and shapely_poly.area > 0:
                                all_polygons.append(shapely_poly)

                                # Calculate "confidence" as normalized size
                                # This is a proxy since we don't have model confidence scores
                                normalized_size = min(1.0, area / 1000)  # Cap at 1.0
                                all_confidences.append(normalized_size)
                        except Exception as e:
                            print(f"Error creating polygon: {e}")

            print(f"Created {len(all_polygons)} valid polygons")

            # Create GeoDataFrame
            if not all_polygons:
                print("No valid polygons found")
                return None

            gdf = gpd.GeoDataFrame(
                {
                    "geometry": all_polygons,
                    "confidence": all_confidences,
                    "class": 1,  # Object class
                },
                crs=crs,
            )

            # Apply non-maximum suppression to remove overlapping polygons
            gdf = self.filter_overlapping_polygons(
                gdf, nms_iou_threshold=nms_iou_threshold
            )

            print(f"Object count after NMS filtering: {len(gdf)}")

            # Apply regularization if requested
            if regularize and len(gdf) > 0:
                # Convert pixel area to geographic units for min_area parameter
                # Estimate pixel size in geographic units
                with rasterio.open(mask_path) as src:
                    pixel_size_x = src.transform[
                        0
                    ]  # width of a pixel in geographic units
                    pixel_size_y = abs(
                        src.transform[4]
                    )  # height of a pixel in geographic units
                    avg_pixel_area = pixel_size_x * pixel_size_y

                # Use 10 pixels as minimum area in geographic units
                min_geo_area = 10 * avg_pixel_area

                # Regularize objects
                gdf = self.regularize_objects(
                    gdf,
                    min_area=min_geo_area,
                    angle_threshold=angle_threshold,
                    rectangularity_threshold=rectangularity_threshold,
                )

            # Save to file
            if output_path:
                if output_path.endswith(".parquet"):
                    gdf.to_parquet(output_path)
                else:
                    gdf.to_file(output_path)
                print(f"Saved {len(gdf)} objects to {output_path}")

            return gdf

    @torch.no_grad()
    def process_raster(
        self,
        raster_path: str,
        output_path: Optional[str] = None,
        batch_size: int = 4,
        filter_edges: bool = True,
        edge_buffer: int = 20,
        band_indexes: Optional[List[int]] = None,
        **kwargs: Any,
    ) -> "gpd.GeoDataFrame":
        """
        Process a raster file to extract objects with customizable parameters.

        Args:
            raster_path: Path to input raster file
            output_path: Path to output GeoJSON or Parquet file (optional)
            batch_size: Batch size for processing
            filter_edges: Whether to filter out objects at the edges of the image
            edge_buffer: Size of edge buffer in pixels to filter out objects (if filter_edges=True)
            band_indexes: List of band indexes to use (if None, use all bands)
            **kwargs: Additional parameters:
                confidence_threshold: Minimum confidence score to keep a detection (0.0-1.0)
                overlap: Overlap between adjacent tiles (0.0-1.0)
                chip_size: Size of image chips for processing (height, width)
                nms_iou_threshold: IoU threshold for non-maximum suppression (0.0-1.0)
                mask_threshold: Threshold for mask binarization (0.0-1.0)
                min_object_area: Minimum area in pixels to keep an object
                simplify_tolerance: Tolerance for polygon simplification

        Returns:
            GeoDataFrame with objects
        """
        # Get parameters from kwargs or use instance defaults
        confidence_threshold = kwargs.get(
            "confidence_threshold", self.confidence_threshold
        )
        overlap = kwargs.get("overlap", self.overlap)
        chip_size = kwargs.get("chip_size", self.chip_size)
        nms_iou_threshold = kwargs.get("nms_iou_threshold", self.nms_iou_threshold)
        mask_threshold = kwargs.get("mask_threshold", self.mask_threshold)
        min_object_area = kwargs.get("min_object_area", self.min_object_area)
        max_object_area = kwargs.get("max_object_area", self.max_object_area)
        simplify_tolerance = kwargs.get("simplify_tolerance", self.simplify_tolerance)

        # Print parameters being used
        print(f"Processing with parameters:")
        print(f"- Confidence threshold: {confidence_threshold}")
        print(f"- Tile overlap: {overlap}")
        print(f"- Chip size: {chip_size}")
        print(f"- NMS IoU threshold: {nms_iou_threshold}")
        print(f"- Mask threshold: {mask_threshold}")
        print(f"- Min object area: {min_object_area}")
        print(f"- Max object area: {max_object_area}")
        print(f"- Simplify tolerance: {simplify_tolerance}")
        print(f"- Filter edge objects: {filter_edges}")
        if filter_edges:
            print(f"- Edge buffer size: {edge_buffer} pixels")

        # Create dataset
        dataset = CustomDataset(
            raster_path=raster_path,
            chip_size=chip_size,
            overlap=overlap,
            band_indexes=band_indexes,
        )
        self.raster_stats = dataset.raster_stats

        # Custom collate function to handle Shapely objects
        def custom_collate(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
            """
            Custom collate function that handles Shapely geometries
            by keeping them as Python objects rather than trying to collate them.
            """
            elem = batch[0]
            if isinstance(elem, dict):
                result = {}
                for key in elem:
                    if key == "bbox":
                        # Don't collate shapely objects, keep as list
                        result[key] = [d[key] for d in batch]
                    else:
                        # For tensors and other collatable types
                        try:
                            result[key] = (
                                torch.utils.data._utils.collate.default_collate(
                                    [d[key] for d in batch]
                                )
                            )
                        except TypeError:
                            # Fall back to list for non-collatable types
                            result[key] = [d[key] for d in batch]
                return result
            else:
                # Default collate for non-dict types
                return torch.utils.data._utils.collate.default_collate(batch)

        # Create dataloader with simple indexing and custom collate
        dataloader = torch.utils.data.DataLoader(
            dataset,
            batch_size=batch_size,
            shuffle=False,
            num_workers=0,
            collate_fn=custom_collate,
        )

        # Process batches
        all_polygons = []
        all_scores = []

        print(f"Processing raster with {len(dataloader)} batches")
        for batch in tqdm(dataloader):
            # Move images to device
            images = batch["image"].to(self.device)
            coords = batch["coords"]  # (i, j) coordinates in pixels
            bboxes = batch[
                "bbox"
            ]  # Geographic bounding boxes - now a list, not a tensor

            # Run inference
            predictions = self.model(images)

            # Process predictions
            for idx, prediction in enumerate(predictions):
                masks = prediction["masks"].cpu().numpy()
                scores = prediction["scores"].cpu().numpy()
                labels = prediction["labels"].cpu().numpy()

                # Skip if no predictions
                if len(scores) == 0:
                    continue

                # Filter by confidence threshold
                valid_indices = scores >= confidence_threshold
                masks = masks[valid_indices]
                scores = scores[valid_indices]
                labels = labels[valid_indices]

                # Skip if no valid predictions
                if len(scores) == 0:
                    continue

                # Get window coordinates
                # The coords might be in different formats depending on batch handling
                if isinstance(coords, list):
                    # If coords is a list of tuples
                    coord_item = coords[idx]
                    if isinstance(coord_item, tuple) and len(coord_item) == 2:
                        i, j = coord_item
                    elif isinstance(coord_item, torch.Tensor):
                        i, j = coord_item.cpu().numpy().tolist()
                    else:
                        print(f"Unexpected coords format: {type(coord_item)}")
                        continue
                elif isinstance(coords, torch.Tensor):
                    # If coords is a tensor of shape [batch_size, 2]
                    i, j = coords[idx].cpu().numpy().tolist()
                else:
                    print(f"Unexpected coords type: {type(coords)}")
                    continue

                # Get window size
                if isinstance(batch["window_size"], list):
                    window_item = batch["window_size"][idx]
                    if isinstance(window_item, tuple) and len(window_item) == 2:
                        window_width, window_height = window_item
                    elif isinstance(window_item, torch.Tensor):
                        window_width, window_height = window_item.cpu().numpy().tolist()
                    else:
                        print(f"Unexpected window_size format: {type(window_item)}")
                        continue
                elif isinstance(batch["window_size"], torch.Tensor):
                    window_width, window_height = (
                        batch["window_size"][idx].cpu().numpy().tolist()
                    )
                else:
                    print(f"Unexpected window_size type: {type(batch['window_size'])}")
                    continue

                # Process masks to polygons
                for mask_idx, mask in enumerate(masks):
                    # Get binary mask
                    binary_mask = mask[0]  # Get binary mask

                    # Convert mask to polygon with custom parameters
                    contours = self.mask_to_polygons(
                        binary_mask,
                        simplify_tolerance=simplify_tolerance,
                        mask_threshold=mask_threshold,
                        min_object_area=min_object_area,
                        max_object_area=max_object_area,
                    )

                    # Skip if no valid polygons
                    if not contours:
                        continue

                    # Transform polygons to geographic coordinates
                    with rasterio.open(raster_path) as src:
                        transform = src.transform

                        for contour in contours:
                            # Convert polygon to global coordinates
                            global_polygon = []
                            for x, y in contour:
                                # Adjust coordinates based on window position
                                gx, gy = transform * (i + x, j + y)
                                global_polygon.append((gx, gy))

                            # Create Shapely polygon
                            if len(global_polygon) >= 3:
                                try:
                                    shapely_poly = Polygon(global_polygon)
                                    if shapely_poly.is_valid and shapely_poly.area > 0:
                                        all_polygons.append(shapely_poly)
                                        all_scores.append(float(scores[mask_idx]))
                                except Exception as e:
                                    print(f"Error creating polygon: {e}")

        # Create GeoDataFrame
        if not all_polygons:
            print("No valid polygons found")
            return None

        gdf = gpd.GeoDataFrame(
            {
                "geometry": all_polygons,
                "confidence": all_scores,
                "class": 1,  # Object class
            },
            crs=dataset.crs,
        )

        # Remove overlapping polygons with custom threshold
        gdf = self.filter_overlapping_polygons(gdf, nms_iou_threshold=nms_iou_threshold)

        # Filter edge objects if requested
        if filter_edges:
            gdf = self.filter_edge_objects(gdf, raster_path, edge_buffer=edge_buffer)

        # Save to file if requested
        if output_path:
            if output_path.endswith(".parquet"):
                gdf.to_parquet(output_path)
            else:
                gdf.to_file(output_path, driver="GeoJSON")
            print(f"Saved {len(gdf)} objects to {output_path}")

        return gdf

    def save_masks_as_geotiff(
        self,
        raster_path: str,
        output_path: Optional[str] = None,
        batch_size: int = 4,
        verbose: bool = False,
        **kwargs: Any,
    ) -> str:
        """
        Process a raster file to extract object masks and save as GeoTIFF.

        Args:
            raster_path: Path to input raster file
            output_path: Path to output GeoTIFF file (optional, default: input_masks.tif)
            batch_size: Batch size for processing
            verbose: Whether to print detailed processing information
            **kwargs: Additional parameters:
                confidence_threshold: Minimum confidence score to keep a detection (0.0-1.0)
                chip_size: Size of image chips for processing (height, width)
                mask_threshold: Threshold for mask binarization (0.0-1.0)

        Returns:
            Path to the saved GeoTIFF file
        """

        # Get parameters from kwargs or use instance defaults
        confidence_threshold = kwargs.get(
            "confidence_threshold", self.confidence_threshold
        )
        chip_size = kwargs.get("chip_size", self.chip_size)
        mask_threshold = kwargs.get("mask_threshold", self.mask_threshold)
        overlap = kwargs.get("overlap", self.overlap)

        # Set default output path if not provided
        if output_path is None:
            output_path = os.path.splitext(raster_path)[0] + "_masks.tif"

        # Print parameters being used
        print(f"Processing masks with parameters:")
        print(f"- Confidence threshold: {confidence_threshold}")
        print(f"- Chip size: {chip_size}")
        print(f"- Mask threshold: {mask_threshold}")

        # Create dataset
        dataset = CustomDataset(
            raster_path=raster_path,
            chip_size=chip_size,
            overlap=overlap,
            verbose=verbose,
        )

        # Store a flag to avoid repetitive messages
        self.raster_stats = dataset.raster_stats
        seen_warnings = {
            "bands": False,
            "resize": {},  # Dictionary to track resize warnings by shape
        }

        # Open original raster to get metadata
        with rasterio.open(raster_path) as src:
            # Create output binary mask raster with same dimensions as input
            output_profile = src.profile.copy()
            output_profile.update(
                dtype=rasterio.uint8,
                count=1,  # Single band for object mask
                compress="lzw",
                nodata=0,
            )

            # Create output mask raster
            with rasterio.open(output_path, "w", **output_profile) as dst:
                # Initialize mask with zeros
                mask_array = np.zeros((src.height, src.width), dtype=np.uint8)

                # Custom collate function to handle Shapely objects
                def custom_collate(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
                    """Custom collate function for DataLoader"""
                    elem = batch[0]
                    if isinstance(elem, dict):
                        result = {}
                        for key in elem:
                            if key == "bbox":
                                # Don't collate shapely objects, keep as list
                                result[key] = [d[key] for d in batch]
                            else:
                                # For tensors and other collatable types
                                try:
                                    result[key] = (
                                        torch.utils.data._utils.collate.default_collate(
                                            [d[key] for d in batch]
                                        )
                                    )
                                except TypeError:
                                    # Fall back to list for non-collatable types
                                    result[key] = [d[key] for d in batch]
                        return result
                    else:
                        # Default collate for non-dict types
                        return torch.utils.data._utils.collate.default_collate(batch)

                # Create dataloader
                dataloader = torch.utils.data.DataLoader(
                    dataset,
                    batch_size=batch_size,
                    shuffle=False,
                    num_workers=0,
                    collate_fn=custom_collate,
                )

                # Process batches
                print(f"Processing raster with {len(dataloader)} batches")
                for batch in tqdm(dataloader):
                    # Move images to device
                    images = batch["image"].to(self.device)
                    coords = batch["coords"]  # (i, j) coordinates in pixels

                    # Run inference
                    with torch.no_grad():
                        predictions = self.model(images)

                    # Process predictions
                    for idx, prediction in enumerate(predictions):
                        masks = prediction["masks"].cpu().numpy()
                        scores = prediction["scores"].cpu().numpy()

                        # Skip if no predictions
                        if len(scores) == 0:
                            continue

                        # Filter by confidence threshold
                        valid_indices = scores >= confidence_threshold
                        masks = masks[valid_indices]
                        scores = scores[valid_indices]

                        # Skip if no valid predictions
                        if len(scores) == 0:
                            continue

                        # Get window coordinates
                        if isinstance(coords, list):
                            coord_item = coords[idx]
                            if isinstance(coord_item, tuple) and len(coord_item) == 2:
                                i, j = coord_item
                            elif isinstance(coord_item, torch.Tensor):
                                i, j = coord_item.cpu().numpy().tolist()
                            else:
                                print(f"Unexpected coords format: {type(coord_item)}")
                                continue
                        elif isinstance(coords, torch.Tensor):
                            i, j = coords[idx].cpu().numpy().tolist()
                        else:
                            print(f"Unexpected coords type: {type(coords)}")
                            continue

                        # Get window size
                        if isinstance(batch["window_size"], list):
                            window_item = batch["window_size"][idx]
                            if isinstance(window_item, tuple) and len(window_item) == 2:
                                window_width, window_height = window_item
                            elif isinstance(window_item, torch.Tensor):
                                window_width, window_height = (
                                    window_item.cpu().numpy().tolist()
                                )
                            else:
                                print(
                                    f"Unexpected window_size format: {type(window_item)}"
                                )
                                continue
                        elif isinstance(batch["window_size"], torch.Tensor):
                            window_width, window_height = (
                                batch["window_size"][idx].cpu().numpy().tolist()
                            )
                        else:
                            print(
                                f"Unexpected window_size type: {type(batch['window_size'])}"
                            )
                            continue

                        # Combine all masks for this window
                        combined_mask = np.zeros(
                            (window_height, window_width), dtype=np.uint8
                        )

                        for mask in masks:
                            # Get the binary mask
                            binary_mask = (mask[0] > mask_threshold).astype(
                                np.uint8
                            ) * 255

                            # Handle size mismatch - resize binary_mask if needed
                            mask_h, mask_w = binary_mask.shape
                            if mask_h != window_height or mask_w != window_width:
                                resize_key = f"{(mask_h, mask_w)}->{(window_height, window_width)}"
                                if resize_key not in seen_warnings["resize"]:
                                    if verbose:
                                        print(
                                            f"Resizing mask from {binary_mask.shape} to {(window_height, window_width)}"
                                        )
                                    else:
                                        if not seen_warnings[
                                            "resize"
                                        ]:  # If this is the first resize warning
                                            print(
                                                f"Resizing masks at image edges (set verbose=True for details)"
                                            )
                                    seen_warnings["resize"][resize_key] = True

                                # Crop or pad the binary mask to match window size
                                resized_mask = np.zeros(
                                    (window_height, window_width), dtype=np.uint8
                                )
                                copy_h = min(mask_h, window_height)
                                copy_w = min(mask_w, window_width)
                                resized_mask[:copy_h, :copy_w] = binary_mask[
                                    :copy_h, :copy_w
                                ]
                                binary_mask = resized_mask

                            # Update combined mask (taking maximum where masks overlap)
                            combined_mask = np.maximum(combined_mask, binary_mask)

                        # Write combined mask to output array
                        # Handle edge cases where window might be smaller than chip size
                        h, w = combined_mask.shape
                        valid_h = min(h, src.height - j)
                        valid_w = min(w, src.width - i)

                        if valid_h > 0 and valid_w > 0:
                            mask_array[j : j + valid_h, i : i + valid_w] = np.maximum(
                                mask_array[j : j + valid_h, i : i + valid_w],
                                combined_mask[:valid_h, :valid_w],
                            )

                # Write the final mask to the output file
                dst.write(mask_array, 1)

        print(f"Object masks saved to {output_path}")
        return output_path

    def regularize_objects(
        self,
        gdf: gpd.GeoDataFrame,
        min_area: int = 10,
        angle_threshold: int = 15,
        orthogonality_threshold: float = 0.3,
        rectangularity_threshold: float = 0.7,
    ) -> gpd.GeoDataFrame:
        """
        Regularize objects to enforce right angles and rectangular shapes.

        Args:
            gdf: GeoDataFrame with objects
            min_area: Minimum area in square units to keep an object
            angle_threshold: Maximum deviation from 90 degrees to consider an angle as orthogonal (degrees)
            orthogonality_threshold: Percentage of angles that must be orthogonal for an object to be regularized
            rectangularity_threshold: Minimum area ratio to Object's oriented bounding box for rectangular simplification

        Returns:
            GeoDataFrame with regularized objects
        """
        import math

        import cv2
        import geopandas as gpd
        import numpy as np
        from shapely.affinity import rotate, translate
        from shapely.geometry import MultiPolygon, Polygon, box
        from tqdm import tqdm

        def get_angle(
            p1: Tuple[float, float], p2: Tuple[float, float], p3: Tuple[float, float]
        ) -> float:
            """Calculate angle between three points in degrees (0-180)"""
            a = np.array(p1)
            b = np.array(p2)
            c = np.array(p3)

            ba = a - b
            bc = c - b

            cosine_angle = np.dot(ba, bc) / (np.linalg.norm(ba) * np.linalg.norm(bc))
            # Handle numerical errors that could push cosine outside [-1, 1]
            cosine_angle = np.clip(cosine_angle, -1.0, 1.0)
            angle = np.degrees(np.arccos(cosine_angle))

            return angle

        def is_orthogonal(angle: float, threshold: int = angle_threshold) -> bool:
            """Check if angle is close to 90 degrees"""
            return abs(angle - 90) <= threshold

        def calculate_dominant_direction(polygon: Polygon) -> float:
            """Find the dominant direction of a polygon using PCA"""
            # Extract coordinates
            coords = np.array(polygon.exterior.coords)

            # Mean center the coordinates
            mean = np.mean(coords, axis=0)
            centered_coords = coords - mean

            # Calculate covariance matrix and its eigenvalues/eigenvectors
            cov_matrix = np.cov(centered_coords.T)
            eigenvalues, eigenvectors = np.linalg.eig(cov_matrix)

            # Get the index of the largest eigenvalue
            largest_idx = np.argmax(eigenvalues)

            # Get the corresponding eigenvector (principal axis)
            principal_axis = eigenvectors[:, largest_idx]

            # Calculate the angle in degrees
            angle_rad = np.arctan2(principal_axis[1], principal_axis[0])
            angle_deg = np.degrees(angle_rad)

            # Normalize to range 0-180
            if angle_deg < 0:
                angle_deg += 180

            return angle_deg

        def create_oriented_envelope(polygon: Polygon, angle_deg: float) -> Polygon:
            """Create an oriented minimum area rectangle for the polygon"""
            # Create a rotated rectangle using OpenCV method (more robust than Shapely methods)
            coords = np.array(polygon.exterior.coords)[:-1].astype(
                np.float32
            )  # Skip the last point (same as first)

            # Use OpenCV's minAreaRect
            rect = cv2.minAreaRect(coords)
            box_points = cv2.boxPoints(rect)

            # Convert to shapely polygon
            oriented_box = Polygon(box_points)

            return oriented_box

        def get_rectangularity(polygon: Polygon, oriented_box: Polygon) -> float:
            """Calculate the rectangularity (area ratio to its oriented bounding box)"""
            if oriented_box.area == 0:
                return 0
            return polygon.area / oriented_box.area

        def check_orthogonality(polygon: Polygon) -> float:
            """Check what percentage of angles in the polygon are orthogonal"""
            coords = list(polygon.exterior.coords)
            if len(coords) <= 4:  # Triangle or point
                return 0

            # Remove last point (same as first)
            coords = coords[:-1]

            orthogonal_count = 0
            total_angles = len(coords)

            for i in range(total_angles):
                p1 = coords[i]
                p2 = coords[(i + 1) % total_angles]
                p3 = coords[(i + 2) % total_angles]

                angle = get_angle(p1, p2, p3)
                if is_orthogonal(angle):
                    orthogonal_count += 1

            return orthogonal_count / total_angles

        def simplify_to_rectangle(polygon: Polygon) -> Polygon:
            """Simplify a polygon to a rectangle using its oriented bounding box"""
            # Get dominant direction
            angle = calculate_dominant_direction(polygon)

            # Create oriented envelope
            rect = create_oriented_envelope(polygon, angle)

            return rect

        if gdf is None or len(gdf) == 0:
            print("No Objects to regularize")
            return gdf

        print(f"Regularizing {len(gdf)} objects...")
        print(f"- Angle threshold: {angle_threshold}° from 90°")
        print(f"- Min orthogonality: {orthogonality_threshold*100}% of angles")
        print(
            f"- Min rectangularity: {rectangularity_threshold*100}% of bounding box area"
        )

        # Create a copy to avoid modifying the original
        result_gdf = gdf.copy()

        # Track statistics
        total_objects = len(gdf)
        regularized_count = 0
        rectangularized_count = 0

        # Process each Object
        for idx, row in tqdm(gdf.iterrows(), total=len(gdf)):
            geom = row.geometry

            # Skip invalid or empty geometries
            if geom is None or geom.is_empty:
                continue

            # Handle MultiPolygons by processing the largest part
            if isinstance(geom, MultiPolygon):
                areas = [p.area for p in geom.geoms]
                if not areas:
                    continue
                geom = list(geom.geoms)[np.argmax(areas)]

            # Filter out tiny Objects
            if geom.area < min_area:
                continue

            # Check orthogonality
            orthogonality = check_orthogonality(geom)

            # Create oriented envelope
            oriented_box = create_oriented_envelope(
                geom, calculate_dominant_direction(geom)
            )

            # Check rectangularity
            rectangularity = get_rectangularity(geom, oriented_box)

            # Decide how to regularize
            if rectangularity >= rectangularity_threshold:
                # Object is already quite rectangular, simplify to a rectangle
                result_gdf.at[idx, "geometry"] = oriented_box
                result_gdf.at[idx, "regularized"] = "rectangle"
                rectangularized_count += 1
            elif orthogonality >= orthogonality_threshold:
                # Object has many orthogonal angles but isn't rectangular
                # Could implement more sophisticated regularization here
                # For now, we'll still use the oriented rectangle
                result_gdf.at[idx, "geometry"] = oriented_box
                result_gdf.at[idx, "regularized"] = "orthogonal"
                regularized_count += 1
            else:
                # Object doesn't have clear orthogonal structure
                # Keep original but flag as unmodified
                result_gdf.at[idx, "regularized"] = "original"

        # Report statistics
        print(f"Regularization completed:")
        print(f"- Total objects: {total_objects}")
        print(
            f"- Rectangular objects: {rectangularized_count} ({rectangularized_count/total_objects*100:.1f}%)"
        )
        print(
            f"- Other regularized objects: {regularized_count} ({regularized_count/total_objects*100:.1f}%)"
        )
        print(
            f"- Unmodified objects: {total_objects-rectangularized_count-regularized_count} ({(total_objects-rectangularized_count-regularized_count)/total_objects*100:.1f}%)"
        )

        return result_gdf

    def visualize_results(
        self,
        raster_path: str,
        gdf: Optional[gpd.GeoDataFrame] = None,
        output_path: Optional[str] = None,
        figsize: Tuple[int, int] = (12, 12),
    ) -> bool:
        """
        Visualize object detection results with proper coordinate transformation.

        This function displays objects on top of the raster image,
        ensuring proper alignment between the GeoDataFrame polygons and the image.

        Args:
            raster_path: Path to input raster
            gdf: GeoDataFrame with object polygons (optional)
            output_path: Path to save visualization (optional)
            figsize: Figure size (width, height) in inches

        Returns:
            bool: True if visualization was successful
        """
        # Check if raster file exists
        if not os.path.exists(raster_path):
            print(f"Error: Raster file '{raster_path}' not found.")
            return False

        # Process raster if GeoDataFrame not provided
        if gdf is None:
            gdf = self.process_raster(raster_path)

        if gdf is None or len(gdf) == 0:
            print("No objects to visualize")
            return False

        # Check if confidence column exists in the GeoDataFrame
        has_confidence = False
        if hasattr(gdf, "columns") and "confidence" in gdf.columns:
            # Try to access a confidence value to confirm it works
            try:
                if len(gdf) > 0:
                    # Try getitem access
                    conf_val = gdf["confidence"].iloc[0]
                    has_confidence = True
                    print(
                        f"Using confidence values (range: {gdf['confidence'].min():.2f} - {gdf['confidence'].max():.2f})"
                    )
            except Exception as e:
                print(f"Confidence column exists but couldn't access values: {e}")
                has_confidence = False
        else:
            print("No confidence column found in GeoDataFrame")
            has_confidence = False

        # Read raster for visualization
        with rasterio.open(raster_path) as src:
            # Read the entire image or a subset if it's very large
            if src.height > 2000 or src.width > 2000:
                # Calculate scale factor to reduce size
                scale = min(2000 / src.height, 2000 / src.width)
                out_shape = (
                    int(src.count),
                    int(src.height * scale),
                    int(src.width * scale),
                )

                # Read and resample
                image = src.read(
                    out_shape=out_shape, resampling=rasterio.enums.Resampling.bilinear
                )

                # Create a scaled transform for the resampled image
                # Calculate scaling factors
                x_scale = src.width / out_shape[2]
                y_scale = src.height / out_shape[1]

                # Get the original transform
                orig_transform = src.transform

                # Create a scaled transform
                scaled_transform = rasterio.transform.Affine(
                    orig_transform.a * x_scale,
                    orig_transform.b,
                    orig_transform.c,
                    orig_transform.d,
                    orig_transform.e * y_scale,
                    orig_transform.f,
                )
            else:
                image = src.read()
                scaled_transform = src.transform

            # Convert to RGB for display
            if image.shape[0] > 3:
                image = image[:3]
            elif image.shape[0] == 1:
                image = np.repeat(image, 3, axis=0)

            # Normalize image for display
            image = image.transpose(1, 2, 0)  # CHW to HWC
            image = image.astype(np.float32)

            if image.max() > 10:  # Likely 0-255 range
                image = image / 255.0

            image = np.clip(image, 0, 1)

            # Get image bounds
            bounds = src.bounds
            crs = src.crs

        # Create figure with appropriate aspect ratio
        aspect_ratio = image.shape[1] / image.shape[0]  # width / height
        plt.figure(figsize=(figsize[0], figsize[0] / aspect_ratio))
        ax = plt.gca()

        # Display image
        ax.imshow(image)

        # Make sure the GeoDataFrame has the same CRS as the raster
        if gdf.crs != crs:
            print(f"Reprojecting GeoDataFrame from {gdf.crs} to {crs}")
            gdf = gdf.to_crs(crs)

        # Set up colors for confidence visualization
        if has_confidence:
            try:
                import matplotlib.cm as cm
                from matplotlib.colors import Normalize

                # Get min/max confidence values
                min_conf = gdf["confidence"].min()
                max_conf = gdf["confidence"].max()

                # Set up normalization and colormap
                norm = Normalize(vmin=min_conf, vmax=max_conf)
                cmap = cm.viridis

                # Create scalar mappable for colorbar
                sm = cm.ScalarMappable(cmap=cmap, norm=norm)
                sm.set_array([])

                # Add colorbar
                cbar = plt.colorbar(
                    sm, ax=ax, orientation="vertical", shrink=0.7, pad=0.01
                )
                cbar.set_label("Confidence Score")
            except Exception as e:
                print(f"Error setting up confidence visualization: {e}")
                has_confidence = False

        # Function to convert coordinates
        def geo_to_pixel(
            geometry: Any, transform: Any
        ) -> Optional[Tuple[List[float], List[float]]]:
            """Convert geometry to pixel coordinates using the provided transform."""
            if geometry.is_empty:
                return None

            if geometry.geom_type == "Polygon":
                # Get exterior coordinates
                exterior_coords = list(geometry.exterior.coords)

                # Convert to pixel coordinates
                pixel_coords = [~transform * (x, y) for x, y in exterior_coords]

                # Split into x and y lists
                pixel_x = [coord[0] for coord in pixel_coords]
                pixel_y = [coord[1] for coord in pixel_coords]

                return pixel_x, pixel_y
            else:
                print(f"Unsupported geometry type: {geometry.geom_type}")
                return None

        # Plot each object
        for idx, row in gdf.iterrows():
            try:
                # Convert polygon to pixel coordinates
                coords = geo_to_pixel(row.geometry, scaled_transform)

                if coords:
                    pixel_x, pixel_y = coords

                    if has_confidence:
                        try:
                            # Get confidence value using different methods
                            # Method 1: Try direct attribute access
                            confidence = None
                            try:
                                confidence = row.confidence
                            except Exception:
                                pass

                            # Method 2: Try dictionary-style access
                            if confidence is None:
                                try:
                                    confidence = row["confidence"]
                                except Exception:
                                    pass

                            # Method 3: Try accessing by index from the GeoDataFrame
                            if confidence is None:
                                try:
                                    confidence = gdf.iloc[idx]["confidence"]
                                except Exception:
                                    pass

                            if confidence is not None:
                                color = cmap(norm(confidence))
                                # Fill polygon with semi-transparent color
                                ax.fill(pixel_x, pixel_y, color=color, alpha=0.5)
                                # Draw border
                                ax.plot(
                                    pixel_x,
                                    pixel_y,
                                    color=color,
                                    linewidth=1,
                                    alpha=0.8,
                                )
                            else:
                                # Fall back to red if confidence value couldn't be accessed
                                ax.plot(pixel_x, pixel_y, color="red", linewidth=1)
                        except Exception as e:
                            print(
                                f"Error using confidence value for polygon {idx}: {e}"
                            )
                            ax.plot(pixel_x, pixel_y, color="red", linewidth=1)
                    else:
                        # No confidence data, just plot outlines in red
                        ax.plot(pixel_x, pixel_y, color="red", linewidth=1)
            except Exception as e:
                print(f"Error plotting polygon {idx}: {e}")

        # Remove axes
        ax.set_xticks([])
        ax.set_yticks([])
        ax.set_title(f"objects (Found: {len(gdf)})")

        # Save if requested
        if output_path:
            plt.tight_layout()
            plt.savefig(output_path, dpi=300, bbox_inches="tight")
            print(f"Visualization saved to {output_path}")

        plt.close()

        # Create a simpler visualization focused just on a subset of objects
        if len(gdf) > 0:
            plt.figure(figsize=figsize)
            ax = plt.gca()

            # Choose a subset of the image to show
            with rasterio.open(raster_path) as src:
                # Get centroid of first object
                sample_geom = gdf.iloc[0].geometry
                centroid = sample_geom.centroid

                # Convert to pixel coordinates
                center_x, center_y = ~src.transform * (centroid.x, centroid.y)

                # Define a window around this object
                window_size = 500  # pixels
                window = rasterio.windows.Window(
                    max(0, int(center_x - window_size / 2)),
                    max(0, int(center_y - window_size / 2)),
                    min(window_size, src.width - int(center_x - window_size / 2)),
                    min(window_size, src.height - int(center_y - window_size / 2)),
                )

                # Read this window
                sample_image = src.read(window=window)

                # Convert to RGB for display
                if sample_image.shape[0] > 3:
                    sample_image = sample_image[:3]
                elif sample_image.shape[0] == 1:
                    sample_image = np.repeat(sample_image, 3, axis=0)

                # Normalize image for display
                sample_image = sample_image.transpose(1, 2, 0)  # CHW to HWC
                sample_image = sample_image.astype(np.float32)

                if sample_image.max() > 10:  # Likely 0-255 range
                    sample_image = sample_image / 255.0

                sample_image = np.clip(sample_image, 0, 1)

                # Display sample image
                ax.imshow(sample_image, extent=[0, window.width, window.height, 0])

                # Get the correct transform for this window
                window_transform = src.window_transform(window)

                # Calculate bounds of the window
                window_bounds = rasterio.windows.bounds(window, src.transform)
                window_box = box(*window_bounds)

                # Filter objects that intersect with this window
                visible_gdf = gdf[gdf.intersects(window_box)]

                # Set up colors for sample view if confidence data exists
                if has_confidence:
                    try:
                        # Reuse the same normalization and colormap from main view
                        sample_sm = cm.ScalarMappable(cmap=cmap, norm=norm)
                        sample_sm.set_array([])

                        # Add colorbar to sample view
                        sample_cbar = plt.colorbar(
                            sample_sm,
                            ax=ax,
                            orientation="vertical",
                            shrink=0.7,
                            pad=0.01,
                        )
                        sample_cbar.set_label("Confidence Score")
                    except Exception as e:
                        print(f"Error setting up sample confidence visualization: {e}")

                # Plot objects in sample view
                for idx, row in visible_gdf.iterrows():
                    try:
                        # Get window-relative pixel coordinates
                        geom = row.geometry

                        # Skip empty geometries
                        if geom.is_empty:
                            continue

                        # Get exterior coordinates
                        exterior_coords = list(geom.exterior.coords)

                        # Convert to pixel coordinates relative to window origin
                        pixel_coords = []
                        for x, y in exterior_coords:
                            px, py = ~src.transform * (x, y)  # Convert to image pixels
                            # Make coordinates relative to window
                            px = px - window.col_off
                            py = py - window.row_off
                            pixel_coords.append((px, py))

                        # Extract x and y coordinates
                        pixel_x = [coord[0] for coord in pixel_coords]
                        pixel_y = [coord[1] for coord in pixel_coords]

                        # Use confidence colors if available
                        if has_confidence:
                            try:
                                # Try different methods to access confidence
                                confidence = None
                                try:
                                    confidence = row.confidence
                                except Exception:
                                    pass

                                if confidence is None:
                                    try:
                                        confidence = row["confidence"]
                                    except Exception:
                                        pass

                                if confidence is None:
                                    try:
                                        confidence = visible_gdf.iloc[idx]["confidence"]
                                    except Exception:
                                        pass

                                if confidence is not None:
                                    color = cmap(norm(confidence))
                                    # Fill polygon with semi-transparent color
                                    ax.fill(pixel_x, pixel_y, color=color, alpha=0.5)
                                    # Draw border
                                    ax.plot(
                                        pixel_x,
                                        pixel_y,
                                        color=color,
                                        linewidth=1.5,
                                        alpha=0.8,
                                    )
                                else:
                                    ax.plot(
                                        pixel_x, pixel_y, color="red", linewidth=1.5
                                    )
                            except Exception as e:
                                print(
                                    f"Error using confidence in sample view for polygon {idx}: {e}"
                                )
                                ax.plot(pixel_x, pixel_y, color="red", linewidth=1.5)
                        else:
                            ax.plot(pixel_x, pixel_y, color="red", linewidth=1.5)
                    except Exception as e:
                        print(f"Error plotting polygon in sample view: {e}")

                # Set title
                ax.set_title(f"Sample Area - objects (Showing: {len(visible_gdf)})")

                # Remove axes
                ax.set_xticks([])
                ax.set_yticks([])

                # Save if requested
                if output_path:
                    sample_output = (
                        os.path.splitext(output_path)[0]
                        + "_sample"
                        + os.path.splitext(output_path)[1]
                    )
                    plt.tight_layout()
                    plt.savefig(sample_output, dpi=300, bbox_inches="tight")
                    print(f"Sample visualization saved to {sample_output}")

    def generate_masks(
        self,
        raster_path: str,
        output_path: Optional[str] = None,
        confidence_threshold: Optional[float] = None,
        mask_threshold: Optional[float] = None,
        min_object_area: int = 10,
        max_object_area: float = float("inf"),
        overlap: float = 0.25,
        batch_size: int = 4,
        band_indexes: Optional[List[int]] = None,
        verbose: bool = False,
        **kwargs: Any,
    ) -> str:
        """
        Save masks with confidence values as a multi-band GeoTIFF.

        Objects with area smaller than min_object_area or larger than max_object_area
        will be filtered out.

        Args:
            raster_path: Path to input raster
            output_path: Path for output GeoTIFF
            confidence_threshold: Minimum confidence score (0.0-1.0)
            mask_threshold: Threshold for mask binarization (0.0-1.0)
            min_object_area: Minimum area (in pixels) for an object to be included
            max_object_area: Maximum area (in pixels) for an object to be included
            overlap: Overlap between tiles (0.0-1.0)
            batch_size: Batch size for processing
            band_indexes: List of band indexes to use (default: all bands)
            verbose: Whether to print detailed processing information

        Returns:
            Path to the saved GeoTIFF
        """
        # Use provided thresholds or fall back to instance defaults
        if confidence_threshold is None:
            confidence_threshold = self.confidence_threshold
        if mask_threshold is None:
            mask_threshold = self.mask_threshold

        chip_size = kwargs.get("chip_size", self.chip_size)

        # Default output path
        if output_path is None:
            output_path = os.path.splitext(raster_path)[0] + "_masks_conf.tif"

        # Process the raster to get individual masks with confidence
        with rasterio.open(raster_path) as src:
            # Create dataset with the specified overlap
            dataset = CustomDataset(
                raster_path=raster_path,
                chip_size=chip_size,
                overlap=overlap,
                band_indexes=band_indexes,
                verbose=verbose,
            )

            # Create output profile
            output_profile = src.profile.copy()
            output_profile.update(
                dtype=rasterio.uint8,
                count=2,  # Two bands: mask and confidence
                compress="lzw",
                nodata=0,
            )

            # Initialize mask and confidence arrays
            mask_array = np.zeros((src.height, src.width), dtype=np.uint8)
            conf_array = np.zeros((src.height, src.width), dtype=np.uint8)

            # Define custom collate function to handle Shapely objects
            def custom_collate(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
                """
                Custom collate function that handles Shapely geometries
                by keeping them as Python objects rather than trying to collate them.
                """
                elem = batch[0]
                if isinstance(elem, dict):
                    result = {}
                    for key in elem:
                        if key == "bbox":
                            # Don't collate shapely objects, keep as list
                            result[key] = [d[key] for d in batch]
                        else:
                            # For tensors and other collatable types
                            try:
                                result[key] = (
                                    torch.utils.data._utils.collate.default_collate(
                                        [d[key] for d in batch]
                                    )
                                )
                            except TypeError:
                                # Fall back to list for non-collatable types
                                result[key] = [d[key] for d in batch]
                    return result
                else:
                    # Default collate for non-dict types
                    return torch.utils.data._utils.collate.default_collate(batch)

            # Create dataloader with custom collate function
            dataloader = torch.utils.data.DataLoader(
                dataset,
                batch_size=batch_size,
                shuffle=False,
                num_workers=0,
                collate_fn=custom_collate,
            )

            # Process batches
            print(f"Processing raster with {len(dataloader)} batches")
            for batch in tqdm(dataloader):
                # Move images to device
                images = batch["image"].to(self.device)
                coords = batch["coords"]  # Tensor of shape [batch_size, 2]

                # Run inference
                with torch.no_grad():
                    predictions = self.model(images)

                # Process predictions
                for idx, prediction in enumerate(predictions):
                    masks = prediction["masks"].cpu().numpy()
                    scores = prediction["scores"].cpu().numpy()

                    # Filter by confidence threshold
                    valid_indices = scores >= confidence_threshold
                    masks = masks[valid_indices]
                    scores = scores[valid_indices]

                    # Skip if no valid predictions
                    if len(masks) == 0:
                        continue

                    # Get window coordinates
                    i, j = coords[idx].cpu().numpy()

                    # Process each mask
                    for mask_idx, mask in enumerate(masks):
                        # Convert to binary mask
                        binary_mask = (mask[0] > mask_threshold).astype(np.uint8) * 255

                        # Check object area - calculate number of pixels in the mask
                        object_area = np.sum(binary_mask > 0)

                        # Skip objects that don't meet area criteria
                        if (
                            object_area < min_object_area
                            or object_area > max_object_area
                        ):
                            if verbose:
                                print(
                                    f"Filtering out object with area {object_area} pixels"
                                )
                            continue

                        conf_value = int(scores[mask_idx] * 255)  # Scale to 0-255

                        # Update the mask and confidence arrays
                        h, w = binary_mask.shape
                        valid_h = min(h, src.height - j)
                        valid_w = min(w, src.width - i)

                        if valid_h > 0 and valid_w > 0:
                            # Use maximum for overlapping regions in the mask
                            mask_array[j : j + valid_h, i : i + valid_w] = np.maximum(
                                mask_array[j : j + valid_h, i : i + valid_w],
                                binary_mask[:valid_h, :valid_w],
                            )

                            # For confidence, only update where mask is positive
                            # and confidence is higher than existing
                            mask_region = binary_mask[:valid_h, :valid_w] > 0
                            if np.any(mask_region):
                                # Only update where mask is positive and new confidence is higher
                                current_conf = conf_array[
                                    j : j + valid_h, i : i + valid_w
                                ]

                                # Where to update confidence (mask positive & higher confidence)
                                update_mask = np.logical_and(
                                    mask_region,
                                    np.logical_or(
                                        current_conf == 0, current_conf < conf_value
                                    ),
                                )

                                if np.any(update_mask):
                                    conf_array[j : j + valid_h, i : i + valid_w][
                                        update_mask
                                    ] = conf_value

            # Write to GeoTIFF
            with rasterio.open(output_path, "w", **output_profile) as dst:
                dst.write(mask_array, 1)
                dst.write(conf_array, 2)

            print(f"Masks with confidence values saved to {output_path}")
            return output_path

    def vectorize_masks(
        self,
        masks_path: str,
        output_path: Optional[str] = None,
        confidence_threshold: float = 0.5,
        min_object_area: int = 100,
        max_object_area: Optional[int] = None,
        n_workers: Optional[int] = None,
        **kwargs: Any,
    ) -> gpd.GeoDataFrame:
        """
        Convert masks with confidence to vector polygons.

        Args:
            masks_path: Path to masks GeoTIFF with confidence band.
            output_path: Path for output GeoJSON.
            confidence_threshold: Minimum confidence score (0.0-1.0). Default: 0.5
            min_object_area: Minimum area in pixels to keep an object. Default: 100
            max_object_area: Maximum area in pixels to keep an object. Default: None
            n_workers: int, default=None
                The number of worker threads to use.
                "None" means single-threaded processing.
                "-1"   means using all available CPU processors.
                Positive integer means using that specific number of threads.
            **kwargs: Additional parameters

        Returns:
            GeoDataFrame with car detections and confidence values
        """
        import cv2  # Lazy import to avoid QGIS opencv conflicts

        def _process_single_component(
            component_mask: np.ndarray,
            conf_data: np.ndarray,
            transform: Any,
            confidence_threshold: float,
            min_object_area: int,
            max_object_area: Optional[int],
        ) -> Optional[Dict[str, Any]]:
            # Get confidence value
            conf_region = conf_data[component_mask > 0]
            if len(conf_region) > 0:
                confidence = np.mean(conf_region) / 255.0
            else:
                confidence = 0.0

            # Skip if confidence is below threshold
            if confidence < confidence_threshold:
                return None

            # Find contours
            contours, _ = cv2.findContours(
                component_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
            )

            results = []

            for contour in contours:
                # Filter by size
                area = cv2.contourArea(contour)
                if area < min_object_area:
                    continue

                if max_object_area is not None and area > max_object_area:
                    continue

                # Get minimum area rectangle
                rect = cv2.minAreaRect(contour)
                box_points = cv2.boxPoints(rect)

                # Convert to geographic coordinates
                geo_points = []
                for x, y in box_points:
                    gx, gy = transform * (x, y)
                    geo_points.append((gx, gy))

                # Create polygon
                poly = Polygon(geo_points)
                results.append((poly, confidence, area))

            return results

        import concurrent.futures
        from functools import partial

        def process_component(
            args: Tuple[int, np.ndarray, np.ndarray, Any, float, int, Optional[int]],
        ) -> Optional[Dict[str, Any]]:
            """
            Helper function to process a single component
            """
            (
                label,
                labeled_mask,
                conf_data,
                transform,
                confidence_threshold,
                min_object_area,
                max_object_area,
            ) = args

            # Create mask for this component
            component_mask = (labeled_mask == label).astype(np.uint8)

            return _process_single_component(
                component_mask,
                conf_data,
                transform,
                confidence_threshold,
                min_object_area,
                max_object_area,
            )

        start_time = time.time()
        print(f"Processing masks from: {masks_path}")

        if n_workers == -1:
            n_workers = os.cpu_count()

        with rasterio.open(masks_path) as src:
            # Read mask and confidence bands
            mask_data = src.read(1)
            conf_data = src.read(2)
            transform = src.transform
            crs = src.crs

            # Convert to binary mask
            binary_mask = mask_data > 0

            # Find connected components
            labeled_mask, num_features = ndimage.label(binary_mask)
            print(f"Found {num_features} connected components")

            # Process each component
            polygons = []
            confidences = []
            pixels = []

            if n_workers is None or n_workers == 1:
                print(
                    "Using single-threaded processing, you can speed up processing by setting n_workers > 1"
                )
                # Add progress bar
                for label in tqdm(
                    range(1, num_features + 1), desc="Processing components"
                ):
                    # Create mask for this component
                    component_mask = (labeled_mask == label).astype(np.uint8)

                    result = _process_single_component(
                        component_mask,
                        conf_data,
                        transform,
                        confidence_threshold,
                        min_object_area,
                        max_object_area,
                    )

                    if result:
                        for poly, confidence, area in result:
                            # Add to lists
                            polygons.append(poly)
                            confidences.append(confidence)
                            pixels.append(area)

            else:
                # Process components in parallel
                print(f"Using {n_workers} workers for parallel processing")

                process_args = [
                    (
                        label,
                        labeled_mask,
                        conf_data,
                        transform,
                        confidence_threshold,
                        min_object_area,
                        max_object_area,
                    )
                    for label in range(1, num_features + 1)
                ]

                with concurrent.futures.ThreadPoolExecutor(
                    max_workers=n_workers
                ) as executor:
                    results = list(
                        tqdm(
                            executor.map(process_component, process_args),
                            total=num_features,
                            desc="Processing components",
                        )
                    )

                    for result in results:
                        if result:
                            for poly, confidence, area in result:
                                # Add to lists
                                polygons.append(poly)
                                confidences.append(confidence)
                                pixels.append(area)

            # Create GeoDataFrame
            if polygons:
                gdf = gpd.GeoDataFrame(
                    {
                        "geometry": polygons,
                        "confidence": confidences,
                        "class": [1] * len(polygons),
                        "pixels": pixels,
                    },
                    crs=crs,
                )

                # Save to file if requested
                if output_path:
                    gdf.to_file(output_path, driver="GeoJSON")
                    print(f"Saved {len(gdf)} objects with confidence to {output_path}")

                end_time = time.time()
                print(f"Total processing time: {end_time - start_time:.2f} seconds")
                return gdf
            else:
                end_time = time.time()
                print(f"Total processing time: {end_time - start_time:.2f} seconds")
                print("No valid polygons found")
                return None

__init__(model_path=None, repo_id=None, model=None, num_classes=2, device=None)

Initialize the object extractor.

Parameters:

Name	Type	Description	Default
model_path	Optional[str]	Path to the .pth model file.	None
repo_id	Optional[str]	Hugging Face repository ID for model download.	None
model	Optional[Any]	Pre-initialized model object (optional).	None
num_classes	int	Number of classes for detection (default: 2).	2
device	Optional[str]	Device to use for inference ('cuda:0', 'cpu', etc.).	None

Source code in geoai/extract.py

def __init__(
    self,
    model_path: Optional[str] = None,
    repo_id: Optional[str] = None,
    model: Optional[Any] = None,
    num_classes: int = 2,
    device: Optional[str] = None,
) -> None:
    """
    Initialize the object extractor.

    Args:
        model_path: Path to the .pth model file.
        repo_id: Hugging Face repository ID for model download.
        model: Pre-initialized model object (optional).
        num_classes: Number of classes for detection (default: 2).
        device: Device to use for inference ('cuda:0', 'cpu', etc.).
    """
    # Set device
    if device is None:
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    else:
        self.device = torch.device(device)

    # Default parameters for object detection - these can be overridden in process_raster
    self.chip_size = (512, 512)  # Size of image chips for processing
    self.overlap = 0.25  # Default overlap between tiles
    self.confidence_threshold = 0.5  # Default confidence threshold
    self.nms_iou_threshold = 0.5  # IoU threshold for non-maximum suppression
    self.min_object_area = 100  # Minimum area in pixels to keep an object
    self.max_object_area = None  # Maximum area in pixels to keep an object
    self.mask_threshold = 0.5  # Threshold for mask binarization
    self.simplify_tolerance = 1.0  # Tolerance for polygon simplification

    # Initialize model
    self.model = self.initialize_model(model, num_classes=num_classes)

    # Download model if needed
    if model_path is None or (not os.path.exists(model_path)):
        model_path = self.download_model_from_hf(model_path, repo_id)

    # Load model weights
    self.load_weights(model_path)

    # Set model to evaluation mode
    self.model.eval()

download_model_from_hf(model_path=None, repo_id=None)

Download the object detection model from Hugging Face.

Parameters:

Name	Type	Description	Default
model_path	Optional[str]	Path to the model file.	None
repo_id	Optional[str]	Hugging Face repository ID.	None

Returns:

Type	Description
str	Path to the downloaded model file

Source code in geoai/extract.py

def download_model_from_hf(
    self, model_path: Optional[str] = None, repo_id: Optional[str] = None
) -> str:
    """
    Download the object detection model from Hugging Face.

    Args:
        model_path: Path to the model file.
        repo_id: Hugging Face repository ID.

    Returns:
        Path to the downloaded model file
    """
    try:

        print("Model path not specified, downloading from Hugging Face...")

        # Define the repository ID and model filename
        if repo_id is None:
            repo_id = "giswqs/geoai"

        if model_path is None:
            model_path = "building_footprints_usa.pth"

        # Download the model
        model_path = hf_hub_download(repo_id=repo_id, filename=model_path)
        print(f"Model downloaded to: {model_path}")

        return model_path

    except Exception as e:
        print(f"Error downloading model from Hugging Face: {e}")
        print("Please specify a local model path or ensure internet connectivity.")
        raise

filter_edge_objects(gdf, raster_path, edge_buffer=10)

Filter out object detections that fall in padding/edge areas of the image.

Parameters:

Name	Type	Description	Default
gdf	GeoDataFrame	GeoDataFrame with object detections	required
raster_path	str	Path to the original raster file	required
edge_buffer	int	Buffer in pixels to consider as edge region	10

Returns:

Type	Description
GeoDataFrame	GeoDataFrame with filtered objects

Source code in geoai/extract.py

def filter_edge_objects(
    self, gdf: gpd.GeoDataFrame, raster_path: str, edge_buffer: int = 10
) -> gpd.GeoDataFrame:
    """
    Filter out object detections that fall in padding/edge areas of the image.

    Args:
        gdf: GeoDataFrame with object detections
        raster_path: Path to the original raster file
        edge_buffer: Buffer in pixels to consider as edge region

    Returns:
        GeoDataFrame with filtered objects
    """
    import rasterio
    from shapely.geometry import box

    # If no objects detected, return empty GeoDataFrame
    if gdf is None or len(gdf) == 0:
        return gdf

    print(f"Objects before filtering: {len(gdf)}")

    with rasterio.open(raster_path) as src:
        # Get raster bounds
        raster_bounds = src.bounds
        raster_width = src.width
        raster_height = src.height

        # Convert edge buffer from pixels to geographic units
        # We need the smallest dimension of a pixel in geographic units
        pixel_width = (raster_bounds[2] - raster_bounds[0]) / raster_width
        pixel_height = (raster_bounds[3] - raster_bounds[1]) / raster_height
        buffer_size = min(pixel_width, pixel_height) * edge_buffer

        # Create a slightly smaller bounding box to exclude edge regions
        inner_bounds = (
            raster_bounds[0] + buffer_size,  # min x (west)
            raster_bounds[1] + buffer_size,  # min y (south)
            raster_bounds[2] - buffer_size,  # max x (east)
            raster_bounds[3] - buffer_size,  # max y (north)
        )

        # Check that inner bounds are valid
        if inner_bounds[0] >= inner_bounds[2] or inner_bounds[1] >= inner_bounds[3]:
            print("Warning: Edge buffer too large, using original bounds")
            inner_box = box(*raster_bounds)
        else:
            inner_box = box(*inner_bounds)

        # Filter out objects that intersect with the edge of the image
        filtered_gdf = gdf[gdf.intersects(inner_box)]

        # Additional check for objects that have >50% of their area outside the valid region
        valid_objects = []
        for idx, row in filtered_gdf.iterrows():
            if row.geometry.intersection(inner_box).area >= 0.5 * row.geometry.area:
                valid_objects.append(idx)

        filtered_gdf = filtered_gdf.loc[valid_objects]

        print(f"Objects after filtering: {len(filtered_gdf)}")

        return filtered_gdf

filter_overlapping_polygons(gdf, **kwargs)

Filter overlapping polygons using non-maximum suppression.

Parameters:

Name	Type	Description	Default
gdf	GeoDataFrame	GeoDataFrame with polygons	required
**kwargs	Any	Optional parameters: nms_iou_threshold: IoU threshold for filtering	{}

Returns:

Type	Description
GeoDataFrame	Filtered GeoDataFrame

Source code in geoai/extract.py

def filter_overlapping_polygons(
    self, gdf: gpd.GeoDataFrame, **kwargs: Any
) -> gpd.GeoDataFrame:
    """
    Filter overlapping polygons using non-maximum suppression.

    Args:
        gdf: GeoDataFrame with polygons
        **kwargs: Optional parameters:
            nms_iou_threshold: IoU threshold for filtering

    Returns:
        Filtered GeoDataFrame
    """
    if len(gdf) <= 1:
        return gdf

    # Get parameters from kwargs or use instance defaults
    iou_threshold = kwargs.get("nms_iou_threshold", self.nms_iou_threshold)

    # Sort by confidence
    gdf = gdf.sort_values("confidence", ascending=False)

    # Fix any invalid geometries
    gdf["geometry"] = gdf["geometry"].apply(
        lambda geom: geom.buffer(0) if not geom.is_valid else geom
    )

    keep_indices = []
    polygons = gdf.geometry.values

    for i in range(len(polygons)):
        if i in keep_indices:
            continue

        keep = True
        for j in keep_indices:
            # Skip invalid geometries
            if not polygons[i].is_valid or not polygons[j].is_valid:
                continue

            # Calculate IoU
            try:
                intersection = polygons[i].intersection(polygons[j]).area
                union = polygons[i].area + polygons[j].area - intersection
                iou = intersection / union if union > 0 else 0

                if iou > iou_threshold:
                    keep = False
                    break
            except Exception:
                # Skip on topology exceptions
                continue

        if keep:
            keep_indices.append(i)

    return gdf.iloc[keep_indices]

generate_masks(raster_path, output_path=None, confidence_threshold=None, mask_threshold=None, min_object_area=10, max_object_area=float('inf'), overlap=0.25, batch_size=4, band_indexes=None, verbose=False, **kwargs)

Save masks with confidence values as a multi-band GeoTIFF.

Objects with area smaller than min_object_area or larger than max_object_area will be filtered out.

Parameters:

Name	Type	Description	Default
raster_path	str	Path to input raster	required
output_path	Optional[str]	Path for output GeoTIFF	None
confidence_threshold	Optional[float]	Minimum confidence score (0.0-1.0)	None
mask_threshold	Optional[float]	Threshold for mask binarization (0.0-1.0)	None
min_object_area	int	Minimum area (in pixels) for an object to be included	10
max_object_area	float	Maximum area (in pixels) for an object to be included	float('inf')
overlap	float	Overlap between tiles (0.0-1.0)	0.25
batch_size	int	Batch size for processing	4
band_indexes	Optional[List[int]]	List of band indexes to use (default: all bands)	None
verbose	bool	Whether to print detailed processing information	False

Returns:

Type	Description
str	Path to the saved GeoTIFF

Source code in geoai/extract.py

def generate_masks(
    self,
    raster_path: str,
    output_path: Optional[str] = None,
    confidence_threshold: Optional[float] = None,
    mask_threshold: Optional[float] = None,
    min_object_area: int = 10,
    max_object_area: float = float("inf"),
    overlap: float = 0.25,
    batch_size: int = 4,
    band_indexes: Optional[List[int]] = None,
    verbose: bool = False,
    **kwargs: Any,
) -> str:
    """
    Save masks with confidence values as a multi-band GeoTIFF.

    Objects with area smaller than min_object_area or larger than max_object_area
    will be filtered out.

    Args:
        raster_path: Path to input raster
        output_path: Path for output GeoTIFF
        confidence_threshold: Minimum confidence score (0.0-1.0)
        mask_threshold: Threshold for mask binarization (0.0-1.0)
        min_object_area: Minimum area (in pixels) for an object to be included
        max_object_area: Maximum area (in pixels) for an object to be included
        overlap: Overlap between tiles (0.0-1.0)
        batch_size: Batch size for processing
        band_indexes: List of band indexes to use (default: all bands)
        verbose: Whether to print detailed processing information

    Returns:
        Path to the saved GeoTIFF
    """
    # Use provided thresholds or fall back to instance defaults
    if confidence_threshold is None:
        confidence_threshold = self.confidence_threshold
    if mask_threshold is None:
        mask_threshold = self.mask_threshold

    chip_size = kwargs.get("chip_size", self.chip_size)

    # Default output path
    if output_path is None:
        output_path = os.path.splitext(raster_path)[0] + "_masks_conf.tif"

    # Process the raster to get individual masks with confidence
    with rasterio.open(raster_path) as src:
        # Create dataset with the specified overlap
        dataset = CustomDataset(
            raster_path=raster_path,
            chip_size=chip_size,
            overlap=overlap,
            band_indexes=band_indexes,
            verbose=verbose,
        )

        # Create output profile
        output_profile = src.profile.copy()
        output_profile.update(
            dtype=rasterio.uint8,
            count=2,  # Two bands: mask and confidence
            compress="lzw",
            nodata=0,
        )

        # Initialize mask and confidence arrays
        mask_array = np.zeros((src.height, src.width), dtype=np.uint8)
        conf_array = np.zeros((src.height, src.width), dtype=np.uint8)

        # Define custom collate function to handle Shapely objects
        def custom_collate(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
            """
            Custom collate function that handles Shapely geometries
            by keeping them as Python objects rather than trying to collate them.
            """
            elem = batch[0]
            if isinstance(elem, dict):
                result = {}
                for key in elem:
                    if key == "bbox":
                        # Don't collate shapely objects, keep as list
                        result[key] = [d[key] for d in batch]
                    else:
                        # For tensors and other collatable types
                        try:
                            result[key] = (
                                torch.utils.data._utils.collate.default_collate(
                                    [d[key] for d in batch]
                                )
                            )
                        except TypeError:
                            # Fall back to list for non-collatable types
                            result[key] = [d[key] for d in batch]
                return result
            else:
                # Default collate for non-dict types
                return torch.utils.data._utils.collate.default_collate(batch)

        # Create dataloader with custom collate function
        dataloader = torch.utils.data.DataLoader(
            dataset,
            batch_size=batch_size,
            shuffle=False,
            num_workers=0,
            collate_fn=custom_collate,
        )

        # Process batches
        print(f"Processing raster with {len(dataloader)} batches")
        for batch in tqdm(dataloader):
            # Move images to device
            images = batch["image"].to(self.device)
            coords = batch["coords"]  # Tensor of shape [batch_size, 2]

            # Run inference
            with torch.no_grad():
                predictions = self.model(images)

            # Process predictions
            for idx, prediction in enumerate(predictions):
                masks = prediction["masks"].cpu().numpy()
                scores = prediction["scores"].cpu().numpy()

                # Filter by confidence threshold
                valid_indices = scores >= confidence_threshold
                masks = masks[valid_indices]
                scores = scores[valid_indices]

                # Skip if no valid predictions
                if len(masks) == 0:
                    continue

                # Get window coordinates
                i, j = coords[idx].cpu().numpy()

                # Process each mask
                for mask_idx, mask in enumerate(masks):
                    # Convert to binary mask
                    binary_mask = (mask[0] > mask_threshold).astype(np.uint8) * 255

                    # Check object area - calculate number of pixels in the mask
                    object_area = np.sum(binary_mask > 0)

                    # Skip objects that don't meet area criteria
                    if (
                        object_area < min_object_area
                        or object_area > max_object_area
                    ):
                        if verbose:
                            print(
                                f"Filtering out object with area {object_area} pixels"
                            )
                        continue

                    conf_value = int(scores[mask_idx] * 255)  # Scale to 0-255

                    # Update the mask and confidence arrays
                    h, w = binary_mask.shape
                    valid_h = min(h, src.height - j)
                    valid_w = min(w, src.width - i)

                    if valid_h > 0 and valid_w > 0:
                        # Use maximum for overlapping regions in the mask
                        mask_array[j : j + valid_h, i : i + valid_w] = np.maximum(
                            mask_array[j : j + valid_h, i : i + valid_w],
                            binary_mask[:valid_h, :valid_w],
                        )

                        # For confidence, only update where mask is positive
                        # and confidence is higher than existing
                        mask_region = binary_mask[:valid_h, :valid_w] > 0
                        if np.any(mask_region):
                            # Only update where mask is positive and new confidence is higher
                            current_conf = conf_array[
                                j : j + valid_h, i : i + valid_w
                            ]

                            # Where to update confidence (mask positive & higher confidence)
                            update_mask = np.logical_and(
                                mask_region,
                                np.logical_or(
                                    current_conf == 0, current_conf < conf_value
                                ),
                            )

                            if np.any(update_mask):
                                conf_array[j : j + valid_h, i : i + valid_w][
                                    update_mask
                                ] = conf_value

        # Write to GeoTIFF
        with rasterio.open(output_path, "w", **output_profile) as dst:
            dst.write(mask_array, 1)
            dst.write(conf_array, 2)

        print(f"Masks with confidence values saved to {output_path}")
        return output_path

initialize_model(model, num_classes=2)

Initialize a deep learning model for object detection.

Parameters:

Name	Type	Description	Default
model	Module	A pre-initialized model object.	required
num_classes	int	Number of classes for detection.	2

Returns:

Type	Description
Any	torch.nn.Module: A deep learning model for object detection.

Source code in geoai/extract.py

def initialize_model(self, model: Optional[Any], num_classes: int = 2) -> Any:
    """Initialize a deep learning model for object detection.

    Args:
        model (torch.nn.Module): A pre-initialized model object.
        num_classes (int): Number of classes for detection.

    Returns:
        torch.nn.Module: A deep learning model for object detection.
    """

    if model is None:  # Initialize Mask R-CNN model with ResNet50 backbone.
        # Standard image mean and std for pre-trained models
        image_mean = [0.485, 0.456, 0.406]
        image_std = [0.229, 0.224, 0.225]

        # Create model with explicit normalization parameters
        model = maskrcnn_resnet50_fpn(
            weights=None,
            progress=False,
            num_classes=num_classes,  # Background + object
            weights_backbone=None,
            # These parameters ensure consistent normalization
            image_mean=image_mean,
            image_std=image_std,
        )

    model.to(self.device)
    return model

load_weights(model_path)

Load weights from file with error handling for different formats.

Parameters:

Name	Type	Description	Default
model_path	str	Path to model weights	required

Source code in geoai/extract.py

def load_weights(self, model_path: str) -> None:
    """
    Load weights from file with error handling for different formats.

    Args:
        model_path: Path to model weights
    """
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model file not found: {model_path}")

    try:
        state_dict = torch.load(model_path, map_location=self.device)

        # Handle different state dict formats
        if isinstance(state_dict, dict):
            if "model" in state_dict:
                state_dict = state_dict["model"]
            elif "state_dict" in state_dict:
                state_dict = state_dict["state_dict"]

        # Try to load state dict
        try:
            self.model.load_state_dict(state_dict)
            print("Model loaded successfully")
        except Exception as e:
            print(f"Error loading model: {e}")
            print("Attempting to fix state_dict keys...")

            # Try to fix state_dict keys (remove module prefix if needed)
            new_state_dict = {}
            for k, v in state_dict.items():
                if k.startswith("module."):
                    new_state_dict[k[7:]] = v
                else:
                    new_state_dict[k] = v

            self.model.load_state_dict(new_state_dict)
            print("Model loaded successfully after key fixing")

    except Exception as e:
        raise RuntimeError(f"Failed to load model: {e}")

mask_to_polygons(mask, **kwargs)

Convert binary mask to polygon contours using OpenCV.

Parameters:

Name	Type	Description	Default
mask	ndarray	Binary mask as numpy array	required
**kwargs	Any	Optional parameters: simplify_tolerance: Tolerance for polygon simplification mask_threshold: Threshold for mask binarization min_object_area: Minimum area in pixels to keep an object max_object_area: Maximum area in pixels to keep an object	{}

Returns:

Type	Description
List[Polygon]	List of polygons as lists of (x, y) coordinates

Source code in geoai/extract.py

def mask_to_polygons(self, mask: np.ndarray, **kwargs: Any) -> List[Polygon]:
    """
    Convert binary mask to polygon contours using OpenCV.

    Args:
        mask: Binary mask as numpy array
        **kwargs: Optional parameters:
            simplify_tolerance: Tolerance for polygon simplification
            mask_threshold: Threshold for mask binarization
            min_object_area: Minimum area in pixels to keep an object
            max_object_area: Maximum area in pixels to keep an object

    Returns:
        List of polygons as lists of (x, y) coordinates
    """
    import cv2  # Lazy import to avoid QGIS opencv conflicts

    # Get parameters from kwargs or use instance defaults
    simplify_tolerance = kwargs.get("simplify_tolerance", self.simplify_tolerance)
    mask_threshold = kwargs.get("mask_threshold", self.mask_threshold)
    min_object_area = kwargs.get("min_object_area", self.min_object_area)
    max_object_area = kwargs.get("max_object_area", self.max_object_area)

    # Ensure binary mask
    mask = (mask > mask_threshold).astype(np.uint8)

    # Optional: apply morphological operations to improve mask quality
    kernel = np.ones((3, 3), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

    # Find contours
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # Convert to list of [x, y] coordinates
    polygons = []
    for contour in contours:
        # Filter out too small contours
        if contour.shape[0] < 3 or cv2.contourArea(contour) < min_object_area:
            continue

        # Filter out too large contours
        if (
            max_object_area is not None
            and cv2.contourArea(contour) > max_object_area
        ):
            continue

        # Simplify contour if it has many points
        if contour.shape[0] > 50:
            epsilon = simplify_tolerance * cv2.arcLength(contour, True)
            contour = cv2.approxPolyDP(contour, epsilon, True)

        # Convert to list of [x, y] coordinates
        polygon = contour.reshape(-1, 2).tolist()
        polygons.append(polygon)

    return polygons

masks_to_vector(mask_path, output_path=None, simplify_tolerance=None, mask_threshold=None, min_object_area=None, max_object_area=None, nms_iou_threshold=None, regularize=True, angle_threshold=15, rectangularity_threshold=0.7)

Convert an object mask GeoTIFF to vector polygons and save as GeoJSON.

Parameters:

Name	Type	Description	Default
mask_path	str	Path to the object masks GeoTIFF	required
output_path	Optional[str]	Path to save the output GeoJSON or Parquet file (default: mask_path with .geojson extension)	None
simplify_tolerance	Optional[float]	Tolerance for polygon simplification (default: self.simplify_tolerance)	None
mask_threshold	Optional[float]	Threshold for mask binarization (default: self.mask_threshold)	None
min_object_area	Optional[int]	Minimum area in pixels to keep an object (default: self.min_object_area)	None
max_object_area	Optional[int]	Minimum area in pixels to keep an object (default: self.max_object_area)	None
nms_iou_threshold	Optional[float]	IoU threshold for non-maximum suppression (default: self.nms_iou_threshold)	None
regularize	bool	Whether to regularize objects to right angles (default: True)	True
angle_threshold	int	Maximum deviation from 90 degrees for regularization (default: 15)	15
rectangularity_threshold	float	Threshold for rectangle simplification (default: 0.7)	0.7

Returns:

Type	Description
GeoDataFrame	GeoDataFrame with objects

Source code in geoai/extract.py

def masks_to_vector(
    self,
    mask_path: str,
    output_path: Optional[str] = None,
    simplify_tolerance: Optional[float] = None,
    mask_threshold: Optional[float] = None,
    min_object_area: Optional[int] = None,
    max_object_area: Optional[int] = None,
    nms_iou_threshold: Optional[float] = None,
    regularize: bool = True,
    angle_threshold: int = 15,
    rectangularity_threshold: float = 0.7,
) -> gpd.GeoDataFrame:
    """
    Convert an object mask GeoTIFF to vector polygons and save as GeoJSON.

    Args:
        mask_path: Path to the object masks GeoTIFF
        output_path: Path to save the output GeoJSON or Parquet file (default: mask_path with .geojson extension)
        simplify_tolerance: Tolerance for polygon simplification (default: self.simplify_tolerance)
        mask_threshold: Threshold for mask binarization (default: self.mask_threshold)
        min_object_area: Minimum area in pixels to keep an object (default: self.min_object_area)
        max_object_area: Minimum area in pixels to keep an object (default: self.max_object_area)
        nms_iou_threshold: IoU threshold for non-maximum suppression (default: self.nms_iou_threshold)
        regularize: Whether to regularize objects to right angles (default: True)
        angle_threshold: Maximum deviation from 90 degrees for regularization (default: 15)
        rectangularity_threshold: Threshold for rectangle simplification (default: 0.7)

    Returns:
        GeoDataFrame with objects
    """
    import cv2  # Lazy import to avoid QGIS opencv conflicts

    # Use class defaults if parameters not provided
    simplify_tolerance = (
        simplify_tolerance
        if simplify_tolerance is not None
        else self.simplify_tolerance
    )
    mask_threshold = (
        mask_threshold if mask_threshold is not None else self.mask_threshold
    )
    min_object_area = (
        min_object_area if min_object_area is not None else self.min_object_area
    )
    max_object_area = (
        max_object_area if max_object_area is not None else self.max_object_area
    )
    nms_iou_threshold = (
        nms_iou_threshold
        if nms_iou_threshold is not None
        else self.nms_iou_threshold
    )

    # Set default output path if not provided
    # if output_path is None:
    #     output_path = os.path.splitext(mask_path)[0] + ".geojson"

    print(f"Converting mask to GeoJSON with parameters:")
    print(f"- Mask threshold: {mask_threshold}")
    print(f"- Min object area: {min_object_area}")
    print(f"- Max object area: {max_object_area}")
    print(f"- Simplify tolerance: {simplify_tolerance}")
    print(f"- NMS IoU threshold: {nms_iou_threshold}")
    print(f"- Regularize objects: {regularize}")
    if regularize:
        print(f"- Angle threshold: {angle_threshold}° from 90°")
        print(f"- Rectangularity threshold: {rectangularity_threshold*100}%")

    # Open the mask raster
    with rasterio.open(mask_path) as src:
        # Read the mask data
        mask_data = src.read(1)
        transform = src.transform
        crs = src.crs

        # Print mask statistics
        print(f"Mask dimensions: {mask_data.shape}")
        print(f"Mask value range: {mask_data.min()} to {mask_data.max()}")

        # Prepare for connected component analysis
        # Binarize the mask based on threshold
        binary_mask = (mask_data > (mask_threshold * 255)).astype(np.uint8)

        # Apply morphological operations for better results (optional)
        kernel = np.ones((3, 3), np.uint8)
        binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, kernel)

        # Find connected components
        num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(
            binary_mask, connectivity=8
        )

        print(
            f"Found {num_labels-1} potential objects"
        )  # Subtract 1 for background

        # Create list to store polygons and confidence values
        all_polygons = []
        all_confidences = []

        # Process each component (skip the first one which is background)
        for i in tqdm(range(1, num_labels)):
            # Extract this object
            area = stats[i, cv2.CC_STAT_AREA]

            # Skip if too small
            if area < min_object_area:
                continue

            # Skip if too large
            if max_object_area is not None and area > max_object_area:
                continue

            # Create a mask for this object
            object_mask = (labels == i).astype(np.uint8)

            # Find contours
            contours, _ = cv2.findContours(
                object_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
            )

            # Process each contour
            for contour in contours:
                # Skip if too few points
                if contour.shape[0] < 3:
                    continue

                # Simplify contour if it has many points
                if contour.shape[0] > 50 and simplify_tolerance > 0:
                    epsilon = simplify_tolerance * cv2.arcLength(contour, True)
                    contour = cv2.approxPolyDP(contour, epsilon, True)

                # Convert to list of (x, y) coordinates
                polygon_points = contour.reshape(-1, 2)

                # Convert pixel coordinates to geographic coordinates
                geo_points = []
                for x, y in polygon_points:
                    gx, gy = transform * (x, y)
                    geo_points.append((gx, gy))

                # Create Shapely polygon
                if len(geo_points) >= 3:
                    try:
                        shapely_poly = Polygon(geo_points)
                        if shapely_poly.is_valid and shapely_poly.area > 0:
                            all_polygons.append(shapely_poly)

                            # Calculate "confidence" as normalized size
                            # This is a proxy since we don't have model confidence scores
                            normalized_size = min(1.0, area / 1000)  # Cap at 1.0
                            all_confidences.append(normalized_size)
                    except Exception as e:
                        print(f"Error creating polygon: {e}")

        print(f"Created {len(all_polygons)} valid polygons")

        # Create GeoDataFrame
        if not all_polygons:
            print("No valid polygons found")
            return None

        gdf = gpd.GeoDataFrame(
            {
                "geometry": all_polygons,
                "confidence": all_confidences,
                "class": 1,  # Object class
            },
            crs=crs,
        )

        # Apply non-maximum suppression to remove overlapping polygons
        gdf = self.filter_overlapping_polygons(
            gdf, nms_iou_threshold=nms_iou_threshold
        )

        print(f"Object count after NMS filtering: {len(gdf)}")

        # Apply regularization if requested
        if regularize and len(gdf) > 0:
            # Convert pixel area to geographic units for min_area parameter
            # Estimate pixel size in geographic units
            with rasterio.open(mask_path) as src:
                pixel_size_x = src.transform[
                    0
                ]  # width of a pixel in geographic units
                pixel_size_y = abs(
                    src.transform[4]
                )  # height of a pixel in geographic units
                avg_pixel_area = pixel_size_x * pixel_size_y

            # Use 10 pixels as minimum area in geographic units
            min_geo_area = 10 * avg_pixel_area

            # Regularize objects
            gdf = self.regularize_objects(
                gdf,
                min_area=min_geo_area,
                angle_threshold=angle_threshold,
                rectangularity_threshold=rectangularity_threshold,
            )

        # Save to file
        if output_path:
            if output_path.endswith(".parquet"):
                gdf.to_parquet(output_path)
            else:
                gdf.to_file(output_path)
            print(f"Saved {len(gdf)} objects to {output_path}")

        return gdf

process_raster(raster_path, output_path=None, batch_size=4, filter_edges=True, edge_buffer=20, band_indexes=None, **kwargs)

Process a raster file to extract objects with customizable parameters.

Parameters:

Name	Type	Description	Default
raster_path	str	Path to input raster file	required
output_path	Optional[str]	Path to output GeoJSON or Parquet file (optional)	None
batch_size	int	Batch size for processing	4
filter_edges	bool	Whether to filter out objects at the edges of the image	True
edge_buffer	int	Size of edge buffer in pixels to filter out objects (if filter_edges=True)	20
band_indexes	Optional[List[int]]	List of band indexes to use (if None, use all bands)	None
**kwargs	Any	Additional parameters: confidence_threshold: Minimum confidence score to keep a detection (0.0-1.0) overlap: Overlap between adjacent tiles (0.0-1.0) chip_size: Size of image chips for processing (height, width) nms_iou_threshold: IoU threshold for non-maximum suppression (0.0-1.0) mask_threshold: Threshold for mask binarization (0.0-1.0) min_object_area: Minimum area in pixels to keep an object simplify_tolerance: Tolerance for polygon simplification	{}

Returns:

Type	Description
GeoDataFrame	GeoDataFrame with objects

Source code in geoai/extract.py

@torch.no_grad()
def process_raster(
    self,
    raster_path: str,
    output_path: Optional[str] = None,
    batch_size: int = 4,
    filter_edges: bool = True,
    edge_buffer: int = 20,
    band_indexes: Optional[List[int]] = None,
    **kwargs: Any,
) -> "gpd.GeoDataFrame":
    """
    Process a raster file to extract objects with customizable parameters.

    Args:
        raster_path: Path to input raster file
        output_path: Path to output GeoJSON or Parquet file (optional)
        batch_size: Batch size for processing
        filter_edges: Whether to filter out objects at the edges of the image
        edge_buffer: Size of edge buffer in pixels to filter out objects (if filter_edges=True)
        band_indexes: List of band indexes to use (if None, use all bands)
        **kwargs: Additional parameters:
            confidence_threshold: Minimum confidence score to keep a detection (0.0-1.0)
            overlap: Overlap between adjacent tiles (0.0-1.0)
            chip_size: Size of image chips for processing (height, width)
            nms_iou_threshold: IoU threshold for non-maximum suppression (0.0-1.0)
            mask_threshold: Threshold for mask binarization (0.0-1.0)
            min_object_area: Minimum area in pixels to keep an object
            simplify_tolerance: Tolerance for polygon simplification

    Returns:
        GeoDataFrame with objects
    """
    # Get parameters from kwargs or use instance defaults
    confidence_threshold = kwargs.get(
        "confidence_threshold", self.confidence_threshold
    )
    overlap = kwargs.get("overlap", self.overlap)
    chip_size = kwargs.get("chip_size", self.chip_size)
    nms_iou_threshold = kwargs.get("nms_iou_threshold", self.nms_iou_threshold)
    mask_threshold = kwargs.get("mask_threshold", self.mask_threshold)
    min_object_area = kwargs.get("min_object_area", self.min_object_area)
    max_object_area = kwargs.get("max_object_area", self.max_object_area)
    simplify_tolerance = kwargs.get("simplify_tolerance", self.simplify_tolerance)

    # Print parameters being used
    print(f"Processing with parameters:")
    print(f"- Confidence threshold: {confidence_threshold}")
    print(f"- Tile overlap: {overlap}")
    print(f"- Chip size: {chip_size}")
    print(f"- NMS IoU threshold: {nms_iou_threshold}")
    print(f"- Mask threshold: {mask_threshold}")
    print(f"- Min object area: {min_object_area}")
    print(f"- Max object area: {max_object_area}")
    print(f"- Simplify tolerance: {simplify_tolerance}")
    print(f"- Filter edge objects: {filter_edges}")
    if filter_edges:
        print(f"- Edge buffer size: {edge_buffer} pixels")

    # Create dataset
    dataset = CustomDataset(
        raster_path=raster_path,
        chip_size=chip_size,
        overlap=overlap,
        band_indexes=band_indexes,
    )
    self.raster_stats = dataset.raster_stats

    # Custom collate function to handle Shapely objects
    def custom_collate(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        Custom collate function that handles Shapely geometries
        by keeping them as Python objects rather than trying to collate them.
        """
        elem = batch[0]
        if isinstance(elem, dict):
            result = {}
            for key in elem:
                if key == "bbox":
                    # Don't collate shapely objects, keep as list
                    result[key] = [d[key] for d in batch]
                else:
                    # For tensors and other collatable types
                    try:
                        result[key] = (
                            torch.utils.data._utils.collate.default_collate(
                                [d[key] for d in batch]
                            )
                        )
                    except TypeError:
                        # Fall back to list for non-collatable types
                        result[key] = [d[key] for d in batch]
            return result
        else:
            # Default collate for non-dict types
            return torch.utils.data._utils.collate.default_collate(batch)

    # Create dataloader with simple indexing and custom collate
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=0,
        collate_fn=custom_collate,
    )

    # Process batches
    all_polygons = []
    all_scores = []

    print(f"Processing raster with {len(dataloader)} batches")
    for batch in tqdm(dataloader):
        # Move images to device
        images = batch["image"].to(self.device)
        coords = batch["coords"]  # (i, j) coordinates in pixels
        bboxes = batch[
            "bbox"
        ]  # Geographic bounding boxes - now a list, not a tensor

        # Run inference
        predictions = self.model(images)

        # Process predictions
        for idx, prediction in enumerate(predictions):
            masks = prediction["masks"].cpu().numpy()
            scores = prediction["scores"].cpu().numpy()
            labels = prediction["labels"].cpu().numpy()

            # Skip if no predictions
            if len(scores) == 0:
                continue

            # Filter by confidence threshold
            valid_indices = scores >= confidence_threshold
            masks = masks[valid_indices]
            scores = scores[valid_indices]
            labels = labels[valid_indices]

            # Skip if no valid predictions
            if len(scores) == 0:
                continue

            # Get window coordinates
            # The coords might be in different formats depending on batch handling
            if isinstance(coords, list):
                # If coords is a list of tuples
                coord_item = coords[idx]
                if isinstance(coord_item, tuple) and len(coord_item) == 2:
                    i, j = coord_item
                elif isinstance(coord_item, torch.Tensor):
                    i, j = coord_item.cpu().numpy().tolist()
                else:
                    print(f"Unexpected coords format: {type(coord_item)}")
                    continue
            elif isinstance(coords, torch.Tensor):
                # If coords is a tensor of shape [batch_size, 2]
                i, j = coords[idx].cpu().numpy().tolist()
            else:
                print(f"Unexpected coords type: {type(coords)}")
                continue

            # Get window size
            if isinstance(batch["window_size"], list):
                window_item = batch["window_size"][idx]
                if isinstance(window_item, tuple) and len(window_item) == 2:
                    window_width, window_height = window_item
                elif isinstance(window_item, torch.Tensor):
                    window_width, window_height = window_item.cpu().numpy().tolist()
                else:
                    print(f"Unexpected window_size format: {type(window_item)}")
                    continue
            elif isinstance(batch["window_size"], torch.Tensor):
                window_width, window_height = (
                    batch["window_size"][idx].cpu().numpy().tolist()
                )
            else:
                print(f"Unexpected window_size type: {type(batch['window_size'])}")
                continue

            # Process masks to polygons
            for mask_idx, mask in enumerate(masks):
                # Get binary mask
                binary_mask = mask[0]  # Get binary mask

                # Convert mask to polygon with custom parameters
                contours = self.mask_to_polygons(
                    binary_mask,
                    simplify_tolerance=simplify_tolerance,
                    mask_threshold=mask_threshold,
                    min_object_area=min_object_area,
                    max_object_area=max_object_area,
                )

                # Skip if no valid polygons
                if not contours:
                    continue

                # Transform polygons to geographic coordinates
                with rasterio.open(raster_path) as src:
                    transform = src.transform

                    for contour in contours:
                        # Convert polygon to global coordinates
                        global_polygon = []
                        for x, y in contour:
                            # Adjust coordinates based on window position
                            gx, gy = transform * (i + x, j + y)
                            global_polygon.append((gx, gy))

                        # Create Shapely polygon
                        if len(global_polygon) >= 3:
                            try:
                                shapely_poly = Polygon(global_polygon)
                                if shapely_poly.is_valid and shapely_poly.area > 0:
                                    all_polygons.append(shapely_poly)
                                    all_scores.append(float(scores[mask_idx]))
                            except Exception as e:
                                print(f"Error creating polygon: {e}")

    # Create GeoDataFrame
    if not all_polygons:
        print("No valid polygons found")
        return None

    gdf = gpd.GeoDataFrame(
        {
            "geometry": all_polygons,
            "confidence": all_scores,
            "class": 1,  # Object class
        },
        crs=dataset.crs,
    )

    # Remove overlapping polygons with custom threshold
    gdf = self.filter_overlapping_polygons(gdf, nms_iou_threshold=nms_iou_threshold)

    # Filter edge objects if requested
    if filter_edges:
        gdf = self.filter_edge_objects(gdf, raster_path, edge_buffer=edge_buffer)

    # Save to file if requested
    if output_path:
        if output_path.endswith(".parquet"):
            gdf.to_parquet(output_path)
        else:
            gdf.to_file(output_path, driver="GeoJSON")
        print(f"Saved {len(gdf)} objects to {output_path}")

    return gdf

regularize_objects(gdf, min_area=10, angle_threshold=15, orthogonality_threshold=0.3, rectangularity_threshold=0.7)

Regularize objects to enforce right angles and rectangular shapes.

Parameters:

Name	Type	Description	Default
gdf	GeoDataFrame	GeoDataFrame with objects	required
min_area	int	Minimum area in square units to keep an object	10
angle_threshold	int	Maximum deviation from 90 degrees to consider an angle as orthogonal (degrees)	15
orthogonality_threshold	float	Percentage of angles that must be orthogonal for an object to be regularized	0.3
rectangularity_threshold	float	Minimum area ratio to Object's oriented bounding box for rectangular simplification	0.7

Returns:

Type	Description
GeoDataFrame	GeoDataFrame with regularized objects

Source code in geoai/extract.py

def regularize_objects(
    self,
    gdf: gpd.GeoDataFrame,
    min_area: int = 10,
    angle_threshold: int = 15,
    orthogonality_threshold: float = 0.3,
    rectangularity_threshold: float = 0.7,
) -> gpd.GeoDataFrame:
    """
    Regularize objects to enforce right angles and rectangular shapes.

    Args:
        gdf: GeoDataFrame with objects
        min_area: Minimum area in square units to keep an object
        angle_threshold: Maximum deviation from 90 degrees to consider an angle as orthogonal (degrees)
        orthogonality_threshold: Percentage of angles that must be orthogonal for an object to be regularized
        rectangularity_threshold: Minimum area ratio to Object's oriented bounding box for rectangular simplification

    Returns:
        GeoDataFrame with regularized objects
    """
    import math

    import cv2
    import geopandas as gpd
    import numpy as np
    from shapely.affinity import rotate, translate
    from shapely.geometry import MultiPolygon, Polygon, box
    from tqdm import tqdm

    def get_angle(
        p1: Tuple[float, float], p2: Tuple[float, float], p3: Tuple[float, float]
    ) -> float:
        """Calculate angle between three points in degrees (0-180)"""
        a = np.array(p1)
        b = np.array(p2)
        c = np.array(p3)

        ba = a - b
        bc = c - b

        cosine_angle = np.dot(ba, bc) / (np.linalg.norm(ba) * np.linalg.norm(bc))
        # Handle numerical errors that could push cosine outside [-1, 1]
        cosine_angle = np.clip(cosine_angle, -1.0, 1.0)
        angle = np.degrees(np.arccos(cosine_angle))

        return angle

    def is_orthogonal(angle: float, threshold: int = angle_threshold) -> bool:
        """Check if angle is close to 90 degrees"""
        return abs(angle - 90) <= threshold

    def calculate_dominant_direction(polygon: Polygon) -> float:
        """Find the dominant direction of a polygon using PCA"""
        # Extract coordinates
        coords = np.array(polygon.exterior.coords)

        # Mean center the coordinates
        mean = np.mean(coords, axis=0)
        centered_coords = coords - mean

        # Calculate covariance matrix and its eigenvalues/eigenvectors
        cov_matrix = np.cov(centered_coords.T)
        eigenvalues, eigenvectors = np.linalg.eig(cov_matrix)

        # Get the index of the largest eigenvalue
        largest_idx = np.argmax(eigenvalues)

        # Get the corresponding eigenvector (principal axis)
        principal_axis = eigenvectors[:, largest_idx]

        # Calculate the angle in degrees
        angle_rad = np.arctan2(principal_axis[1], principal_axis[0])
        angle_deg = np.degrees(angle_rad)

        # Normalize to range 0-180
        if angle_deg < 0:
            angle_deg += 180

        return angle_deg

    def create_oriented_envelope(polygon: Polygon, angle_deg: float) -> Polygon:
        """Create an oriented minimum area rectangle for the polygon"""
        # Create a rotated rectangle using OpenCV method (more robust than Shapely methods)
        coords = np.array(polygon.exterior.coords)[:-1].astype(
            np.float32
        )  # Skip the last point (same as first)

        # Use OpenCV's minAreaRect
        rect = cv2.minAreaRect(coords)
        box_points = cv2.boxPoints(rect)

        # Convert to shapely polygon
        oriented_box = Polygon(box_points)

        return oriented_box

    def get_rectangularity(polygon: Polygon, oriented_box: Polygon) -> float:
        """Calculate the rectangularity (area ratio to its oriented bounding box)"""
        if oriented_box.area == 0:
            return 0
        return polygon.area / oriented_box.area

    def check_orthogonality(polygon: Polygon) -> float:
        """Check what percentage of angles in the polygon are orthogonal"""
        coords = list(polygon.exterior.coords)
        if len(coords) <= 4:  # Triangle or point
            return 0

        # Remove last point (same as first)
        coords = coords[:-1]

        orthogonal_count = 0
        total_angles = len(coords)

        for i in range(total_angles):
            p1 = coords[i]
            p2 = coords[(i + 1) % total_angles]
            p3 = coords[(i + 2) % total_angles]

            angle = get_angle(p1, p2, p3)
            if is_orthogonal(angle):
                orthogonal_count += 1

        return orthogonal_count / total_angles

    def simplify_to_rectangle(polygon: Polygon) -> Polygon:
        """Simplify a polygon to a rectangle using its oriented bounding box"""
        # Get dominant direction
        angle = calculate_dominant_direction(polygon)

        # Create oriented envelope
        rect = create_oriented_envelope(polygon, angle)

        return rect

    if gdf is None or len(gdf) == 0:
        print("No Objects to regularize")
        return gdf

    print(f"Regularizing {len(gdf)} objects...")
    print(f"- Angle threshold: {angle_threshold}° from 90°")
    print(f"- Min orthogonality: {orthogonality_threshold*100}% of angles")
    print(
        f"- Min rectangularity: {rectangularity_threshold*100}% of bounding box area"
    )

    # Create a copy to avoid modifying the original
    result_gdf = gdf.copy()

    # Track statistics
    total_objects = len(gdf)
    regularized_count = 0
    rectangularized_count = 0

    # Process each Object
    for idx, row in tqdm(gdf.iterrows(), total=len(gdf)):
        geom = row.geometry

        # Skip invalid or empty geometries
        if geom is None or geom.is_empty:
            continue

        # Handle MultiPolygons by processing the largest part
        if isinstance(geom, MultiPolygon):
            areas = [p.area for p in geom.geoms]
            if not areas:
                continue
            geom = list(geom.geoms)[np.argmax(areas)]

        # Filter out tiny Objects
        if geom.area < min_area:
            continue

        # Check orthogonality
        orthogonality = check_orthogonality(geom)

        # Create oriented envelope
        oriented_box = create_oriented_envelope(
            geom, calculate_dominant_direction(geom)
        )

        # Check rectangularity
        rectangularity = get_rectangularity(geom, oriented_box)

        # Decide how to regularize
        if rectangularity >= rectangularity_threshold:
            # Object is already quite rectangular, simplify to a rectangle
            result_gdf.at[idx, "geometry"] = oriented_box
            result_gdf.at[idx, "regularized"] = "rectangle"
            rectangularized_count += 1
        elif orthogonality >= orthogonality_threshold:
            # Object has many orthogonal angles but isn't rectangular
            # Could implement more sophisticated regularization here
            # For now, we'll still use the oriented rectangle
            result_gdf.at[idx, "geometry"] = oriented_box
            result_gdf.at[idx, "regularized"] = "orthogonal"
            regularized_count += 1
        else:
            # Object doesn't have clear orthogonal structure
            # Keep original but flag as unmodified
            result_gdf.at[idx, "regularized"] = "original"

    # Report statistics
    print(f"Regularization completed:")
    print(f"- Total objects: {total_objects}")
    print(
        f"- Rectangular objects: {rectangularized_count} ({rectangularized_count/total_objects*100:.1f}%)"
    )
    print(
        f"- Other regularized objects: {regularized_count} ({regularized_count/total_objects*100:.1f}%)"
    )
    print(
        f"- Unmodified objects: {total_objects-rectangularized_count-regularized_count} ({(total_objects-rectangularized_count-regularized_count)/total_objects*100:.1f}%)"
    )

    return result_gdf

save_masks_as_geotiff(raster_path, output_path=None, batch_size=4, verbose=False, **kwargs)

Process a raster file to extract object masks and save as GeoTIFF.

Parameters:

Name	Type	Description	Default
raster_path	str	Path to input raster file	required
output_path	Optional[str]	Path to output GeoTIFF file (optional, default: input_masks.tif)	None
batch_size	int	Batch size for processing	4
verbose	bool	Whether to print detailed processing information	False
**kwargs	Any	Additional parameters: confidence_threshold: Minimum confidence score to keep a detection (0.0-1.0) chip_size: Size of image chips for processing (height, width) mask_threshold: Threshold for mask binarization (0.0-1.0)	{}

Returns:

Type	Description
str	Path to the saved GeoTIFF file

Source code in geoai/extract.py

def save_masks_as_geotiff(
    self,
    raster_path: str,
    output_path: Optional[str] = None,
    batch_size: int = 4,
    verbose: bool = False,
    **kwargs: Any,
) -> str:
    """
    Process a raster file to extract object masks and save as GeoTIFF.

    Args:
        raster_path: Path to input raster file
        output_path: Path to output GeoTIFF file (optional, default: input_masks.tif)
        batch_size: Batch size for processing
        verbose: Whether to print detailed processing information
        **kwargs: Additional parameters:
            confidence_threshold: Minimum confidence score to keep a detection (0.0-1.0)
            chip_size: Size of image chips for processing (height, width)
            mask_threshold: Threshold for mask binarization (0.0-1.0)

    Returns:
        Path to the saved GeoTIFF file
    """

    # Get parameters from kwargs or use instance defaults
    confidence_threshold = kwargs.get(
        "confidence_threshold", self.confidence_threshold
    )
    chip_size = kwargs.get("chip_size", self.chip_size)
    mask_threshold = kwargs.get("mask_threshold", self.mask_threshold)
    overlap = kwargs.get("overlap", self.overlap)

    # Set default output path if not provided
    if output_path is None:
        output_path = os.path.splitext(raster_path)[0] + "_masks.tif"

    # Print parameters being used
    print(f"Processing masks with parameters:")
    print(f"- Confidence threshold: {confidence_threshold}")
    print(f"- Chip size: {chip_size}")
    print(f"- Mask threshold: {mask_threshold}")

    # Create dataset
    dataset = CustomDataset(
        raster_path=raster_path,
        chip_size=chip_size,
        overlap=overlap,
        verbose=verbose,
    )

    # Store a flag to avoid repetitive messages
    self.raster_stats = dataset.raster_stats
    seen_warnings = {
        "bands": False,
        "resize": {},  # Dictionary to track resize warnings by shape
    }

    # Open original raster to get metadata
    with rasterio.open(raster_path) as src:
        # Create output binary mask raster with same dimensions as input
        output_profile = src.profile.copy()
        output_profile.update(
            dtype=rasterio.uint8,
            count=1,  # Single band for object mask
            compress="lzw",
            nodata=0,
        )

        # Create output mask raster
        with rasterio.open(output_path, "w", **output_profile) as dst:
            # Initialize mask with zeros
            mask_array = np.zeros((src.height, src.width), dtype=np.uint8)

            # Custom collate function to handle Shapely objects
            def custom_collate(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
                """Custom collate function for DataLoader"""
                elem = batch[0]
                if isinstance(elem, dict):
                    result = {}
                    for key in elem:
                        if key == "bbox":
                            # Don't collate shapely objects, keep as list
                            result[key] = [d[key] for d in batch]
                        else:
                            # For tensors and other collatable types
                            try:
                                result[key] = (
                                    torch.utils.data._utils.collate.default_collate(
                                        [d[key] for d in batch]
                                    )
                                )
                            except TypeError:
                                # Fall back to list for non-collatable types
                                result[key] = [d[key] for d in batch]
                    return result
                else:
                    # Default collate for non-dict types
                    return torch.utils.data._utils.collate.default_collate(batch)

            # Create dataloader
            dataloader = torch.utils.data.DataLoader(
                dataset,
                batch_size=batch_size,
                shuffle=False,
                num_workers=0,
                collate_fn=custom_collate,
            )

            # Process batches
            print(f"Processing raster with {len(dataloader)} batches")
            for batch in tqdm(dataloader):
                # Move images to device
                images = batch["image"].to(self.device)
                coords = batch["coords"]  # (i, j) coordinates in pixels

                # Run inference
                with torch.no_grad():
                    predictions = self.model(images)

                # Process predictions
                for idx, prediction in enumerate(predictions):
                    masks = prediction["masks"].cpu().numpy()
                    scores = prediction["scores"].cpu().numpy()

                    # Skip if no predictions
                    if len(scores) == 0:
                        continue

                    # Filter by confidence threshold
                    valid_indices = scores >= confidence_threshold
                    masks = masks[valid_indices]
                    scores = scores[valid_indices]

                    # Skip if no valid predictions
                    if len(scores) == 0:
                        continue

                    # Get window coordinates
                    if isinstance(coords, list):
                        coord_item = coords[idx]
                        if isinstance(coord_item, tuple) and len(coord_item) == 2:
                            i, j = coord_item
                        elif isinstance(coord_item, torch.Tensor):
                            i, j = coord_item.cpu().numpy().tolist()
                        else:
                            print(f"Unexpected coords format: {type(coord_item)}")
                            continue
                    elif isinstance(coords, torch.Tensor):
                        i, j = coords[idx].cpu().numpy().tolist()
                    else:
                        print(f"Unexpected coords type: {type(coords)}")
                        continue

                    # Get window size
                    if isinstance(batch["window_size"], list):
                        window_item = batch["window_size"][idx]
                        if isinstance(window_item, tuple) and len(window_item) == 2:
                            window_width, window_height = window_item
                        elif isinstance(window_item, torch.Tensor):
                            window_width, window_height = (
                                window_item.cpu().numpy().tolist()
                            )
                        else:
                            print(
                                f"Unexpected window_size format: {type(window_item)}"
                            )
                            continue
                    elif isinstance(batch["window_size"], torch.Tensor):
                        window_width, window_height = (
                            batch["window_size"][idx].cpu().numpy().tolist()
                        )
                    else:
                        print(
                            f"Unexpected window_size type: {type(batch['window_size'])}"
                        )
                        continue

                    # Combine all masks for this window
                    combined_mask = np.zeros(
                        (window_height, window_width), dtype=np.uint8
                    )

                    for mask in masks:
                        # Get the binary mask
                        binary_mask = (mask[0] > mask_threshold).astype(
                            np.uint8
                        ) * 255

                        # Handle size mismatch - resize binary_mask if needed
                        mask_h, mask_w = binary_mask.shape
                        if mask_h != window_height or mask_w != window_width:
                            resize_key = f"{(mask_h, mask_w)}->{(window_height, window_width)}"
                            if resize_key not in seen_warnings["resize"]:
                                if verbose:
                                    print(
                                        f"Resizing mask from {binary_mask.shape} to {(window_height, window_width)}"
                                    )
                                else:
                                    if not seen_warnings[
                                        "resize"
                                    ]:  # If this is the first resize warning
                                        print(
                                            f"Resizing masks at image edges (set verbose=True for details)"
                                        )
                                seen_warnings["resize"][resize_key] = True

                            # Crop or pad the binary mask to match window size
                            resized_mask = np.zeros(
                                (window_height, window_width), dtype=np.uint8
                            )
                            copy_h = min(mask_h, window_height)
                            copy_w = min(mask_w, window_width)
                            resized_mask[:copy_h, :copy_w] = binary_mask[
                                :copy_h, :copy_w
                            ]
                            binary_mask = resized_mask

                        # Update combined mask (taking maximum where masks overlap)
                        combined_mask = np.maximum(combined_mask, binary_mask)

                    # Write combined mask to output array
                    # Handle edge cases where window might be smaller than chip size
                    h, w = combined_mask.shape
                    valid_h = min(h, src.height - j)
                    valid_w = min(w, src.width - i)

                    if valid_h > 0 and valid_w > 0:
                        mask_array[j : j + valid_h, i : i + valid_w] = np.maximum(
                            mask_array[j : j + valid_h, i : i + valid_w],
                            combined_mask[:valid_h, :valid_w],
                        )

            # Write the final mask to the output file
            dst.write(mask_array, 1)

    print(f"Object masks saved to {output_path}")
    return output_path

vectorize_masks(masks_path, output_path=None, confidence_threshold=0.5, min_object_area=100, max_object_area=None, n_workers=None, **kwargs)

Convert masks with confidence to vector polygons.

Parameters:

Name	Type	Description	Default
masks_path	str	Path to masks GeoTIFF with confidence band.	required
output_path	Optional[str]	Path for output GeoJSON.	None
confidence_threshold	float	Minimum confidence score (0.0-1.0). Default: 0.5	0.5
min_object_area	int	Minimum area in pixels to keep an object. Default: 100	100
max_object_area	Optional[int]	Maximum area in pixels to keep an object. Default: None	None
n_workers	Optional[int]	int, default=None The number of worker threads to use. "None" means single-threaded processing. "-1" means using all available CPU processors. Positive integer means using that specific number of threads.	None
**kwargs	Any	Additional parameters	{}

Returns:

Type	Description
GeoDataFrame	GeoDataFrame with car detections and confidence values

Source code in geoai/extract.py

def vectorize_masks(
    self,
    masks_path: str,
    output_path: Optional[str] = None,
    confidence_threshold: float = 0.5,
    min_object_area: int = 100,
    max_object_area: Optional[int] = None,
    n_workers: Optional[int] = None,
    **kwargs: Any,
) -> gpd.GeoDataFrame:
    """
    Convert masks with confidence to vector polygons.

    Args:
        masks_path: Path to masks GeoTIFF with confidence band.
        output_path: Path for output GeoJSON.
        confidence_threshold: Minimum confidence score (0.0-1.0). Default: 0.5
        min_object_area: Minimum area in pixels to keep an object. Default: 100
        max_object_area: Maximum area in pixels to keep an object. Default: None
        n_workers: int, default=None
            The number of worker threads to use.
            "None" means single-threaded processing.
            "-1"   means using all available CPU processors.
            Positive integer means using that specific number of threads.
        **kwargs: Additional parameters

    Returns:
        GeoDataFrame with car detections and confidence values
    """
    import cv2  # Lazy import to avoid QGIS opencv conflicts

    def _process_single_component(
        component_mask: np.ndarray,
        conf_data: np.ndarray,
        transform: Any,
        confidence_threshold: float,
        min_object_area: int,
        max_object_area: Optional[int],
    ) -> Optional[Dict[str, Any]]:
        # Get confidence value
        conf_region = conf_data[component_mask > 0]
        if len(conf_region) > 0:
            confidence = np.mean(conf_region) / 255.0
        else:
            confidence = 0.0

        # Skip if confidence is below threshold
        if confidence < confidence_threshold:
            return None

        # Find contours
        contours, _ = cv2.findContours(
            component_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )

        results = []

        for contour in contours:
            # Filter by size
            area = cv2.contourArea(contour)
            if area < min_object_area:
                continue

            if max_object_area is not None and area > max_object_area:
                continue

            # Get minimum area rectangle
            rect = cv2.minAreaRect(contour)
            box_points = cv2.boxPoints(rect)

            # Convert to geographic coordinates
            geo_points = []
            for x, y in box_points:
                gx, gy = transform * (x, y)
                geo_points.append((gx, gy))

            # Create polygon
            poly = Polygon(geo_points)
            results.append((poly, confidence, area))

        return results

    import concurrent.futures
    from functools import partial

    def process_component(
        args: Tuple[int, np.ndarray, np.ndarray, Any, float, int, Optional[int]],
    ) -> Optional[Dict[str, Any]]:
        """
        Helper function to process a single component
        """
        (
            label,
            labeled_mask,
            conf_data,
            transform,
            confidence_threshold,
            min_object_area,
            max_object_area,
        ) = args

        # Create mask for this component
        component_mask = (labeled_mask == label).astype(np.uint8)

        return _process_single_component(
            component_mask,
            conf_data,
            transform,
            confidence_threshold,
            min_object_area,
            max_object_area,
        )

    start_time = time.time()
    print(f"Processing masks from: {masks_path}")

    if n_workers == -1:
        n_workers = os.cpu_count()

    with rasterio.open(masks_path) as src:
        # Read mask and confidence bands
        mask_data = src.read(1)
        conf_data = src.read(2)
        transform = src.transform
        crs = src.crs

        # Convert to binary mask
        binary_mask = mask_data > 0

        # Find connected components
        labeled_mask, num_features = ndimage.label(binary_mask)
        print(f"Found {num_features} connected components")

        # Process each component
        polygons = []
        confidences = []
        pixels = []

        if n_workers is None or n_workers == 1:
            print(
                "Using single-threaded processing, you can speed up processing by setting n_workers > 1"
            )
            # Add progress bar
            for label in tqdm(
                range(1, num_features + 1), desc="Processing components"
            ):
                # Create mask for this component
                component_mask = (labeled_mask == label).astype(np.uint8)

                result = _process_single_component(
                    component_mask,
                    conf_data,
                    transform,
                    confidence_threshold,
                    min_object_area,
                    max_object_area,
                )

                if result:
                    for poly, confidence, area in result:
                        # Add to lists
                        polygons.append(poly)
                        confidences.append(confidence)
                        pixels.append(area)

        else:
            # Process components in parallel
            print(f"Using {n_workers} workers for parallel processing")

            process_args = [
                (
                    label,
                    labeled_mask,
                    conf_data,
                    transform,
                    confidence_threshold,
                    min_object_area,
                    max_object_area,
                )
                for label in range(1, num_features + 1)
            ]

            with concurrent.futures.ThreadPoolExecutor(
                max_workers=n_workers
            ) as executor:
                results = list(
                    tqdm(
                        executor.map(process_component, process_args),
                        total=num_features,
                        desc="Processing components",
                    )
                )

                for result in results:
                    if result:
                        for poly, confidence, area in result:
                            # Add to lists
                            polygons.append(poly)
                            confidences.append(confidence)
                            pixels.append(area)

        # Create GeoDataFrame
        if polygons:
            gdf = gpd.GeoDataFrame(
                {
                    "geometry": polygons,
                    "confidence": confidences,
                    "class": [1] * len(polygons),
                    "pixels": pixels,
                },
                crs=crs,
            )

            # Save to file if requested
            if output_path:
                gdf.to_file(output_path, driver="GeoJSON")
                print(f"Saved {len(gdf)} objects with confidence to {output_path}")

            end_time = time.time()
            print(f"Total processing time: {end_time - start_time:.2f} seconds")
            return gdf
        else:
            end_time = time.time()
            print(f"Total processing time: {end_time - start_time:.2f} seconds")
            print("No valid polygons found")
            return None

visualize_results(raster_path, gdf=None, output_path=None, figsize=(12, 12))

Visualize object detection results with proper coordinate transformation.

This function displays objects on top of the raster image, ensuring proper alignment between the GeoDataFrame polygons and the image.

Parameters:

Name	Type	Description	Default
raster_path	str	Path to input raster	required
gdf	Optional[GeoDataFrame]	GeoDataFrame with object polygons (optional)	None
output_path	Optional[str]	Path to save visualization (optional)	None
figsize	Tuple[int, int]	Figure size (width, height) in inches	(12, 12)

Returns:

Name	Type	Description
bool	bool	True if visualization was successful

Source code in geoai/extract.py

def visualize_results(
    self,
    raster_path: str,
    gdf: Optional[gpd.GeoDataFrame] = None,
    output_path: Optional[str] = None,
    figsize: Tuple[int, int] = (12, 12),
) -> bool:
    """
    Visualize object detection results with proper coordinate transformation.

    This function displays objects on top of the raster image,
    ensuring proper alignment between the GeoDataFrame polygons and the image.

    Args:
        raster_path: Path to input raster
        gdf: GeoDataFrame with object polygons (optional)
        output_path: Path to save visualization (optional)
        figsize: Figure size (width, height) in inches

    Returns:
        bool: True if visualization was successful
    """
    # Check if raster file exists
    if not os.path.exists(raster_path):
        print(f"Error: Raster file '{raster_path}' not found.")
        return False

    # Process raster if GeoDataFrame not provided
    if gdf is None:
        gdf = self.process_raster(raster_path)

    if gdf is None or len(gdf) == 0:
        print("No objects to visualize")
        return False

    # Check if confidence column exists in the GeoDataFrame
    has_confidence = False
    if hasattr(gdf, "columns") and "confidence" in gdf.columns:
        # Try to access a confidence value to confirm it works
        try:
            if len(gdf) > 0:
                # Try getitem access
                conf_val = gdf["confidence"].iloc[0]
                has_confidence = True
                print(
                    f"Using confidence values (range: {gdf['confidence'].min():.2f} - {gdf['confidence'].max():.2f})"
                )
        except Exception as e:
            print(f"Confidence column exists but couldn't access values: {e}")
            has_confidence = False
    else:
        print("No confidence column found in GeoDataFrame")
        has_confidence = False

    # Read raster for visualization
    with rasterio.open(raster_path) as src:
        # Read the entire image or a subset if it's very large
        if src.height > 2000 or src.width > 2000:
            # Calculate scale factor to reduce size
            scale = min(2000 / src.height, 2000 / src.width)
            out_shape = (
                int(src.count),
                int(src.height * scale),
                int(src.width * scale),
            )

            # Read and resample
            image = src.read(
                out_shape=out_shape, resampling=rasterio.enums.Resampling.bilinear
            )

            # Create a scaled transform for the resampled image
            # Calculate scaling factors
            x_scale = src.width / out_shape[2]
            y_scale = src.height / out_shape[1]

            # Get the original transform
            orig_transform = src.transform

            # Create a scaled transform
            scaled_transform = rasterio.transform.Affine(
                orig_transform.a * x_scale,
                orig_transform.b,
                orig_transform.c,
                orig_transform.d,
                orig_transform.e * y_scale,
                orig_transform.f,
            )
        else:
            image = src.read()
            scaled_transform = src.transform

        # Convert to RGB for display
        if image.shape[0] > 3:
            image = image[:3]
        elif image.shape[0] == 1:
            image = np.repeat(image, 3, axis=0)

        # Normalize image for display
        image = image.transpose(1, 2, 0)  # CHW to HWC
        image = image.astype(np.float32)

        if image.max() > 10:  # Likely 0-255 range
            image = image / 255.0

        image = np.clip(image, 0, 1)

        # Get image bounds
        bounds = src.bounds
        crs = src.crs

    # Create figure with appropriate aspect ratio
    aspect_ratio = image.shape[1] / image.shape[0]  # width / height
    plt.figure(figsize=(figsize[0], figsize[0] / aspect_ratio))
    ax = plt.gca()

    # Display image
    ax.imshow(image)

    # Make sure the GeoDataFrame has the same CRS as the raster
    if gdf.crs != crs:
        print(f"Reprojecting GeoDataFrame from {gdf.crs} to {crs}")
        gdf = gdf.to_crs(crs)

    # Set up colors for confidence visualization
    if has_confidence:
        try:
            import matplotlib.cm as cm
            from matplotlib.colors import Normalize

            # Get min/max confidence values
            min_conf = gdf["confidence"].min()
            max_conf = gdf["confidence"].max()

            # Set up normalization and colormap
            norm = Normalize(vmin=min_conf, vmax=max_conf)
            cmap = cm.viridis

            # Create scalar mappable for colorbar
            sm = cm.ScalarMappable(cmap=cmap, norm=norm)
            sm.set_array([])

            # Add colorbar
            cbar = plt.colorbar(
                sm, ax=ax, orientation="vertical", shrink=0.7, pad=0.01
            )
            cbar.set_label("Confidence Score")
        except Exception as e:
            print(f"Error setting up confidence visualization: {e}")
            has_confidence = False

    # Function to convert coordinates
    def geo_to_pixel(
        geometry: Any, transform: Any
    ) -> Optional[Tuple[List[float], List[float]]]:
        """Convert geometry to pixel coordinates using the provided transform."""
        if geometry.is_empty:
            return None

        if geometry.geom_type == "Polygon":
            # Get exterior coordinates
            exterior_coords = list(geometry.exterior.coords)

            # Convert to pixel coordinates
            pixel_coords = [~transform * (x, y) for x, y in exterior_coords]

            # Split into x and y lists
            pixel_x = [coord[0] for coord in pixel_coords]
            pixel_y = [coord[1] for coord in pixel_coords]

            return pixel_x, pixel_y
        else:
            print(f"Unsupported geometry type: {geometry.geom_type}")
            return None

    # Plot each object
    for idx, row in gdf.iterrows():
        try:
            # Convert polygon to pixel coordinates
            coords = geo_to_pixel(row.geometry, scaled_transform)

            if coords:
                pixel_x, pixel_y = coords

                if has_confidence:
                    try:
                        # Get confidence value using different methods
                        # Method 1: Try direct attribute access
                        confidence = None
                        try:
                            confidence = row.confidence
                        except Exception:
                            pass

                        # Method 2: Try dictionary-style access
                        if confidence is None:
                            try:
                                confidence = row["confidence"]
                            except Exception:
                                pass

                        # Method 3: Try accessing by index from the GeoDataFrame
                        if confidence is None:
                            try:
                                confidence = gdf.iloc[idx]["confidence"]
                            except Exception:
                                pass

                        if confidence is not None:
                            color = cmap(norm(confidence))
                            # Fill polygon with semi-transparent color
                            ax.fill(pixel_x, pixel_y, color=color, alpha=0.5)
                            # Draw border
                            ax.plot(
                                pixel_x,
                                pixel_y,
                                color=color,
                                linewidth=1,
                                alpha=0.8,
                            )
                        else:
                            # Fall back to red if confidence value couldn't be accessed
                            ax.plot(pixel_x, pixel_y, color="red", linewidth=1)
                    except Exception as e:
                        print(
                            f"Error using confidence value for polygon {idx}: {e}"
                        )
                        ax.plot(pixel_x, pixel_y, color="red", linewidth=1)
                else:
                    # No confidence data, just plot outlines in red
                    ax.plot(pixel_x, pixel_y, color="red", linewidth=1)
        except Exception as e:
            print(f"Error plotting polygon {idx}: {e}")

    # Remove axes
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_title(f"objects (Found: {len(gdf)})")

    # Save if requested
    if output_path:
        plt.tight_layout()
        plt.savefig(output_path, dpi=300, bbox_inches="tight")
        print(f"Visualization saved to {output_path}")

    plt.close()

    # Create a simpler visualization focused just on a subset of objects
    if len(gdf) > 0:
        plt.figure(figsize=figsize)
        ax = plt.gca()

        # Choose a subset of the image to show
        with rasterio.open(raster_path) as src:
            # Get centroid of first object
            sample_geom = gdf.iloc[0].geometry
            centroid = sample_geom.centroid

            # Convert to pixel coordinates
            center_x, center_y = ~src.transform * (centroid.x, centroid.y)

            # Define a window around this object
            window_size = 500  # pixels
            window = rasterio.windows.Window(
                max(0, int(center_x - window_size / 2)),
                max(0, int(center_y - window_size / 2)),
                min(window_size, src.width - int(center_x - window_size / 2)),
                min(window_size, src.height - int(center_y - window_size / 2)),
            )

            # Read this window
            sample_image = src.read(window=window)

            # Convert to RGB for display
            if sample_image.shape[0] > 3:
                sample_image = sample_image[:3]
            elif sample_image.shape[0] == 1:
                sample_image = np.repeat(sample_image, 3, axis=0)

            # Normalize image for display
            sample_image = sample_image.transpose(1, 2, 0)  # CHW to HWC
            sample_image = sample_image.astype(np.float32)

            if sample_image.max() > 10:  # Likely 0-255 range
                sample_image = sample_image / 255.0

            sample_image = np.clip(sample_image, 0, 1)

            # Display sample image
            ax.imshow(sample_image, extent=[0, window.width, window.height, 0])

            # Get the correct transform for this window
            window_transform = src.window_transform(window)

            # Calculate bounds of the window
            window_bounds = rasterio.windows.bounds(window, src.transform)
            window_box = box(*window_bounds)

            # Filter objects that intersect with this window
            visible_gdf = gdf[gdf.intersects(window_box)]

            # Set up colors for sample view if confidence data exists
            if has_confidence:
                try:
                    # Reuse the same normalization and colormap from main view
                    sample_sm = cm.ScalarMappable(cmap=cmap, norm=norm)
                    sample_sm.set_array([])

                    # Add colorbar to sample view
                    sample_cbar = plt.colorbar(
                        sample_sm,
                        ax=ax,
                        orientation="vertical",
                        shrink=0.7,
                        pad=0.01,
                    )
                    sample_cbar.set_label("Confidence Score")
                except Exception as e:
                    print(f"Error setting up sample confidence visualization: {e}")

            # Plot objects in sample view
            for idx, row in visible_gdf.iterrows():
                try:
                    # Get window-relative pixel coordinates
                    geom = row.geometry

                    # Skip empty geometries
                    if geom.is_empty:
                        continue

                    # Get exterior coordinates
                    exterior_coords = list(geom.exterior.coords)

                    # Convert to pixel coordinates relative to window origin
                    pixel_coords = []
                    for x, y in exterior_coords:
                        px, py = ~src.transform * (x, y)  # Convert to image pixels
                        # Make coordinates relative to window
                        px = px - window.col_off
                        py = py - window.row_off
                        pixel_coords.append((px, py))

                    # Extract x and y coordinates
                    pixel_x = [coord[0] for coord in pixel_coords]
                    pixel_y = [coord[1] for coord in pixel_coords]

                    # Use confidence colors if available
                    if has_confidence:
                        try:
                            # Try different methods to access confidence
                            confidence = None
                            try:
                                confidence = row.confidence
                            except Exception:
                                pass

                            if confidence is None:
                                try:
                                    confidence = row["confidence"]
                                except Exception:
                                    pass

                            if confidence is None:
                                try:
                                    confidence = visible_gdf.iloc[idx]["confidence"]
                                except Exception:
                                    pass

                            if confidence is not None:
                                color = cmap(norm(confidence))
                                # Fill polygon with semi-transparent color
                                ax.fill(pixel_x, pixel_y, color=color, alpha=0.5)
                                # Draw border
                                ax.plot(
                                    pixel_x,
                                    pixel_y,
                                    color=color,
                                    linewidth=1.5,
                                    alpha=0.8,
                                )
                            else:
                                ax.plot(
                                    pixel_x, pixel_y, color="red", linewidth=1.5
                                )
                        except Exception as e:
                            print(
                                f"Error using confidence in sample view for polygon {idx}: {e}"
                            )
                            ax.plot(pixel_x, pixel_y, color="red", linewidth=1.5)
                    else:
                        ax.plot(pixel_x, pixel_y, color="red", linewidth=1.5)
                except Exception as e:
                    print(f"Error plotting polygon in sample view: {e}")

            # Set title
            ax.set_title(f"Sample Area - objects (Showing: {len(visible_gdf)})")

            # Remove axes
            ax.set_xticks([])
            ax.set_yticks([])

            # Save if requested
            if output_path:
                sample_output = (
                    os.path.splitext(output_path)[0]
                    + "_sample"
                    + os.path.splitext(output_path)[1]
                )
                plt.tight_layout()
                plt.savefig(sample_output, dpi=300, bbox_inches="tight")
                print(f"Sample visualization saved to {sample_output}")

ParkingSplotDetector

Bases: ObjectDetector

Car detection using a pre-trained Mask R-CNN model.

This class extends the ObjectDetector class with additional methods for car detection.

Source code in geoai/extract.py

class ParkingSplotDetector(ObjectDetector):
    """
    Car detection using a pre-trained Mask R-CNN model.

    This class extends the `ObjectDetector` class with additional methods for car detection.
    """

    def __init__(
        self,
        model_path: str = "parking_spot_detection.pth",
        repo_id: Optional[str] = None,
        model: Optional[Any] = None,
        num_classes: int = 3,
        device: Optional[str] = None,
    ) -> None:
        """
        Initialize the object extractor.

        Args:
            model_path: Path to the .pth model file.
            repo_id: Repo ID for loading models from the Hub.
            model: Custom model to use for inference.
            num_classes: Number of classes for the model. Default: 3
            device: Device to use for inference ('cuda:0', 'cpu', etc.).
        """
        super().__init__(
            model_path=model_path,
            repo_id=repo_id,
            model=model,
            num_classes=num_classes,
            device=device,
        )

__init__(model_path='parking_spot_detection.pth', repo_id=None, model=None, num_classes=3, device=None)

Initialize the object extractor.

Parameters:

Name	Type	Description	Default
model_path	str	Path to the .pth model file.	'parking_spot_detection.pth'
repo_id	Optional[str]	Repo ID for loading models from the Hub.	None
model	Optional[Any]	Custom model to use for inference.	None
num_classes	int	Number of classes for the model. Default: 3	3
device	Optional[str]	Device to use for inference ('cuda:0', 'cpu', etc.).	None

Source code in geoai/extract.py

def __init__(
    self,
    model_path: str = "parking_spot_detection.pth",
    repo_id: Optional[str] = None,
    model: Optional[Any] = None,
    num_classes: int = 3,
    device: Optional[str] = None,
) -> None:
    """
    Initialize the object extractor.

    Args:
        model_path: Path to the .pth model file.
        repo_id: Repo ID for loading models from the Hub.
        model: Custom model to use for inference.
        num_classes: Number of classes for the model. Default: 3
        device: Device to use for inference ('cuda:0', 'cpu', etc.).
    """
    super().__init__(
        model_path=model_path,
        repo_id=repo_id,
        model=model,
        num_classes=num_classes,
        device=device,
    )

ShipDetector

Bases: ObjectDetector

Ship detection using a pre-trained Mask R-CNN model.

This class extends the ObjectDetector class with additional methods for ship detection."

Source code in geoai/extract.py

class ShipDetector(ObjectDetector):
    """
    Ship detection using a pre-trained Mask R-CNN model.

    This class extends the
    `ObjectDetector` class with additional methods for ship detection."
    """

    def __init__(
        self,
        model_path: str = "ship_detection.pth",
        repo_id: Optional[str] = None,
        model: Optional[Any] = None,
        device: Optional[str] = None,
    ) -> None:
        """
        Initialize the object extractor.

        Args:
            model_path: Path to the .pth model file.
            repo_id: Repo ID for loading models from the Hub.
            model: Custom model to use for inference.
            device: Device to use for inference ('cuda:0', 'cpu', etc.).
        """
        super().__init__(
            model_path=model_path, repo_id=repo_id, model=model, device=device
        )

__init__(model_path='ship_detection.pth', repo_id=None, model=None, device=None)

Initialize the object extractor.

Parameters:

Name	Type	Description	Default
model_path	str	Path to the .pth model file.	'ship_detection.pth'
repo_id	Optional[str]	Repo ID for loading models from the Hub.	None
model	Optional[Any]	Custom model to use for inference.	None
device	Optional[str]	Device to use for inference ('cuda:0', 'cpu', etc.).	None

Source code in geoai/extract.py

def __init__(
    self,
    model_path: str = "ship_detection.pth",
    repo_id: Optional[str] = None,
    model: Optional[Any] = None,
    device: Optional[str] = None,
) -> None:
    """
    Initialize the object extractor.

    Args:
        model_path: Path to the .pth model file.
        repo_id: Repo ID for loading models from the Hub.
        model: Custom model to use for inference.
        device: Device to use for inference ('cuda:0', 'cpu', etc.).
    """
    super().__init__(
        model_path=model_path, repo_id=repo_id, model=model, device=device
    )

SolarPanelDetector

Bases: ObjectDetector

Solar panel detection using a pre-trained Mask R-CNN model.

This class extends the ObjectDetector class with additional methods for solar panel detection."

Source code in geoai/extract.py

class SolarPanelDetector(ObjectDetector):
    """
    Solar panel detection using a pre-trained Mask R-CNN model.

    This class extends the
    `ObjectDetector` class with additional methods for solar panel detection."
    """

    def __init__(
        self,
        model_path: str = "solar_panel_detection.pth",
        repo_id: Optional[str] = None,
        model: Optional[Any] = None,
        device: Optional[str] = None,
    ) -> None:
        """
        Initialize the object extractor.

        Args:
            model_path: Path to the .pth model file.
            repo_id: Repo ID for loading models from the Hub.
            model: Custom model to use for inference.
            device: Device to use for inference ('cuda:0', 'cpu', etc.).
        """
        super().__init__(
            model_path=model_path, repo_id=repo_id, model=model, device=device
        )

__init__(model_path='solar_panel_detection.pth', repo_id=None, model=None, device=None)

Initialize the object extractor.

Parameters:

Name	Type	Description	Default
model_path	str	Path to the .pth model file.	'solar_panel_detection.pth'
repo_id	Optional[str]	Repo ID for loading models from the Hub.	None
model	Optional[Any]	Custom model to use for inference.	None
device	Optional[str]	Device to use for inference ('cuda:0', 'cpu', etc.).	None

Source code in geoai/extract.py

def __init__(
    self,
    model_path: str = "solar_panel_detection.pth",
    repo_id: Optional[str] = None,
    model: Optional[Any] = None,
    device: Optional[str] = None,
) -> None:
    """
    Initialize the object extractor.

    Args:
        model_path: Path to the .pth model file.
        repo_id: Repo ID for loading models from the Hub.
        model: Custom model to use for inference.
        device: Device to use for inference ('cuda:0', 'cpu', etc.).
    """
    super().__init__(
        model_path=model_path, repo_id=repo_id, model=model, device=device
    )

adaptive_regularization(building_polygons, simplify_tolerance=0.5, area_threshold=0.9, preserve_shape=True)

Adaptively regularizes building footprints based on their characteristics.

This approach determines the best regularization method for each building.

Parameters:

Name	Type	Description	Default
building_polygons	Union[GeoDataFrame, List[Polygon]]	GeoDataFrame or list of shapely Polygons	required
simplify_tolerance	float	Distance tolerance for simplification	0.5
area_threshold	float	Minimum acceptable area ratio	0.9
preserve_shape	bool	Whether to preserve overall shape for complex buildings	True

Returns:

Type	Description
Union[GeoDataFrame, List[Polygon]]	GeoDataFrame or list of shapely Polygons with regularized building footprints

Source code in geoai/utils/geometry.py

def adaptive_regularization(
    building_polygons: Union[gpd.GeoDataFrame, List[Polygon]],
    simplify_tolerance: float = 0.5,
    area_threshold: float = 0.9,
    preserve_shape: bool = True,
) -> Union[gpd.GeoDataFrame, List[Polygon]]:
    """
    Adaptively regularizes building footprints based on their characteristics.

    This approach determines the best regularization method for each building.

    Args:
        building_polygons: GeoDataFrame or list of shapely Polygons
        simplify_tolerance: Distance tolerance for simplification
        area_threshold: Minimum acceptable area ratio
        preserve_shape: Whether to preserve overall shape for complex buildings

    Returns:
        GeoDataFrame or list of shapely Polygons with regularized building footprints
    """
    from shapely.affinity import rotate
    from shapely.geometry import Polygon

    # Analyze the overall dataset to set appropriate parameters
    if is_gdf := isinstance(building_polygons, gpd.GeoDataFrame):
        geom_objects = building_polygons.geometry
    else:
        geom_objects = building_polygons

    results = []

    for building in geom_objects:
        # Skip invalid geometries
        if not hasattr(building, "exterior") or building.is_empty:
            results.append(building)
            continue

        # Measure building complexity
        complexity = building.length / (4 * np.sqrt(building.area))

        # Determine if the building has a clear principal direction
        coords = np.array(building.exterior.coords)[:-1]
        segments = np.diff(np.vstack([coords, coords[0]]), axis=0)
        segment_lengths = np.sqrt(segments[:, 0] ** 2 + segments[:, 1] ** 2)
        angles = np.arctan2(segments[:, 1], segments[:, 0]) * 180 / np.pi

        # Normalize angles to 0-180 range and get histogram
        norm_angles = angles % 180
        hist, bins = np.histogram(
            norm_angles, bins=18, range=(0, 180), weights=segment_lengths
        )

        # Calculate direction clarity (ratio of longest direction to total)
        direction_clarity = np.max(hist) / np.sum(hist) if np.sum(hist) > 0 else 0

        # Choose regularization method based on building characteristics
        if complexity < 1.2 and direction_clarity > 0.5:
            # Simple building with clear direction: use rotated rectangle
            bin_max = np.argmax(hist)
            bin_centers = (bins[:-1] + bins[1:]) / 2
            dominant_angle = bin_centers[bin_max]

            # Rotate to align with coordinate system
            rotated = rotate(building, -dominant_angle, origin="centroid")

            # Create bounding box in rotated space
            bounds = rotated.bounds
            rect = Polygon(
                [
                    (bounds[0], bounds[1]),
                    (bounds[2], bounds[1]),
                    (bounds[2], bounds[3]),
                    (bounds[0], bounds[3]),
                ]
            )

            # Rotate back
            result = rotate(rect, dominant_angle, origin="centroid")

            # Quality check
            if (
                result.area / building.area < area_threshold
                or result.area / building.area > (1.0 / area_threshold)
            ):
                # Too much area change, use simplified original
                result = building.simplify(simplify_tolerance, preserve_topology=True)

        else:
            # Complex building or no clear direction: preserve shape
            if preserve_shape:
                # Simplify with topology preservation
                result = building.simplify(simplify_tolerance, preserve_topology=True)
            else:
                # Fall back to convex hull for very complex shapes
                result = building.convex_hull

        results.append(result)

    # Return in same format as input
    if is_gdf:
        return gpd.GeoDataFrame(geometry=results, crs=building_polygons.crs)
    else:
        return results

add_geometric_properties(data, properties=None, area_unit='m2', length_unit='m')

Calculates geometric properties and adds them to the GeoDataFrame.

This function calculates various geometric properties of features in a GeoDataFrame and adds them as new columns without modifying existing attributes.

Parameters:

Name	Type	Description	Default
data	GeoDataFrame	GeoDataFrame containing vector features.	required
properties	Optional[List[str]]	List of geometric properties to calculate. Options include: 'area', 'length', 'perimeter', 'centroid_x', 'centroid_y', 'bounds', 'convex_hull_area', 'orientation', 'complexity', 'area_bbox', 'area_convex', 'area_filled', 'major_length', 'minor_length', 'eccentricity', 'diameter_area', 'extent', 'solidity', 'elongation'. Defaults to ['area', 'length'] if None.	None
area_unit	str	String specifying the unit for area calculation ('m2', 'km2', 'ha'). Defaults to 'm2'.	'm2'
length_unit	str	String specifying the unit for length calculation ('m', 'km'). Defaults to 'm'.	'm'

Returns:

Type	Description
GeoDataFrame	geopandas.GeoDataFrame: A copy of the input GeoDataFrame with added
GeoDataFrame	geometric property columns.

Source code in geoai/utils/vector.py

def add_geometric_properties(
    data: gpd.GeoDataFrame,
    properties: Optional[List[str]] = None,
    area_unit: str = "m2",
    length_unit: str = "m",
) -> gpd.GeoDataFrame:
    """Calculates geometric properties and adds them to the GeoDataFrame.

    This function calculates various geometric properties of features in a
    GeoDataFrame and adds them as new columns without modifying existing attributes.

    Args:
        data: GeoDataFrame containing vector features.
        properties: List of geometric properties to calculate. Options include:
            'area', 'length', 'perimeter', 'centroid_x', 'centroid_y', 'bounds',
            'convex_hull_area', 'orientation', 'complexity', 'area_bbox',
            'area_convex', 'area_filled', 'major_length', 'minor_length',
            'eccentricity', 'diameter_area', 'extent', 'solidity',
            'elongation'.
            Defaults to ['area', 'length'] if None.
        area_unit: String specifying the unit for area calculation ('m2', 'km2',
            'ha'). Defaults to 'm2'.
        length_unit: String specifying the unit for length calculation ('m', 'km').
            Defaults to 'm'.

    Returns:
        geopandas.GeoDataFrame: A copy of the input GeoDataFrame with added
        geometric property columns.
    """
    from shapely.ops import unary_union

    if isinstance(data, str):
        from .raster import read_vector

        data = read_vector(data)

    # Make a copy to avoid modifying the original
    result = data.copy()

    # Default properties to calculate
    if properties is None:
        properties = [
            "area",
            "length",
            "perimeter",
            "convex_hull_area",
            "orientation",
            "complexity",
            "area_bbox",
            "area_convex",
            "area_filled",
            "major_length",
            "minor_length",
            "eccentricity",
            "diameter_area",
            "extent",
            "solidity",
            "elongation",
        ]

    # Make sure we're working with a GeoDataFrame with a valid CRS

    if not isinstance(result, gpd.GeoDataFrame):
        raise ValueError("Input must be a GeoDataFrame")

    if result.crs is None:
        raise ValueError(
            "GeoDataFrame must have a defined coordinate reference system (CRS)"
        )

    # Ensure we're working with a projected CRS for accurate measurements
    if result.crs.is_geographic:
        # Reproject to a suitable projected CRS for accurate measurements
        result = result.to_crs(result.estimate_utm_crs())

    # Basic area calculation with unit conversion
    if "area" in properties:
        # Calculate area (only for polygons)
        result["area"] = result.geometry.apply(
            lambda geom: geom.area if isinstance(geom, (Polygon, MultiPolygon)) else 0
        )

        # Convert to requested units
        if area_unit == "km2":
            result["area"] = result["area"] / 1_000_000  # m² to km²
            result.rename(columns={"area": "area_km2"}, inplace=True)
        elif area_unit == "ha":
            result["area"] = result["area"] / 10_000  # m² to hectares
            result.rename(columns={"area": "area_ha"}, inplace=True)
        else:  # Default is m²
            result.rename(columns={"area": "area_m2"}, inplace=True)

    # Length calculation with unit conversion
    if "length" in properties:
        # Calculate length (works for lines and polygon boundaries)
        result["length"] = result.geometry.length

        # Convert to requested units
        if length_unit == "km":
            result["length"] = result["length"] / 1_000  # m to km
            result.rename(columns={"length": "length_km"}, inplace=True)
        else:  # Default is m
            result.rename(columns={"length": "length_m"}, inplace=True)

    # Perimeter calculation (for polygons)
    if "perimeter" in properties:
        result["perimeter"] = result.geometry.apply(
            lambda geom: (
                geom.boundary.length if isinstance(geom, (Polygon, MultiPolygon)) else 0
            )
        )

        # Convert to requested units
        if length_unit == "km":
            result["perimeter"] = result["perimeter"] / 1_000  # m to km
            result.rename(columns={"perimeter": "perimeter_km"}, inplace=True)
        else:  # Default is m
            result.rename(columns={"perimeter": "perimeter_m"}, inplace=True)

    # Centroid coordinates
    if "centroid_x" in properties or "centroid_y" in properties:
        centroids = result.geometry.centroid

        if "centroid_x" in properties:
            result["centroid_x"] = centroids.x

        if "centroid_y" in properties:
            result["centroid_y"] = centroids.y

    # Bounding box properties
    if "bounds" in properties:
        bounds = result.geometry.bounds
        result["minx"] = bounds.minx
        result["miny"] = bounds.miny
        result["maxx"] = bounds.maxx
        result["maxy"] = bounds.maxy

    # Area of bounding box
    if "area_bbox" in properties:
        bounds = result.geometry.bounds
        result["area_bbox"] = (bounds.maxx - bounds.minx) * (bounds.maxy - bounds.miny)

        # Convert to requested units
        if area_unit == "km2":
            result["area_bbox"] = result["area_bbox"] / 1_000_000
            result.rename(columns={"area_bbox": "area_bbox_km2"}, inplace=True)
        elif area_unit == "ha":
            result["area_bbox"] = result["area_bbox"] / 10_000
            result.rename(columns={"area_bbox": "area_bbox_ha"}, inplace=True)
        else:  # Default is m²
            result.rename(columns={"area_bbox": "area_bbox_m2"}, inplace=True)

    # Area of convex hull
    if "area_convex" in properties or "convex_hull_area" in properties:
        result["area_convex"] = result.geometry.convex_hull.area

        # Convert to requested units
        if area_unit == "km2":
            result["area_convex"] = result["area_convex"] / 1_000_000
            result.rename(columns={"area_convex": "area_convex_km2"}, inplace=True)
        elif area_unit == "ha":
            result["area_convex"] = result["area_convex"] / 10_000
            result.rename(columns={"area_convex": "area_convex_ha"}, inplace=True)
        else:  # Default is m²
            result.rename(columns={"area_convex": "area_convex_m2"}, inplace=True)

        # For backward compatibility
        if "convex_hull_area" in properties and "area_convex" not in properties:
            result["convex_hull_area"] = result["area_convex"]
            if area_unit == "km2":
                result.rename(
                    columns={"convex_hull_area": "convex_hull_area_km2"}, inplace=True
                )
            elif area_unit == "ha":
                result.rename(
                    columns={"convex_hull_area": "convex_hull_area_ha"}, inplace=True
                )
            else:
                result.rename(
                    columns={"convex_hull_area": "convex_hull_area_m2"}, inplace=True
                )

    # Area of filled geometry (no holes)
    if "area_filled" in properties:

        def get_filled_area(geom):
            if not isinstance(geom, (Polygon, MultiPolygon)):
                return 0

            if isinstance(geom, MultiPolygon):
                # For MultiPolygon, fill all constituent polygons
                filled_polys = [Polygon(p.exterior) for p in geom.geoms]
                return unary_union(filled_polys).area
            else:
                # For single Polygon, create a new one with just the exterior ring
                return Polygon(geom.exterior).area

        result["area_filled"] = result.geometry.apply(get_filled_area)

        # Convert to requested units
        if area_unit == "km2":
            result["area_filled"] = result["area_filled"] / 1_000_000
            result.rename(columns={"area_filled": "area_filled_km2"}, inplace=True)
        elif area_unit == "ha":
            result["area_filled"] = result["area_filled"] / 10_000
            result.rename(columns={"area_filled": "area_filled_ha"}, inplace=True)
        else:  # Default is m²
            result.rename(columns={"area_filled": "area_filled_m2"}, inplace=True)

    # Axes lengths, eccentricity, orientation, and elongation
    if any(
        p in properties
        for p in [
            "major_length",
            "minor_length",
            "eccentricity",
            "orientation",
            "elongation",
        ]
    ):

        def get_axes_properties(geom):
            # Skip non-polygons
            if not isinstance(geom, (Polygon, MultiPolygon)):
                return None, None, None, None, None

            # Handle multipolygons by using the largest polygon
            if isinstance(geom, MultiPolygon):
                # Get the polygon with the largest area
                geom = sorted(list(geom.geoms), key=lambda p: p.area, reverse=True)[0]

            try:
                # Get the minimum rotated rectangle
                rect = geom.minimum_rotated_rectangle

                # Extract coordinates
                coords = list(rect.exterior.coords)[
                    :-1
                ]  # Remove the duplicated last point

                if len(coords) < 4:
                    return None, None, None, None, None

                # Calculate lengths of all four sides
                sides = []
                for i in range(len(coords)):
                    p1 = coords[i]
                    p2 = coords[(i + 1) % len(coords)]
                    dx = p2[0] - p1[0]
                    dy = p2[1] - p1[1]
                    length = np.sqrt(dx**2 + dy**2)
                    angle = np.degrees(np.arctan2(dy, dx)) % 180
                    sides.append((length, angle, p1, p2))

                # Group sides by length (allowing for small differences due to floating point precision)
                # This ensures we correctly identify the rectangle's dimensions
                sides_grouped = {}
                tolerance = 1e-6  # Tolerance for length comparison

                for s in sides:
                    length, angle = s[0], s[1]
                    matched = False

                    for key in sides_grouped:
                        if abs(length - key) < tolerance:
                            sides_grouped[key].append(s)
                            matched = True
                            break

                    if not matched:
                        sides_grouped[length] = [s]

                # Get unique lengths (should be 2 for a rectangle, parallel sides have equal length)
                unique_lengths = sorted(sides_grouped.keys(), reverse=True)

                if len(unique_lengths) != 2:
                    # If we don't get exactly 2 unique lengths, something is wrong with the rectangle
                    # Fall back to simpler method using bounds
                    bounds = rect.bounds
                    width = bounds[2] - bounds[0]
                    height = bounds[3] - bounds[1]
                    major_length = max(width, height)
                    minor_length = min(width, height)
                    orientation = 0 if width > height else 90
                else:
                    major_length = unique_lengths[0]
                    minor_length = unique_lengths[1]
                    # Get orientation from the major axis
                    orientation = sides_grouped[major_length][0][1]

                # Calculate eccentricity
                if major_length > 0:
                    # Eccentricity for an ellipse: e = sqrt(1 - (b²/a²))
                    # where a is the semi-major axis and b is the semi-minor axis
                    eccentricity = np.sqrt(
                        1 - ((minor_length / 2) ** 2 / (major_length / 2) ** 2)
                    )
                else:
                    eccentricity = 0

                # Calculate elongation (ratio of minor to major axis)
                elongation = major_length / minor_length if major_length > 0 else 1

                return major_length, minor_length, eccentricity, orientation, elongation

            except Exception as e:
                # For debugging
                # print(f"Error calculating axes: {e}")
                return None, None, None, None, None

        # Apply the function and split the results
        axes_data = result.geometry.apply(get_axes_properties)

        if "major_length" in properties:
            result["major_length"] = axes_data.apply(lambda x: x[0] if x else None)
            # Convert to requested units
            if length_unit == "km":
                result["major_length"] = result["major_length"] / 1_000
                result.rename(columns={"major_length": "major_length_km"}, inplace=True)
            else:
                result.rename(columns={"major_length": "major_length_m"}, inplace=True)

        if "minor_length" in properties:
            result["minor_length"] = axes_data.apply(lambda x: x[1] if x else None)
            # Convert to requested units
            if length_unit == "km":
                result["minor_length"] = result["minor_length"] / 1_000
                result.rename(columns={"minor_length": "minor_length_km"}, inplace=True)
            else:
                result.rename(columns={"minor_length": "minor_length_m"}, inplace=True)

        if "eccentricity" in properties:
            result["eccentricity"] = axes_data.apply(lambda x: x[2] if x else None)

        if "orientation" in properties:
            result["orientation"] = axes_data.apply(lambda x: x[3] if x else None)

        if "elongation" in properties:
            result["elongation"] = axes_data.apply(lambda x: x[4] if x else None)

    # Equivalent diameter based on area
    if "diameter_area" in properties:

        def get_equivalent_diameter(geom):
            if not isinstance(geom, (Polygon, MultiPolygon)) or geom.area <= 0:
                return None
            # Diameter of a circle with the same area: d = 2 * sqrt(A / π)
            return 2 * np.sqrt(geom.area / np.pi)

        result["diameter_area"] = result.geometry.apply(get_equivalent_diameter)

        # Convert to requested units
        if length_unit == "km":
            result["diameter_area"] = result["diameter_area"] / 1_000
            result.rename(
                columns={"diameter_area": "equivalent_diameter_area_km"},
                inplace=True,
            )
        else:
            result.rename(
                columns={"diameter_area": "equivalent_diameter_area_m"},
                inplace=True,
            )

    # Extent (ratio of shape area to bounding box area)
    if "extent" in properties:

        def get_extent(geom):
            if not isinstance(geom, (Polygon, MultiPolygon)) or geom.area <= 0:
                return None

            bounds = geom.bounds
            bbox_area = (bounds[2] - bounds[0]) * (bounds[3] - bounds[1])

            if bbox_area > 0:
                return geom.area / bbox_area
            return None

        result["extent"] = result.geometry.apply(get_extent)

    # Solidity (ratio of shape area to convex hull area)
    if "solidity" in properties:

        def get_solidity(geom):
            if not isinstance(geom, (Polygon, MultiPolygon)) or geom.area <= 0:
                return None

            convex_hull_area = geom.convex_hull.area

            if convex_hull_area > 0:
                return geom.area / convex_hull_area
            return None

        result["solidity"] = result.geometry.apply(get_solidity)

    # Complexity (ratio of perimeter to area)
    if "complexity" in properties:

        def calc_complexity(geom):
            if isinstance(geom, (Polygon, MultiPolygon)) and geom.area > 0:
                # Shape index: P / (2 * sqrt(π * A))
                # Normalized to 1 for a circle, higher for more complex shapes
                return geom.boundary.length / (2 * np.sqrt(np.pi * geom.area))
            return None

        result["complexity"] = result.geometry.apply(calc_complexity)

    return result

analyze_vector_attributes(vector_path, attribute_name)

Analyze a specific attribute in a vector dataset and create a histogram.

Parameters:

Name	Type	Description	Default
vector_path	str	Path to the vector file	required
attribute_name	str	Name of the attribute to analyze	required

Returns:

Name	Type	Description
dict	Optional[Dict[str, Any]]	Dictionary containing analysis results for the attribute

Source code in geoai/utils/vector.py

def analyze_vector_attributes(
    vector_path: str, attribute_name: str
) -> Optional[Dict[str, Any]]:
    """Analyze a specific attribute in a vector dataset and create a histogram.

    Args:
        vector_path (str): Path to the vector file
        attribute_name (str): Name of the attribute to analyze

    Returns:
        dict: Dictionary containing analysis results for the attribute
    """
    try:
        gdf = gpd.read_file(vector_path)

        # Check if attribute exists
        if attribute_name not in gdf.columns:
            print(f"Attribute '{attribute_name}' not found in the dataset")
            return None

        # Get the attribute series
        attr = gdf[attribute_name]

        # Perform different analyses based on data type
        if pd.api.types.is_numeric_dtype(attr):
            # Numeric attribute
            analysis = {
                "attribute": attribute_name,
                "type": "numeric",
                "count": attr.count(),
                "null_count": attr.isna().sum(),
                "min": attr.min(),
                "max": attr.max(),
                "mean": attr.mean(),
                "median": attr.median(),
                "std": attr.std(),
                "unique_values": attr.nunique(),
            }

            # Create histogram
            plt.figure(figsize=(10, 6))
            plt.hist(attr.dropna(), bins=20, alpha=0.7, color="blue")
            plt.title(f"Histogram of {attribute_name}")
            plt.xlabel(attribute_name)
            plt.ylabel("Frequency")
            plt.grid(True, alpha=0.3)
            plt.show()

        else:
            # Categorical attribute
            analysis = {
                "attribute": attribute_name,
                "type": "categorical",
                "count": attr.count(),
                "null_count": attr.isna().sum(),
                "unique_values": attr.nunique(),
                "value_counts": attr.value_counts().to_dict(),
            }

            # Create bar plot for top categories
            top_n = min(10, attr.nunique())
            plt.figure(figsize=(10, 6))
            attr.value_counts().head(top_n).plot(kind="bar", color="skyblue")
            plt.title(f"Top {top_n} values for {attribute_name}")
            plt.xlabel(attribute_name)
            plt.ylabel("Count")
            plt.xticks(rotation=45)
            plt.grid(True, alpha=0.3)
            plt.tight_layout()
            plt.show()

        return analysis

    except Exception as e:
        print(f"Error analyzing attribute: {str(e)}")
        return None

batch_vector_to_raster(vector_path, output_dir, attribute_field=None, reference_rasters=None, bounds_list=None, output_filename_pattern='{vector_name}_{index}', pixel_size=1.0, all_touched=False, fill_value=0, dtype=np.uint8, nodata=None)

Batch convert vector data to multiple rasters based on different extents or reference rasters.

Parameters:

Name	Type	Description	Default
vector_path	str or GeoDataFrame	Path to the input vector file or a GeoDataFrame.	required
output_dir	str	Directory to save output raster files.	required
attribute_field	str	Field name in the vector data to use for pixel values.	None
reference_rasters	list	List of paths to reference rasters for dimensions, transform and CRS.	None
bounds_list	list	List of bounds tuples (left, bottom, right, top) to use if reference_rasters not provided.	None
output_filename_pattern	str	Pattern for output filenames. Can include {vector_name} and {index} placeholders.	'{vector_name}_{index}'
pixel_size	float or tuple	Pixel size to use if reference_rasters not provided.	1.0
all_touched	bool	If True, all pixels touched by geometries will be burned in.	False
fill_value	int	Value to fill the raster with before burning in features.	0
dtype	dtype	Data type of the output raster.	uint8
nodata	int	No data value for the output raster.	None

Returns:

Type	Description
List[str]	List[str]: List of paths to the created raster files.

Source code in geoai/utils/raster.py

def batch_vector_to_raster(
    vector_path,
    output_dir,
    attribute_field=None,
    reference_rasters=None,
    bounds_list=None,
    output_filename_pattern="{vector_name}_{index}",
    pixel_size=1.0,
    all_touched=False,
    fill_value=0,
    dtype=np.uint8,
    nodata=None,
) -> List[str]:
    """
    Batch convert vector data to multiple rasters based on different extents or reference rasters.

    Args:
        vector_path (str or GeoDataFrame): Path to the input vector file or a GeoDataFrame.
        output_dir (str): Directory to save output raster files.
        attribute_field (str): Field name in the vector data to use for pixel values.
        reference_rasters (list): List of paths to reference rasters for dimensions, transform and CRS.
        bounds_list (list): List of bounds tuples (left, bottom, right, top) to use if reference_rasters not provided.
        output_filename_pattern (str): Pattern for output filenames.
            Can include {vector_name} and {index} placeholders.
        pixel_size (float or tuple): Pixel size to use if reference_rasters not provided.
        all_touched (bool): If True, all pixels touched by geometries will be burned in.
        fill_value (int): Value to fill the raster with before burning in features.
        dtype (numpy.dtype): Data type of the output raster.
        nodata (int): No data value for the output raster.

    Returns:
        List[str]: List of paths to the created raster files.
    """
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)

    # Load vector data if it's a path
    if isinstance(vector_path, str):
        gdf = gpd.read_file(vector_path)
        vector_name = os.path.splitext(os.path.basename(vector_path))[0]
    else:
        gdf = vector_path
        vector_name = "vector"

    # Check input parameters
    if reference_rasters is None and bounds_list is None:
        raise ValueError("Either reference_rasters or bounds_list must be provided.")

    # Use reference_rasters if provided, otherwise use bounds_list
    if reference_rasters is not None:
        sources = reference_rasters
        is_raster_reference = True
    else:
        sources = bounds_list
        is_raster_reference = False

    # Create output filenames
    output_files = []

    # Process each source (reference raster or bounds)
    for i, source in enumerate(tqdm(sources, desc="Processing")):
        # Generate output filename
        output_filename = output_filename_pattern.format(
            vector_name=vector_name, index=i
        )
        if not output_filename.endswith(".tif"):
            output_filename += ".tif"
        output_path = os.path.join(output_dir, output_filename)

        if is_raster_reference:
            # Use reference raster
            vector_to_raster(
                vector_path=gdf,
                output_path=output_path,
                reference_raster=source,
                attribute_field=attribute_field,
                all_touched=all_touched,
                fill_value=fill_value,
                dtype=dtype,
                nodata=nodata,
            )
        else:
            # Use bounds
            vector_to_raster(
                vector_path=gdf,
                output_path=output_path,
                bounds=source,
                pixel_size=pixel_size,
                attribute_field=attribute_field,
                all_touched=all_touched,
                fill_value=fill_value,
                dtype=dtype,
                nodata=nodata,
            )

        output_files.append(output_path)

    return output_files

bbox_to_xy(src_fp, coords, coord_crs='epsg:4326', **kwargs)

Converts a list of coordinates to pixel coordinates, i.e., (col, row) coordinates. Note that map bbox coords is [minx, miny, maxx, maxy] from bottomleft to topright While rasterio bbox coords is [minx, max, maxx, min] from topleft to bottomright

Parameters:

Name	Type	Description	Default
src_fp	str	The source raster file path.	required
coords	list	A list of coordinates in the format of [[minx, miny, maxx, maxy], [minx, miny, maxx, maxy], ...]	required
coord_crs	str	The coordinate CRS of the input coordinates. Defaults to "epsg:4326".	'epsg:4326'

Returns:

Name	Type	Description
list	List[float]	A list of pixel coordinates in the format of [[minx, maxy, maxx, miny], ...] from top left to bottom right.

Source code in geoai/utils/conversion.py

def bbox_to_xy(
    src_fp: str, coords: List[float], coord_crs: str = "epsg:4326", **kwargs: Any
) -> List[float]:
    """Converts a list of coordinates to pixel coordinates, i.e., (col, row) coordinates.
        Note that map bbox coords is [minx, miny, maxx, maxy] from bottomleft to topright
        While rasterio bbox coords is [minx, max, maxx, min] from topleft to bottomright

    Args:
        src_fp (str): The source raster file path.
        coords (list): A list of coordinates in the format of [[minx, miny, maxx, maxy], [minx, miny, maxx, maxy], ...]
        coord_crs (str, optional): The coordinate CRS of the input coordinates. Defaults to "epsg:4326".

    Returns:
        list: A list of pixel coordinates in the format of [[minx, maxy, maxx, miny], ...] from top left to bottom right.
    """

    if isinstance(coords, str):
        gdf = gpd.read_file(coords)
        coords = gdf.geometry.bounds.values.tolist()
        if gdf.crs is not None:
            coord_crs = f"epsg:{gdf.crs.to_epsg()}"
    elif isinstance(coords, np.ndarray):
        coords = coords.tolist()
    if isinstance(coords, dict):
        import json

        geojson = json.dumps(coords)
        gdf = gpd.read_file(geojson, driver="GeoJSON")
        coords = gdf.geometry.bounds.values.tolist()

    elif not isinstance(coords, list):
        raise ValueError("coords must be a list of coordinates.")

    if not isinstance(coords[0], list):
        coords = [coords]

    new_coords = []

    with rasterio.open(src_fp) as src:
        width = src.width
        height = src.height

        for coord in coords:
            minx, miny, maxx, maxy = coord

            if coord_crs != src.crs:
                from rasterio.warp import transform as transform_coords

                minx_list, miny_list = transform_coords(
                    coord_crs, src.crs, [minx], [miny], **kwargs
                )
                maxx_list, maxy_list = transform_coords(
                    coord_crs, src.crs, [maxx], [maxy], **kwargs
                )

                minx, miny = minx_list[0], miny_list[0]
                maxx, maxy = maxx_list[0], maxy_list[0]

                rows1, cols1 = rasterio.transform.rowcol(
                    src.transform, minx, miny, **kwargs
                )
                rows2, cols2 = rasterio.transform.rowcol(
                    src.transform, maxx, maxy, **kwargs
                )

                new_coords.append([cols1, rows1, cols2, rows2])

            else:
                new_coords.append([minx, miny, maxx, maxy])

    result = []

    for coord in new_coords:
        minx, miny, maxx, maxy = coord

        if (
            minx >= 0
            and miny >= 0
            and maxx >= 0
            and maxy >= 0
            and minx < width
            and miny < height
            and maxx < width
            and maxy < height
        ):
            # Note that map bbox coords is [minx, miny, maxx, maxy] from bottomleft to topright
            # While rasterio bbox coords is [minx, max, maxx, min] from topleft to bottomright
            result.append([minx, maxy, maxx, miny])

    if len(result) == 0:
        print("No valid pixel coordinates found.")
        return None
    elif len(result) == 1:
        return result[0]
    elif len(result) < len(coords):
        print("Some coordinates are out of the image boundary.")

    return result

boxes_to_vector(coords, src_crs, dst_crs='EPSG:4326', output=None, **kwargs)

Convert a list of bounding box coordinates to vector data.

Parameters:

Name	Type	Description	Default
coords	list	A list of bounding box coordinates in the format [[left, top, right, bottom], [left, top, right, bottom], ...].	required
src_crs	int or str	The EPSG code or proj4 string representing the source coordinate reference system (CRS) of the input coordinates.	required
dst_crs	int or str	The EPSG code or proj4 string representing the destination CRS to reproject the data (default is "EPSG:4326").	'EPSG:4326'
output	str or None	The full file path (including the directory and filename without the extension) where the vector data should be saved. If None (default), the function returns the GeoDataFrame without saving it to a file.	None
**kwargs	Any	Additional keyword arguments to pass to geopandas.GeoDataFrame.to_file() when saving the vector data.	{}

Returns:

Type	Description
GeoDataFrame	geopandas.GeoDataFrame or None: The GeoDataFrame with the converted vector data if output is None, otherwise None if the data is saved to a file.

Source code in geoai/utils/vector.py

def boxes_to_vector(
    coords: Union[List[List[float]], np.ndarray],
    src_crs: str,
    dst_crs: str = "EPSG:4326",
    output: Optional[str] = None,
    **kwargs: Any,
) -> gpd.GeoDataFrame:
    """
    Convert a list of bounding box coordinates to vector data.

    Args:
        coords (list): A list of bounding box coordinates in the format [[left, top, right, bottom], [left, top, right, bottom], ...].
        src_crs (int or str): The EPSG code or proj4 string representing the source coordinate reference system (CRS) of the input coordinates.
        dst_crs (int or str, optional): The EPSG code or proj4 string representing the destination CRS to reproject the data (default is "EPSG:4326").
        output (str or None, optional): The full file path (including the directory and filename without the extension) where the vector data should be saved.
                                       If None (default), the function returns the GeoDataFrame without saving it to a file.
        **kwargs: Additional keyword arguments to pass to geopandas.GeoDataFrame.to_file() when saving the vector data.

    Returns:
        geopandas.GeoDataFrame or None: The GeoDataFrame with the converted vector data if output is None, otherwise None if the data is saved to a file.
    """

    from shapely.geometry import box

    # Create a list of Shapely Polygon objects based on the provided coordinates
    polygons = [box(*coord) for coord in coords]

    # Create a GeoDataFrame with the Shapely Polygon objects
    gdf = gpd.GeoDataFrame({"geometry": polygons}, crs=src_crs)

    # Reproject the GeoDataFrame to the specified EPSG code
    gdf_reprojected = gdf.to_crs(dst_crs)

    if output is not None:
        gdf_reprojected.to_file(output, **kwargs)
    else:
        return gdf_reprojected

calc_stats(dataset, divide_by=1.0)

Calculate the statistics (mean and std) for the entire dataset.

This function is adapted from the plot_batch() function in the torchgeo library at https://torchgeo.readthedocs.io/en/stable/tutorials/earth_surface_water.html. Credit to the torchgeo developers for the original implementation.

Warning: This is an approximation. The correct value should take into account the mean for the whole dataset for computing individual stds.

Parameters:

Name	Type	Description	Default
dataset	RasterDataset	The dataset to calculate statistics for.	required
divide_by	float	The value to divide the image data by. Defaults to 1.0.	1.0

Returns:

Type	Description
Tuple[ndarray, ndarray]	Tuple[np.ndarray, np.ndarray]: The mean and standard deviation for each band.

Source code in geoai/utils/raster.py

def calc_stats(dataset, divide_by: float = 1.0) -> Tuple[np.ndarray, np.ndarray]:
    """
    Calculate the statistics (mean and std) for the entire dataset.

    This function is adapted from the plot_batch() function in the torchgeo library at
    https://torchgeo.readthedocs.io/en/stable/tutorials/earth_surface_water.html.
    Credit to the torchgeo developers for the original implementation.

    Warning: This is an approximation. The correct value should take into account the
    mean for the whole dataset for computing individual stds.

    Args:
        dataset (RasterDataset): The dataset to calculate statistics for.
        divide_by (float, optional): The value to divide the image data by. Defaults to 1.0.

    Returns:
        Tuple[np.ndarray, np.ndarray]: The mean and standard deviation for each band.
    """

    # To avoid loading the entire dataset in memory, we will loop through each img
    # The filenames will be retrieved from the dataset's rtree index
    files = [
        item.object
        for item in dataset.index.intersection(dataset.index.bounds, objects=True)
    ]

    # Resetting statistics
    accum_mean = 0
    accum_std = 0

    for file in files:
        img = rasterio.open(file).read() / divide_by  # type: ignore
        accum_mean += img.reshape((img.shape[0], -1)).mean(axis=1)
        accum_std += img.reshape((img.shape[0], -1)).std(axis=1)

    # at the end, we shall have 2 vectors with length n=chnls
    # we will average them considering the number of images
    return accum_mean / len(files), accum_std / len(files)

clip_raster_by_bbox(input_raster, output_raster, bbox, bands=None, bbox_type='geo', bbox_crs=None)

Clip a raster dataset using a bounding box and optionally select specific bands.

Parameters:

Name	Type	Description	Default
input_raster	str	Path to the input raster file.	required
output_raster	str	Path where the clipped raster will be saved.	required
bbox	tuple	Bounding box coordinates either as: - Geographic coordinates (minx, miny, maxx, maxy) if bbox_type="geo" - Pixel indices (min_row, min_col, max_row, max_col) if bbox_type="pixel"	required
bands	list	List of band indices to keep (1-based indexing). If None, all bands will be kept.	None
bbox_type	str	Type of bounding box coordinates. Either "geo" for geographic coordinates or "pixel" for row/column indices. Default is "geo".	'geo'
bbox_crs	str or dict	CRS of the bbox if different from the raster CRS. Can be provided as EPSG code (e.g., "EPSG:4326") or as a proj4 string. Only applies when bbox_type="geo". If None, assumes bbox is in the same CRS as the raster.	None

Returns:

Name	Type	Description
str	str	Path to the clipped output raster.

Raises:

Type	Description
ImportError	If required dependencies are not installed.
ValueError	If the bbox is invalid, bands are out of range, or bbox_type is invalid.
RuntimeError	If the clipping operation fails.

Examples:

Clip using geographic coordinates in the same CRS as the raster

>>> clip_raster_by_bbox('input.tif', 'clipped_geo.tif', (100, 200, 300, 400))
'clipped_geo.tif'

Clip using WGS84 coordinates when the raster is in a different CRS

>>> clip_raster_by_bbox('input.tif', 'clipped_wgs84.tif', (-122.5, 37.7, -122.4, 37.8),
...                     bbox_crs="EPSG:4326")
'clipped_wgs84.tif'

Clip using row/column indices

>>> clip_raster_by_bbox('input.tif', 'clipped_pixel.tif', (50, 100, 150, 200),
...                     bbox_type="pixel")
'clipped_pixel.tif'

Clip with band selection

>>> clip_raster_by_bbox('input.tif', 'clipped_bands.tif', (100, 200, 300, 400),
...                     bands=[1, 3])
'clipped_bands.tif'

Source code in geoai/utils/raster.py

def clip_raster_by_bbox(
    input_raster: str,
    output_raster: str,
    bbox: List[float],
    bands: Optional[List[int]] = None,
    bbox_type: str = "geo",
    bbox_crs: Optional[str] = None,
) -> str:
    """
    Clip a raster dataset using a bounding box and optionally select specific bands.

    Args:
        input_raster (str): Path to the input raster file.
        output_raster (str): Path where the clipped raster will be saved.
        bbox (tuple): Bounding box coordinates either as:
                     - Geographic coordinates (minx, miny, maxx, maxy) if bbox_type="geo"
                     - Pixel indices (min_row, min_col, max_row, max_col) if bbox_type="pixel"
        bands (list, optional): List of band indices to keep (1-based indexing).
                               If None, all bands will be kept.
        bbox_type (str, optional): Type of bounding box coordinates. Either "geo" for
                                  geographic coordinates or "pixel" for row/column indices.
                                  Default is "geo".
        bbox_crs (str or dict, optional): CRS of the bbox if different from the raster CRS.
                                         Can be provided as EPSG code (e.g., "EPSG:4326") or
                                         as a proj4 string. Only applies when bbox_type="geo".
                                         If None, assumes bbox is in the same CRS as the raster.

    Returns:
        str: Path to the clipped output raster.

    Raises:
        ImportError: If required dependencies are not installed.
        ValueError: If the bbox is invalid, bands are out of range, or bbox_type is invalid.
        RuntimeError: If the clipping operation fails.

    Examples:
        Clip using geographic coordinates in the same CRS as the raster
        >>> clip_raster_by_bbox('input.tif', 'clipped_geo.tif', (100, 200, 300, 400))
        'clipped_geo.tif'

        Clip using WGS84 coordinates when the raster is in a different CRS
        >>> clip_raster_by_bbox('input.tif', 'clipped_wgs84.tif', (-122.5, 37.7, -122.4, 37.8),
        ...                     bbox_crs="EPSG:4326")
        'clipped_wgs84.tif'

        Clip using row/column indices
        >>> clip_raster_by_bbox('input.tif', 'clipped_pixel.tif', (50, 100, 150, 200),
        ...                     bbox_type="pixel")
        'clipped_pixel.tif'

        Clip with band selection
        >>> clip_raster_by_bbox('input.tif', 'clipped_bands.tif', (100, 200, 300, 400),
        ...                     bands=[1, 3])
        'clipped_bands.tif'
    """
    from rasterio.transform import from_bounds
    from rasterio.warp import transform_bounds

    # Validate bbox_type
    if bbox_type not in ["geo", "pixel"]:
        raise ValueError("bbox_type must be either 'geo' or 'pixel'")

    # Validate bbox
    if len(bbox) != 4:
        raise ValueError("bbox must contain exactly 4 values")

    # Open the source raster
    with rasterio.open(input_raster) as src:
        # Get the source CRS
        src_crs = src.crs

        # Handle different bbox types
        if bbox_type == "geo":
            minx, miny, maxx, maxy = bbox

            # Validate geographic bbox
            if minx >= maxx or miny >= maxy:
                raise ValueError(
                    "Invalid geographic bbox. Expected (minx, miny, maxx, maxy) where minx < maxx and miny < maxy"
                )

            # If bbox_crs is provided and different from the source CRS, transform the bbox
            if bbox_crs is not None and bbox_crs != src_crs:
                try:
                    # Transform bbox coordinates from bbox_crs to src_crs
                    minx, miny, maxx, maxy = transform_bounds(
                        bbox_crs, src_crs, minx, miny, maxx, maxy
                    )
                except Exception as e:
                    raise ValueError(
                        f"Failed to transform bbox from {bbox_crs} to {src_crs}: {str(e)}"
                    )

            # Calculate the pixel window from geographic coordinates
            window = src.window(minx, miny, maxx, maxy)

            # Use the same bounds for the output transform
            output_bounds = (minx, miny, maxx, maxy)

        else:  # bbox_type == "pixel"
            min_row, min_col, max_row, max_col = bbox

            # Validate pixel bbox
            if min_row >= max_row or min_col >= max_col:
                raise ValueError(
                    "Invalid pixel bbox. Expected (min_row, min_col, max_row, max_col) where min_row < max_row and min_col < max_col"
                )

            if (
                min_row < 0
                or min_col < 0
                or max_row > src.height
                or max_col > src.width
            ):
                raise ValueError(
                    f"Pixel indices out of bounds. Raster dimensions are {src.height} rows x {src.width} columns"
                )

            # Create a window from pixel coordinates
            window = Window(min_col, min_row, max_col - min_col, max_row - min_row)

            # Calculate the geographic bounds for this window
            window_transform = src.window_transform(window)
            output_bounds = rasterio.transform.array_bounds(
                window.height, window.width, window_transform
            )
            # Reorder to (minx, miny, maxx, maxy)
            output_bounds = (
                output_bounds[0],
                output_bounds[1],
                output_bounds[2],
                output_bounds[3],
            )

        # Get window dimensions
        window_width = int(window.width)
        window_height = int(window.height)

        # Check if the window is valid
        if window_width <= 0 or window_height <= 0:
            raise ValueError("Bounding box results in an empty window")

        # Handle band selection
        if bands is None:
            # Use all bands
            bands_to_read = list(range(1, src.count + 1))
        else:
            # Validate band indices
            if not all(1 <= b <= src.count for b in bands):
                raise ValueError(f"Band indices must be between 1 and {src.count}")
            bands_to_read = bands

        # Calculate new transform for the clipped raster
        new_transform = from_bounds(
            output_bounds[0],
            output_bounds[1],
            output_bounds[2],
            output_bounds[3],
            window_width,
            window_height,
        )

        # Create a metadata dictionary for the output
        out_meta = src.meta.copy()
        out_meta.update(
            {
                "height": window_height,
                "width": window_width,
                "transform": new_transform,
                "count": len(bands_to_read),
            }
        )

        # Read the data for the selected bands
        data = []
        for band_idx in bands_to_read:
            band_data = src.read(band_idx, window=window)
            data.append(band_data)

        # Stack the bands into a single array
        if len(data) > 1:
            clipped_data = np.stack(data)
        else:
            clipped_data = data[0][np.newaxis, :, :]

        # Write the output raster
        with rasterio.open(output_raster, "w", **out_meta) as dst:
            dst.write(clipped_data)

    return output_raster

compute_class_weights(labels_dir, num_classes, ignore_index=-100, custom_multipliers=None, max_weight=50.0, use_inverse_frequency=True)

Compute class weights for imbalanced datasets with optional custom multipliers and maximum weight cap.

Parameters:

Name	Type	Description	Default
labels_dir	str	Directory containing label files	required
num_classes	int	Number of classes	required
ignore_index	Union[int, bool]	Class index to ignore when computing weights. - If int: specific class to ignore (pixels will be excluded from weight calc) - If False: no class ignored (all classes contribute to weights)	-100
custom_multipliers	Optional[Dict[int, float]]	Custom multipliers for specific classes after inverse frequency calculation. Format: {class_id: multiplier} Example: {1: 0.5, 7: 2.0} - reduce class 1 weight by half, double class 7 weight	None
max_weight	float	Maximum allowed weight value to prevent extreme values (default: 50.0)	50.0
use_inverse_frequency	bool	Whether to compute inverse frequency weights. - True (default): Compute inverse frequency weights, then apply custom multipliers - False: Use uniform weights (1.0) for all classes, then apply custom multipliers	True

Returns:

Type	Description
Tensor	Tensor of class weights (num_classes,) with custom adjustments and maximum weight cap applied

Source code in geoai/landcover_train.py

def compute_class_weights(
    labels_dir: str,
    num_classes: int,
    ignore_index: Union[int, bool] = -100,
    custom_multipliers: Optional[Dict[int, float]] = None,
    max_weight: float = 50.0,
    use_inverse_frequency: bool = True,
) -> torch.Tensor:
    """
    Compute class weights for imbalanced datasets with optional custom multipliers and maximum weight cap.

    Args:
        labels_dir: Directory containing label files
        num_classes: Number of classes
        ignore_index: Class index to ignore when computing weights.
            - If int: specific class to ignore (pixels will be excluded from weight calc)
            - If False: no class ignored (all classes contribute to weights)
        custom_multipliers: Custom multipliers for specific classes after inverse frequency calculation.
            Format: {class_id: multiplier}
            Example: {1: 0.5, 7: 2.0} - reduce class 1 weight by half, double class 7 weight
        max_weight: Maximum allowed weight value to prevent extreme values (default: 50.0)
        use_inverse_frequency: Whether to compute inverse frequency weights.
            - True (default): Compute inverse frequency weights, then apply custom multipliers
            - False: Use uniform weights (1.0) for all classes, then apply custom multipliers

    Returns:
        Tensor of class weights (num_classes,) with custom adjustments and maximum weight cap applied
    """
    import os
    import rasterio
    from collections import Counter

    # Count pixels for each class
    class_counts = Counter()
    total_pixels = 0

    # Get all label files
    label_extensions = (".tif", ".tiff", ".png", ".jpg", ".jpeg")
    label_files = [
        os.path.join(labels_dir, f)
        for f in os.listdir(labels_dir)
        if f.lower().endswith(label_extensions)
    ]

    print(f"Computing class weights from {len(label_files)} label files...")

    for label_file in label_files:
        try:
            with rasterio.open(label_file) as src:
                label_data = src.read(1)
                for class_id in range(num_classes):
                    if isinstance(ignore_index, int) and class_id == ignore_index:
                        continue
                    count = (label_data == class_id).sum()
                    class_counts[class_id] += int(count)
                    total_pixels += int(count)
        except Exception as e:
            print(f"Warning: Could not read {label_file}: {e}")
            continue

    if total_pixels == 0:
        raise ValueError("No valid pixels found in label files")

    # Initialize weights
    weights = torch.ones(num_classes)

    if use_inverse_frequency:
        # Compute inverse frequency weights
        for class_id in range(num_classes):
            if isinstance(ignore_index, int) and class_id == ignore_index:
                weights[class_id] = 0.0
            elif class_counts[class_id] > 0:
                # Inverse frequency: total_pixels / class_pixels
                weights[class_id] = total_pixels / class_counts[class_id]
            else:
                weights[class_id] = 0.0

        # Normalize to have mean weight of 1.0
        non_zero_weights = weights[weights > 0]
        if len(non_zero_weights) > 0:
            weights = weights / non_zero_weights.mean()
    else:
        # Use uniform weights (all 1.0)
        for class_id in range(num_classes):
            if isinstance(ignore_index, int) and class_id == ignore_index:
                weights[class_id] = 0.0

    # Apply custom multipliers if provided
    if custom_multipliers:
        print(f"\nApplying custom multipliers: {custom_multipliers}")
        for class_id, multiplier in custom_multipliers.items():
            if class_id < 0 or class_id >= num_classes:
                print(f"Warning: Invalid class_id {class_id}, skipping")
                continue

            original_weight = weights[class_id].item()
            weights[class_id] = weights[class_id] * multiplier
            print(
                f"  Class {class_id}: {original_weight:.4f} × {multiplier} = {weights[class_id].item():.4f}"
            )
    else:
        print("\nNo custom multipliers provided, using computed weights as-is")

    # Apply maximum weight cap to prevent extreme values
    weights_capped = False
    print(f"\nApplying maximum weight cap of {max_weight}...")
    for class_id in range(num_classes):
        if weights[class_id] > max_weight:
            print(
                f"  Class {class_id}: {weights[class_id].item():.4f} → {max_weight} (capped)"
            )
            weights[class_id] = max_weight
            weights_capped = True

    if not weights_capped:
        print("  No weights exceeded the cap")

    print(f"\nClass pixel counts: {dict(class_counts)}")
    print(f"\nFinal class weights:")
    for class_id in range(num_classes):
        pixel_count = class_counts.get(class_id, 0)
        percent = (pixel_count / total_pixels * 100) if total_pixels > 0 else 0
        print(
            f"  Class {class_id}: weight={weights[class_id].item():.4f}, "
            f"pixels={pixel_count:,} ({percent:.2f}%)"
        )

    if isinstance(ignore_index, int) and 0 <= ignore_index < num_classes:
        print(f"\nNote: Class {ignore_index} (ignore_index) has weight 0.0")

    return weights

coords_to_xy(src_fp, coords, coord_crs='epsg:4326', return_out_of_bounds=False, **kwargs)

Converts a list or array of coordinates to pixel coordinates, i.e., (col, row) coordinates.

Parameters:

Name	Type	Description	Default
src_fp	str	The source raster file path.	required
coords	ndarray	A 2D or 3D array of coordinates. Can be of shape [[x1, y1], [x2, y2], ...] or [[[x1, y1]], [[x2, y2]], ...].	required
coord_crs	str	The coordinate CRS of the input coordinates. Defaults to "epsg:4326".	'epsg:4326'
return_out_of_bounds	bool	Whether to return out-of-bounds coordinates. Defaults to False.	False
**kwargs	Any	Additional keyword arguments to pass to rasterio.transform.rowcol.	{}

Returns:

Type	Description
ndarray	A 2D or 3D array of pixel coordinates in the same format as the input.

Source code in geoai/utils/conversion.py

def coords_to_xy(
    src_fp: str,
    coords: np.ndarray,
    coord_crs: str = "epsg:4326",
    return_out_of_bounds: bool = False,
    **kwargs: Any,
) -> np.ndarray:
    """Converts a list or array of coordinates to pixel coordinates, i.e., (col, row) coordinates.

    Args:
        src_fp: The source raster file path.
        coords: A 2D or 3D array of coordinates. Can be of shape [[x1, y1], [x2, y2], ...]
                or [[[x1, y1]], [[x2, y2]], ...].
        coord_crs: The coordinate CRS of the input coordinates. Defaults to "epsg:4326".
        return_out_of_bounds: Whether to return out-of-bounds coordinates. Defaults to False.
        **kwargs: Additional keyword arguments to pass to rasterio.transform.rowcol.

    Returns:
        A 2D or 3D array of pixel coordinates in the same format as the input.
    """
    from rasterio.warp import transform as transform_coords

    out_of_bounds = []
    if isinstance(coords, np.ndarray):
        input_is_3d = coords.ndim == 3  # Check if the input is a 3D array
    else:
        input_is_3d = False

    # Flatten the 3D array to 2D if necessary
    if input_is_3d:
        original_shape = coords.shape  # Store the original shape
        coords = coords.reshape(-1, 2)  # Flatten to 2D

    # Convert ndarray to a list if necessary
    if isinstance(coords, np.ndarray):
        coords = coords.tolist()

    xs, ys = zip(*coords)
    with rasterio.open(src_fp) as src:
        width = src.width
        height = src.height
        if coord_crs != src.crs:
            xs, ys = transform_coords(coord_crs, src.crs, xs, ys, **kwargs)
        rows, cols = rasterio.transform.rowcol(src.transform, xs, ys, **kwargs)

    result = [[col, row] for col, row in zip(cols, rows)]

    output = []

    for i, (x, y) in enumerate(result):
        if x >= 0 and y >= 0 and x < width and y < height:
            output.append([x, y])
        else:
            out_of_bounds.append(i)

    # Convert the output back to the original shape if input was 3D
    output = np.array(output)
    if input_is_3d:
        output = output.reshape(original_shape)

    # Handle cases where no valid pixel coordinates are found
    if len(output) == 0:
        print("No valid pixel coordinates found.")
    elif len(output) < len(coords):
        print("Some coordinates are out of the image boundary.")

    if return_out_of_bounds:
        return output, out_of_bounds
    else:
        return output

create_vector_data(m=None, properties=None, time_format='%Y%m%dT%H%M%S', column_widths=(9, 3), map_height='600px', out_dir=None, filename_prefix='', file_ext='geojson', add_mapillary=False, style='photo', radius=5e-05, width=300, height=420, frame_border=0, **kwargs)

Generates a widget-based interface for creating and managing vector data on a map.

This function creates an interactive widget interface that allows users to draw features (points, lines, polygons) on a map, assign properties to these features, and export them as GeoJSON files. The interface includes a map, a sidebar for property management, and buttons for saving, exporting, and resetting the data.

Parameters:

Name	Type	Description	Default
m	Map	An existing Map object. If not provided, a default map with basemaps and drawing controls will be created. Defaults to None.	None
properties	Dict[str, List[Any]]	A dictionary where keys are property names and values are lists of possible values for each property. These properties can be assigned to the drawn features. Defaults to None.	None
time_format	str	The format string for the timestamp used in the exported filename. Defaults to "%Y%m%dT%H%M%S".	'%Y%m%dT%H%M%S'
column_widths	Optional[List[int]]	A list of two integers specifying the relative widths of the map and sidebar columns. Defaults to (9, 3).	(9, 3)
map_height	str	The height of the map widget. Defaults to "600px".	'600px'
out_dir	str	The directory where the exported GeoJSON files will be saved. If not provided, the current working directory is used. Defaults to None.	None
filename_prefix	str	A prefix to be added to the exported filename. Defaults to "".	''
file_ext	str	The file extension for the exported file. Defaults to "geojson".	'geojson'
add_mapillary	bool	Whether to add a Mapillary image widget that displays the nearest image to the clicked point on the map. Defaults to False.	False
style	str	The style of the Mapillary image widget. Can be "classic", "photo", or "split". Defaults to "photo".	'photo'
radius	float	The radius (in degrees) used to search for the nearest Mapillary image. Defaults to 0.00005 degrees.	5e-05
width	int	The width of the Mapillary image widget. Defaults to 300.	300
height	int	The height of the Mapillary image widget. Defaults to 420.	420
frame_border	int	The width of the frame border for the Mapillary image widget. Defaults to 0.	0
**kwargs	Any	Additional keyword arguments that may be passed to the function.	{}

Returns:

Type	Description
Any	widgets.VBox: A vertical box widget containing the map, sidebar, and control buttons.

Example

properties = { ... "Type": ["Residential", "Commercial", "Industrial"], ... "Area": [100, 200, 300], ... } widget = create_vector_data(properties=properties) display(widget) # Display the widget in a Jupyter notebook

Source code in geoai/geoai.py

def create_vector_data(
    m: Optional[LeafMap] = None,
    properties: Optional[Dict[str, List[Any]]] = None,
    time_format: str = "%Y%m%dT%H%M%S",
    column_widths: Optional[List[int]] = (9, 3),
    map_height: str = "600px",
    out_dir: Optional[str] = None,
    filename_prefix: str = "",
    file_ext: str = "geojson",
    add_mapillary: bool = False,
    style: str = "photo",
    radius: float = 0.00005,
    width: int = 300,
    height: int = 420,
    frame_border: int = 0,
    **kwargs: Any,
) -> Any:
    """Generates a widget-based interface for creating and managing vector data on a map.

    This function creates an interactive widget interface that allows users to draw features
    (points, lines, polygons) on a map, assign properties to these features, and export them
    as GeoJSON files. The interface includes a map, a sidebar for property management, and
    buttons for saving, exporting, and resetting the data.

    Args:
        m (Map, optional): An existing Map object. If not provided, a default map with
            basemaps and drawing controls will be created. Defaults to None.
        properties (Dict[str, List[Any]], optional): A dictionary where keys are property names
            and values are lists of possible values for each property. These properties can be
            assigned to the drawn features. Defaults to None.
        time_format (str, optional): The format string for the timestamp used in the exported
            filename. Defaults to "%Y%m%dT%H%M%S".
        column_widths (Optional[List[int]], optional): A list of two integers specifying the
            relative widths of the map and sidebar columns. Defaults to (9, 3).
        map_height (str, optional): The height of the map widget. Defaults to "600px".
        out_dir (str, optional): The directory where the exported GeoJSON files will be saved.
            If not provided, the current working directory is used. Defaults to None.
        filename_prefix (str, optional): A prefix to be added to the exported filename.
            Defaults to "".
        file_ext (str, optional): The file extension for the exported file. Defaults to "geojson".
        add_mapillary (bool, optional): Whether to add a Mapillary image widget that displays the
            nearest image to the clicked point on the map. Defaults to False.
        style (str, optional): The style of the Mapillary image widget. Can be "classic", "photo",
            or "split". Defaults to "photo".
        radius (float, optional): The radius (in degrees) used to search for the nearest Mapillary
            image. Defaults to 0.00005 degrees.
        width (int, optional): The width of the Mapillary image widget. Defaults to 300.
        height (int, optional): The height of the Mapillary image widget. Defaults to 420.
        frame_border (int, optional): The width of the frame border for the Mapillary image widget.
            Defaults to 0.
        **kwargs (Any): Additional keyword arguments that may be passed to the function.

    Returns:
        widgets.VBox: A vertical box widget containing the map, sidebar, and control buttons.

    Example:
        >>> properties = {
        ...     "Type": ["Residential", "Commercial", "Industrial"],
        ...     "Area": [100, 200, 300],
        ... }
        >>> widget = create_vector_data(properties=properties)
        >>> display(widget)  # Display the widget in a Jupyter notebook
    """
    return maplibregl.create_vector_data(
        m=m,
        properties=properties,
        time_format=time_format,
        column_widths=column_widths,
        map_height=map_height,
        out_dir=out_dir,
        filename_prefix=filename_prefix,
        file_ext=file_ext,
        add_mapillary=add_mapillary,
        style=style,
        radius=radius,
        width=width,
        height=height,
        frame_border=frame_border,
        **kwargs,
    )

dict_to_image(data_dict, output=None, **kwargs)

Convert a dictionary containing spatial data to a rasterio dataset or save it to a file. The dictionary should contain the following keys: "crs", "bounds", and "image". It can be generated from a TorchGeo dataset sampler.

This function transforms a dictionary with CRS, bounding box, and image data into a rasterio DatasetReader using leafmap's array_to_image utility after first converting to a rioxarray DataArray.

Parameters:

Name	Type	Description	Default
data_dict	Dict[str, Any]	A dictionary containing: - 'crs': A pyproj CRS object - 'bounds': A BoundingBox object with minx, maxx, miny, maxy attributes and optionally mint, maxt for temporal bounds - 'image': A tensor or array-like object with image data	required
output	Optional[str]	Optional path to save the image to a file. If not provided, the image will be returned as a rasterio DatasetReader object.	None
**kwargs	Any	Additional keyword arguments to pass to leafmap.array_to_image. Common options include: - colormap: str, name of the colormap (e.g., 'viridis', 'terrain') - vmin: float, minimum value for colormap scaling - vmax: float, maximum value for colormap scaling	{}

Returns:

Type	Description
Union[str, Any]	A rasterio DatasetReader object that can be used for visualization or
Union[str, Any]	further processing.

Examples:

>>> image = dict_to_image(
...     {'crs': CRS.from_epsg(26911), 'bounds': bbox, 'image': tensor},
...     colormap='terrain'
... )
>>> fig, ax = plt.subplots(figsize=(10, 10))
>>> show(image, ax=ax)

Source code in geoai/utils/conversion.py

def dict_to_image(
    data_dict: Dict[str, Any], output: Optional[str] = None, **kwargs: Any
) -> Union[str, Any]:
    """Convert a dictionary containing spatial data to a rasterio dataset or save it to
    a file. The dictionary should contain the following keys: "crs", "bounds", and "image".
    It can be generated from a TorchGeo dataset sampler.

    This function transforms a dictionary with CRS, bounding box, and image data
    into a rasterio DatasetReader using leafmap's array_to_image utility after
    first converting to a rioxarray DataArray.

    Args:
        data_dict: A dictionary containing:
            - 'crs': A pyproj CRS object
            - 'bounds': A BoundingBox object with minx, maxx, miny, maxy attributes
              and optionally mint, maxt for temporal bounds
            - 'image': A tensor or array-like object with image data
        output: Optional path to save the image to a file. If not provided, the image
            will be returned as a rasterio DatasetReader object.
        **kwargs: Additional keyword arguments to pass to leafmap.array_to_image.
            Common options include:
            - colormap: str, name of the colormap (e.g., 'viridis', 'terrain')
            - vmin: float, minimum value for colormap scaling
            - vmax: float, maximum value for colormap scaling

    Returns:
        A rasterio DatasetReader object that can be used for visualization or
        further processing.

    Examples:
        >>> image = dict_to_image(
        ...     {'crs': CRS.from_epsg(26911), 'bounds': bbox, 'image': tensor},
        ...     colormap='terrain'
        ... )
        >>> fig, ax = plt.subplots(figsize=(10, 10))
        >>> show(image, ax=ax)
    """
    da = dict_to_rioxarray(data_dict)

    if output is not None:
        out_dir = os.path.abspath(os.path.dirname(output))
        if not os.path.exists(out_dir):
            os.makedirs(out_dir, exist_ok=True)
        da.rio.to_raster(output)
        return output
    else:
        image = leafmap.array_to_image(da, **kwargs)
        return image

dict_to_rioxarray(data_dict)

Convert a dictionary to a xarray DataArray. The dictionary should contain the following keys: "crs", "bounds", and "image". It can be generated from a TorchGeo dataset sampler.

Parameters:

Name	Type	Description	Default
data_dict	Dict	The dictionary containing the data.	required

Returns:

Type	Description
DataArray	xr.DataArray: The xarray DataArray.

Source code in geoai/utils/conversion.py

def dict_to_rioxarray(data_dict: Dict) -> xr.DataArray:
    """Convert a dictionary to a xarray DataArray. The dictionary should contain the
    following keys: "crs", "bounds", and "image". It can be generated from a TorchGeo
    dataset sampler.

    Args:
        data_dict (Dict): The dictionary containing the data.

    Returns:
        xr.DataArray: The xarray DataArray.
    """

    from collections import namedtuple

    from affine import Affine

    BoundingBox = namedtuple("BoundingBox", ["minx", "maxx", "miny", "maxy"])

    # Extract components from the dictionary
    crs = data_dict["crs"]
    bounds = data_dict["bounds"]
    image_tensor = data_dict["image"]

    if hasattr(bounds, "left"):
        bounds = BoundingBox(bounds.left, bounds.right, bounds.bottom, bounds.top)

    # Convert tensor to numpy array if needed
    if hasattr(image_tensor, "numpy"):
        # For PyTorch tensors
        image_array = image_tensor.numpy()
    else:
        # If it's already a numpy array or similar
        image_array = np.array(image_tensor)

    # Calculate pixel resolution
    width = image_array.shape[2]  # Width is the size of the last dimension
    height = image_array.shape[1]  # Height is the size of the middle dimension

    res_x = (bounds.maxx - bounds.minx) / width
    res_y = (bounds.maxy - bounds.miny) / height

    # Create the transform matrix
    transform = Affine(res_x, 0.0, bounds.minx, 0.0, -res_y, bounds.maxy)

    # Create dimensions
    x_coords = np.linspace(bounds.minx + res_x / 2, bounds.maxx - res_x / 2, width)
    y_coords = np.linspace(bounds.maxy - res_y / 2, bounds.miny + res_y / 2, height)

    # If time dimension exists in the bounds
    if hasattr(bounds, "mint") and hasattr(bounds, "maxt"):
        # Create a single time value or range if needed
        t_coords = [
            bounds.mint
        ]  # Or np.linspace(bounds.mint, bounds.maxt, num_time_steps)

        # Create DataArray with time dimension
        dims = (
            ("band", "y", "x")
            if image_array.shape[0] <= 10
            else ("time", "band", "y", "x")
        )

        if dims[0] == "band":
            # For multi-band single time
            da = xr.DataArray(
                image_array,
                dims=dims,
                coords={
                    "band": np.arange(1, image_array.shape[0] + 1),
                    "y": y_coords,
                    "x": x_coords,
                },
            )
        else:
            # For multi-time multi-band
            da = xr.DataArray(
                image_array,
                dims=dims,
                coords={
                    "time": t_coords,
                    "band": np.arange(1, image_array.shape[1] + 1),
                    "y": y_coords,
                    "x": x_coords,
                },
            )
    else:
        # Create DataArray without time dimension
        da = xr.DataArray(
            image_array,
            dims=("band", "y", "x"),
            coords={
                "band": np.arange(1, image_array.shape[0] + 1),
                "y": y_coords,
                "x": x_coords,
            },
        )

    # Set spatial attributes
    da.rio.write_crs(crs, inplace=True)
    da.rio.write_transform(transform, inplace=True)

    return da

download_file(url, output_path=None, overwrite=False, unzip=True)

Download a file from a given URL with a progress bar. Optionally unzip the file if it's a ZIP archive.

Parameters:

Name	Type	Description	Default
url	str	The URL of the file to download.	required
output_path	str	The path where the downloaded file will be saved. If not provided, the filename from the URL will be used.	None
overwrite	bool	Whether to overwrite the file if it already exists.	False
unzip	bool	Whether to unzip the file if it is a ZIP archive.	True

Returns:

Name	Type	Description
str	str	The path to the downloaded file or the extracted directory.

Source code in geoai/utils/download.py

def download_file(
    url: str,
    output_path: Optional[str] = None,
    overwrite: bool = False,
    unzip: bool = True,
) -> str:
    """
    Download a file from a given URL with a progress bar.
    Optionally unzip the file if it's a ZIP archive.

    Args:
        url (str): The URL of the file to download.
        output_path (str, optional): The path where the downloaded file will be saved.
            If not provided, the filename from the URL will be used.
        overwrite (bool, optional): Whether to overwrite the file if it already exists.
        unzip (bool, optional): Whether to unzip the file if it is a ZIP archive.

    Returns:
        str: The path to the downloaded file or the extracted directory.
    """

    import zipfile

    from tqdm import tqdm

    if output_path is None:
        output_path = os.path.basename(url)

    if os.path.exists(output_path) and not overwrite:
        print(f"File already exists: {output_path}")
    else:
        # Download the file with a progress bar
        response = requests.get(url, stream=True, timeout=50)
        response.raise_for_status()
        total_size = int(response.headers.get("content-length", 0))

        with (
            open(output_path, "wb") as file,
            tqdm(
                desc=f"Downloading {os.path.basename(output_path)}",
                total=total_size,
                unit="B",
                unit_scale=True,
                unit_divisor=1024,
            ) as progress_bar,
        ):
            for chunk in response.iter_content(chunk_size=1024):
                if chunk:
                    file.write(chunk)
                    progress_bar.update(len(chunk))

    # If the file is a ZIP archive and unzip is True
    if unzip and zipfile.is_zipfile(output_path):
        extract_dir = os.path.splitext(output_path)[0]
        if not os.path.exists(extract_dir) or overwrite:
            with zipfile.ZipFile(output_path, "r") as zip_ref:
                zip_ref.extractall(extract_dir)
            print(f"Extracted to: {extract_dir}")
        return extract_dir

    return output_path

download_model_from_hf(model_path, repo_id=None)

Download the object detection model from Hugging Face.

Parameters:

Name	Type	Description	Default
model_path	str	Path to the model file.	required
repo_id	Optional[str]	Hugging Face repository ID.	None

Returns:

Type	Description
str	Path to the downloaded model file

Source code in geoai/utils/download.py

def download_model_from_hf(model_path: str, repo_id: Optional[str] = None) -> str:
    """
    Download the object detection model from Hugging Face.

    Args:
        model_path: Path to the model file.
        repo_id: Hugging Face repository ID.

    Returns:
        Path to the downloaded model file
    """
    from huggingface_hub import hf_hub_download

    try:

        # Define the repository ID and model filename
        if repo_id is None:
            print(
                "Repo is not specified, using default Hugging Face repo_id: giswqs/geoai"
            )
            repo_id = "giswqs/geoai"

        # Download the model
        model_path = hf_hub_download(repo_id=repo_id, filename=model_path)
        print(f"Model downloaded to: {model_path}")

        return model_path

    except Exception as e:
        print(f"Error downloading model from Hugging Face: {e}")
        print("Please specify a local model path or ensure internet connectivity.")
        raise

edit_vector_data(m=None, filename=None, properties=None, time_format='%Y%m%dT%H%M%S', column_widths=(9, 3), map_height='600px', out_dir=None, filename_prefix='', file_ext='geojson', add_mapillary=False, style='photo', radius=5e-05, width=300, height=420, frame_border=0, controls=None, position='top-right', fit_bounds_options=None, **kwargs)

Generates a widget-based interface for creating and managing vector data on a map.

This function creates an interactive widget interface that allows users to draw features (points, lines, polygons) on a map, assign properties to these features, and export them as GeoJSON files. The interface includes a map, a sidebar for property management, and buttons for saving, exporting, and resetting the data.

Parameters:

Name	Type	Description	Default
m	Map	An existing Map object. If not provided, a default map with basemaps and drawing controls will be created. Defaults to None.	None
filename	str or GeoDataFrame	The path to a GeoJSON file or a GeoDataFrame containing the vector data to be edited. Defaults to None.	None
properties	Dict[str, List[Any]]	A dictionary where keys are property names and values are lists of possible values for each property. These properties can be assigned to the drawn features. Defaults to None.	None
time_format	str	The format string for the timestamp used in the exported filename. Defaults to "%Y%m%dT%H%M%S".	'%Y%m%dT%H%M%S'
column_widths	Optional[List[int]]	A list of two integers specifying the relative widths of the map and sidebar columns. Defaults to (9, 3).	(9, 3)
map_height	str	The height of the map widget. Defaults to "600px".	'600px'
out_dir	str	The directory where the exported GeoJSON files will be saved. If not provided, the current working directory is used. Defaults to None.	None
filename_prefix	str	A prefix to be added to the exported filename. Defaults to "".	''
file_ext	str	The file extension for the exported file. Defaults to "geojson".	'geojson'
add_mapillary	bool	Whether to add a Mapillary image widget that displays the nearest image to the clicked point on the map. Defaults to False.	False
style	str	The style of the Mapillary image widget. Can be "classic", "photo", or "split". Defaults to "photo".	'photo'
radius	float	The radius (in degrees) used to search for the nearest Mapillary image. Defaults to 0.00005 degrees.	5e-05
width	int	The width of the Mapillary image widget. Defaults to 300.	300
height	int	The height of the Mapillary image widget. Defaults to 420.	420
frame_border	int	The width of the frame border for the Mapillary image widget. Defaults to 0.	0
controls	Optional[List[str]]	The drawing controls to be added to the map. Defaults to ["point", "polygon", "line_string", "trash"].	None
position	str	The position of the drawing controls on the map. Defaults to "top-right".	'top-right'
**kwargs	Any	Additional keyword arguments that may be passed to the function.	{}

Returns:

Type	Description
Any	widgets.VBox: A vertical box widget containing the map, sidebar, and control buttons.

Source code in geoai/geoai.py

def edit_vector_data(
    m: Optional[LeafMap] = None,
    filename: str = None,
    properties: Optional[Dict[str, List[Any]]] = None,
    time_format: str = "%Y%m%dT%H%M%S",
    column_widths: Optional[List[int]] = (9, 3),
    map_height: str = "600px",
    out_dir: Optional[str] = None,
    filename_prefix: str = "",
    file_ext: str = "geojson",
    add_mapillary: bool = False,
    style: str = "photo",
    radius: float = 0.00005,
    width: int = 300,
    height: int = 420,
    frame_border: int = 0,
    controls: Optional[List[str]] = None,
    position: str = "top-right",
    fit_bounds_options: Optional[Dict] = None,
    **kwargs: Any,
) -> Any:
    """Generates a widget-based interface for creating and managing vector data on a map.

    This function creates an interactive widget interface that allows users to draw features
    (points, lines, polygons) on a map, assign properties to these features, and export them
    as GeoJSON files. The interface includes a map, a sidebar for property management, and
    buttons for saving, exporting, and resetting the data.

    Args:
        m (Map, optional): An existing Map object. If not provided, a default map with
            basemaps and drawing controls will be created. Defaults to None.
        filename (str or gpd.GeoDataFrame): The path to a GeoJSON file or a GeoDataFrame
            containing the vector data to be edited. Defaults to None.
        properties (Dict[str, List[Any]], optional): A dictionary where keys are property names
            and values are lists of possible values for each property. These properties can be
            assigned to the drawn features. Defaults to None.
        time_format (str, optional): The format string for the timestamp used in the exported
            filename. Defaults to "%Y%m%dT%H%M%S".
        column_widths (Optional[List[int]], optional): A list of two integers specifying the
            relative widths of the map and sidebar columns. Defaults to (9, 3).
        map_height (str, optional): The height of the map widget. Defaults to "600px".
        out_dir (str, optional): The directory where the exported GeoJSON files will be saved.
            If not provided, the current working directory is used. Defaults to None.
        filename_prefix (str, optional): A prefix to be added to the exported filename.
            Defaults to "".
        file_ext (str, optional): The file extension for the exported file. Defaults to "geojson".
        add_mapillary (bool, optional): Whether to add a Mapillary image widget that displays the
            nearest image to the clicked point on the map. Defaults to False.
        style (str, optional): The style of the Mapillary image widget. Can be "classic", "photo",
            or "split". Defaults to "photo".
        radius (float, optional): The radius (in degrees) used to search for the nearest Mapillary
            image. Defaults to 0.00005 degrees.
        width (int, optional): The width of the Mapillary image widget. Defaults to 300.
        height (int, optional): The height of the Mapillary image widget. Defaults to 420.
        frame_border (int, optional): The width of the frame border for the Mapillary image widget.
            Defaults to 0.
        controls (Optional[List[str]], optional): The drawing controls to be added to the map.
            Defaults to ["point", "polygon", "line_string", "trash"].
        position (str, optional): The position of the drawing controls on the map. Defaults to "top-right".
        **kwargs (Any): Additional keyword arguments that may be passed to the function.

    Returns:
        widgets.VBox: A vertical box widget containing the map, sidebar, and control buttons.
    """
    return maplibregl.edit_vector_data(
        m=m,
        filename=filename,
        properties=properties,
        time_format=time_format,
        column_widths=column_widths,
        map_height=map_height,
        out_dir=out_dir,
        filename_prefix=filename_prefix,
        file_ext=file_ext,
        add_mapillary=add_mapillary,
        style=style,
        radius=radius,
        width=width,
        height=height,
        frame_border=frame_border,
        controls=controls,
        position=position,
        fit_bounds_options=fit_bounds_options,
        **kwargs,
    )

empty_cache()

Empty the cache of the current device.

Source code in geoai/utils/device.py

def empty_cache() -> None:
    """Empty the cache of the current device."""
    import torch

    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
        torch.mps.empty_cache()

evaluate_sparse_iou(model, images_dir, labels_dir, num_classes, num_channels=3, batch_size=8, background_class=0, ignore_index=False, device=None, verbose=True)

Evaluate a trained model using sparse labels IoU.

This function is designed for incomplete/sparse ground truth where background (0) means "unlabeled" rather than "definitely not this class". Predictions in background areas are NOT penalized as false positives.

Use this for post-training evaluation when your training masks are incomplete.

Parameters:

Name	Type	Description	Default
model	Module	Trained segmentation model	required
images_dir	str	Directory containing validation images	required
labels_dir	str	Directory containing validation labels	required
num_classes	int	Number of classes	required
num_channels	int	Number of input channels (default: 3)	3
batch_size	int	Batch size for evaluation (default: 8)	8
background_class	int	Class ID for background/unlabeled pixels (default: 0)	0
ignore_index	Union[int, bool]	Class to ignore during evaluation. - If int: specific class index to ignore - If False: no class ignored (default)	False
device	Optional[device]	Torch device (auto-detected if None)	None
verbose	bool	Print detailed results (default: True)	True

Returns:

Type	Description
Dict[str, Any]	Dictionary containing:
Dict[str, Any]	'mean_sparse_iou': Mean IoU across all non-background classes
Dict[str, Any]	'per_class_iou': Dict of class_id -> IoU
Dict[str, Any]	'per_class_recall': Dict of class_id -> recall (sensitivity)
Dict[str, Any]	'per_class_precision': Dict of class_id -> precision
Dict[str, Any]	'mean_recall': Mean recall across classes
Dict[str, Any]	'mean_precision': Mean precision across classes

Example

model = torch.load("best_model.pth") results = evaluate_sparse_iou( ... model=model, ... images_dir="tiles/images", ... labels_dir="tiles/labels", ... num_classes=18, ... background_class=0, ... ) print(f"Sparse IoU: {results['mean_sparse_iou']:.4f}")

Source code in geoai/landcover_train.py

def evaluate_sparse_iou(
    model: torch.nn.Module,
    images_dir: str,
    labels_dir: str,
    num_classes: int,
    num_channels: int = 3,
    batch_size: int = 8,
    background_class: int = 0,
    ignore_index: Union[int, bool] = False,
    device: Optional[torch.device] = None,
    verbose: bool = True,
) -> Dict[str, Any]:
    """
    Evaluate a trained model using sparse labels IoU.

    This function is designed for incomplete/sparse ground truth where
    background (0) means "unlabeled" rather than "definitely not this class".
    Predictions in background areas are NOT penalized as false positives.

    Use this for post-training evaluation when your training masks are incomplete.

    Args:
        model: Trained segmentation model
        images_dir: Directory containing validation images
        labels_dir: Directory containing validation labels
        num_classes: Number of classes
        num_channels: Number of input channels (default: 3)
        batch_size: Batch size for evaluation (default: 8)
        background_class: Class ID for background/unlabeled pixels (default: 0)
        ignore_index: Class to ignore during evaluation.
            - If int: specific class index to ignore
            - If False: no class ignored (default)
        device: Torch device (auto-detected if None)
        verbose: Print detailed results (default: True)

    Returns:
        Dictionary containing:
        - 'mean_sparse_iou': Mean IoU across all non-background classes
        - 'per_class_iou': Dict of class_id -> IoU
        - 'per_class_recall': Dict of class_id -> recall (sensitivity)
        - 'per_class_precision': Dict of class_id -> precision
        - 'mean_recall': Mean recall across classes
        - 'mean_precision': Mean precision across classes

    Example:
        >>> model = torch.load("best_model.pth")
        >>> results = evaluate_sparse_iou(
        ...     model=model,
        ...     images_dir="tiles/images",
        ...     labels_dir="tiles/labels",
        ...     num_classes=18,
        ...     background_class=0,
        ... )
        >>> print(f"Sparse IoU: {results['mean_sparse_iou']:.4f}")
    """
    import os
    import rasterio
    from torch.utils.data import DataLoader, Dataset

    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model.to(device)
    model.eval()

    # Get all image and label files
    image_extensions = (".tif", ".tiff", ".png", ".jpg", ".jpeg")
    image_files = sorted(
        [
            os.path.join(images_dir, f)
            for f in os.listdir(images_dir)
            if f.lower().endswith(image_extensions)
        ]
    )
    label_files = sorted(
        [
            os.path.join(labels_dir, f)
            for f in os.listdir(labels_dir)
            if f.lower().endswith(image_extensions)
        ]
    )

    if len(image_files) != len(label_files):
        raise ValueError(
            f"Mismatch: {len(image_files)} images vs {len(label_files)} labels"
        )

    if verbose:
        print(f"\n{'='*60}")
        print("SPARSE LABELS IoU EVALUATION")
        print(f"{'='*60}")
        print(f"Evaluating {len(image_files)} image-label pairs")
        print(f"Background class: {background_class} (predictions here NOT penalized)")
        print(f"Number of classes: {num_classes}")

    # Accumulate predictions and targets
    all_preds = []
    all_targets = []

    with torch.no_grad():
        for i, (img_path, label_path) in enumerate(zip(image_files, label_files)):
            # Load image
            with rasterio.open(img_path) as src:
                image = src.read()[:num_channels]  # (C, H, W)

            # Load label
            with rasterio.open(label_path) as src:
                label = src.read(1)  # (H, W)

            # Normalize image to 0-1 range
            image = image.astype("float32")
            if image.max() > 1.0:
                image = image / 255.0

            # Convert to tensor and add batch dimension
            image_tensor = torch.from_numpy(image).unsqueeze(0).to(device)
            label_tensor = torch.from_numpy(label).unsqueeze(0)

            # Get prediction
            output = model(image_tensor)
            pred = torch.argmax(output, dim=1).cpu()

            all_preds.append(pred)
            all_targets.append(label_tensor)

            if verbose and (i + 1) % 50 == 0:
                print(f"   Processed {i + 1}/{len(image_files)} tiles...")

    # Concatenate all predictions and targets
    all_preds = torch.cat(all_preds, dim=0)
    all_targets = torch.cat(all_targets, dim=0)

    # Calculate sparse IoU
    mean_iou, per_class_ious, recalls, precisions = landcover_iou(
        pred=all_preds,
        target=all_targets,
        num_classes=num_classes,
        ignore_index=ignore_index,
        mode="sparse_labels",
        background_class=background_class,
    )

    # Build results dictionary
    per_class_iou_dict = {}
    per_class_recall_dict = {}
    per_class_precision_dict = {}

    class_idx = 0
    for cls in range(num_classes):
        if cls == background_class:
            continue
        if isinstance(ignore_index, int) and cls == ignore_index:
            continue

        per_class_iou_dict[cls] = per_class_ious[cls]
        per_class_recall_dict[cls] = (
            recalls[class_idx] if class_idx < len(recalls) else 0.0
        )
        per_class_precision_dict[cls] = (
            precisions[class_idx] if class_idx < len(precisions) else 0.0
        )
        class_idx += 1

    # Calculate means (excluding background)
    valid_recalls = [r for r in recalls if r > 0]
    valid_precisions = [p for p in precisions if p > 0]
    mean_recall = sum(valid_recalls) / len(valid_recalls) if valid_recalls else 0.0
    mean_precision = (
        sum(valid_precisions) / len(valid_precisions) if valid_precisions else 0.0
    )

    results = {
        "mean_sparse_iou": mean_iou,
        "per_class_iou": per_class_iou_dict,
        "per_class_recall": per_class_recall_dict,
        "per_class_precision": per_class_precision_dict,
        "mean_recall": mean_recall,
        "mean_precision": mean_precision,
    }

    if verbose:
        print("\nSPARSE LABELS IoU RESULTS:")
        print(f"   (Predictions in background areas NOT counted as false positives)")
        print(f"\n   {'Class':<8} {'IoU':>8} {'Recall':>8} {'Precision':>10}")
        print(f"   {'-'*36}")
        for cls in sorted(per_class_iou_dict.keys()):
            iou = per_class_iou_dict.get(cls, 0.0)
            recall = per_class_recall_dict.get(cls, 0.0)
            precision = per_class_precision_dict.get(cls, 0.0)
            print(f"   {cls:<8} {iou:>8.4f} {recall:>8.4f} {precision:>10.4f}")

        print(f"   {'-'*36}")
        print(
            f"   {'MEAN':<8} {mean_iou:>8.4f} {mean_recall:>8.4f} {mean_precision:>10.4f}"
        )
        print("\nSparse IoU evaluation complete!")
        print(f"{'='*60}\n")

    return results

export_landcover_tiles(in_raster, out_folder, in_class_data=None, tile_size=256, stride=128, class_value_field='class', buffer_radius=0, max_tiles=None, quiet=False, all_touched=True, create_overview=False, skip_empty_tiles=False, min_feature_ratio=False, metadata_format='PASCAL_VOC')

Export GeoTIFF tiles optimized for landcover classification training.

This function extends the base export_geotiff_tiles with enhanced filtering capabilities specifically designed for discrete landcover classification. It can filter out tiles dominated by background pixels to improve training data quality and reduce dataset size.

Parameters:

Name	Type	Description	Default
in_raster	str	Path to input raster (image to tile)	required
out_folder	str	Output directory for tiles	required
in_class_data	Optional[Union[str, GeoDataFrame]]	Path to vector mask or GeoDataFrame (optional for image-only export)	None
tile_size	int	Size of output tiles in pixels (default: 256)	256
stride	int	Stride for sliding window (default: 128)	128
class_value_field	str	Field name containing class values (default: "class")	'class'
buffer_radius	float	Buffer radius around features in pixels (default: 0)	0
max_tiles	Optional[int]	Maximum number of tiles to export (default: None)	None
quiet	bool	Suppress progress output (default: False)	False
all_touched	bool	Include pixels touched by geometry (default: True)	True
create_overview	bool	Create overview image showing tile locations (default: False)	False
skip_empty_tiles	bool	Skip tiles with no features (default: False)	False
min_feature_ratio	Union[bool, float]	Minimum ratio of non-background pixels required to keep tile - False: Disable ratio filtering (default) - 0.0-1.0: Minimum ratio threshold (e.g., 0.1 = 10% features required)	False
metadata_format	str	Annotation format ("PASCAL_VOC" or "YOLO")	'PASCAL_VOC'

Returns:

Type	Description
Dict[str, Any]	Dictionary containing: - tiles_exported: Number of tiles successfully exported - tiles_skipped_empty: Number of completely empty tiles skipped - tiles_skipped_ratio: Number of tiles filtered by min_feature_ratio - output_dirs: Dictionary with paths to images and labels directories

Examples:

Original behavior (no filtering)

export_landcover_tiles( "input.tif", "output", "mask.shp", skip_empty_tiles=True )

Light filtering (keep tiles with ≥5% features)

export_landcover_tiles( "input.tif", "output", "mask.shp", skip_empty_tiles=True, min_feature_ratio=0.05 )

Moderate filtering (keep tiles with ≥15% features)

export_landcover_tiles( "input.tif", "output", "mask.shp", skip_empty_tiles=True, min_feature_ratio=0.15 )

Note

This function is designed for discrete landcover classification where class 0 typically represents background/no data. The min_feature_ratio parameter counts non-zero pixels as "features".

Source code in geoai/landcover_utils.py

def export_landcover_tiles(
    in_raster: str,
    out_folder: str,
    in_class_data: Optional[Union[str, gpd.GeoDataFrame]] = None,
    tile_size: int = 256,
    stride: int = 128,
    class_value_field: str = "class",
    buffer_radius: float = 0,
    max_tiles: Optional[int] = None,
    quiet: bool = False,
    all_touched: bool = True,
    create_overview: bool = False,
    skip_empty_tiles: bool = False,
    min_feature_ratio: Union[bool, float] = False,
    metadata_format: str = "PASCAL_VOC",
) -> Dict[str, Any]:
    """
    Export GeoTIFF tiles optimized for landcover classification training.

    This function extends the base export_geotiff_tiles with enhanced filtering
    capabilities specifically designed for discrete landcover classification.
    It can filter out tiles dominated by background pixels to improve training
    data quality and reduce dataset size.

    Args:
        in_raster: Path to input raster (image to tile)
        out_folder: Output directory for tiles
        in_class_data: Path to vector mask or GeoDataFrame (optional for image-only export)
        tile_size: Size of output tiles in pixels (default: 256)
        stride: Stride for sliding window (default: 128)
        class_value_field: Field name containing class values (default: "class")
        buffer_radius: Buffer radius around features in pixels (default: 0)
        max_tiles: Maximum number of tiles to export (default: None)
        quiet: Suppress progress output (default: False)
        all_touched: Include pixels touched by geometry (default: True)
        create_overview: Create overview image showing tile locations (default: False)
        skip_empty_tiles: Skip tiles with no features (default: False)
        min_feature_ratio: Minimum ratio of non-background pixels required to keep tile
            - False: Disable ratio filtering (default)
            - 0.0-1.0: Minimum ratio threshold (e.g., 0.1 = 10% features required)
        metadata_format: Annotation format ("PASCAL_VOC" or "YOLO")

    Returns:
        Dictionary containing:
            - tiles_exported: Number of tiles successfully exported
            - tiles_skipped_empty: Number of completely empty tiles skipped
            - tiles_skipped_ratio: Number of tiles filtered by min_feature_ratio
            - output_dirs: Dictionary with paths to images and labels directories

    Examples:
        # Original behavior (no filtering)
        export_landcover_tiles(
            "input.tif",
            "output",
            "mask.shp",
            skip_empty_tiles=True
        )

        # Light filtering (keep tiles with ≥5% features)
        export_landcover_tiles(
            "input.tif",
            "output",
            "mask.shp",
            skip_empty_tiles=True,
            min_feature_ratio=0.05
        )

        # Moderate filtering (keep tiles with ≥15% features)
        export_landcover_tiles(
            "input.tif",
            "output",
            "mask.shp",
            skip_empty_tiles=True,
            min_feature_ratio=0.15
        )

    Note:
        This function is designed for discrete landcover classification where
        class 0 typically represents background/no data. The min_feature_ratio
        parameter counts non-zero pixels as "features".
    """

    # Validate min_feature_ratio parameter
    if min_feature_ratio is not False:
        if not isinstance(min_feature_ratio, (int, float)):
            warnings.warn(
                f"min_feature_ratio must be a number between 0.0 and 1.0, got {type(min_feature_ratio)}. "
                "Disabling ratio filtering."
            )
            min_feature_ratio = False
        elif not (0.0 <= min_feature_ratio <= 1.0):
            warnings.warn(
                f"min_feature_ratio must be between 0.0 and 1.0, got {min_feature_ratio}. "
                "Disabling ratio filtering."
            )
            min_feature_ratio = False

    # Create output directories
    out_folder = Path(out_folder)
    out_folder.mkdir(parents=True, exist_ok=True)

    images_dir = out_folder / "images"
    labels_dir = out_folder / "labels"
    images_dir.mkdir(exist_ok=True)
    labels_dir.mkdir(exist_ok=True)

    if metadata_format == "PASCAL_VOC":
        ann_dir = out_folder / "annotations"
        ann_dir.mkdir(exist_ok=True)

    # Initialize statistics
    stats = {
        "tiles_exported": 0,
        "tiles_skipped_empty": 0,
        "tiles_skipped_ratio": 0,
        "output_dirs": {"images": str(images_dir), "labels": str(labels_dir)},
    }

    # Open raster
    with rasterio.open(in_raster) as src:
        height, width = src.shape

        # Detect if in_class_data is raster or vector
        is_class_data_raster = False
        class_src = None
        gdf = None
        mask_array = None

        if in_class_data is not None:
            if isinstance(in_class_data, str):
                file_ext = Path(in_class_data).suffix.lower()
                if file_ext in [
                    ".tif",
                    ".tiff",
                    ".img",
                    ".jp2",
                    ".png",
                    ".bmp",
                    ".gif",
                ]:
                    try:
                        # Try to open as raster
                        class_src = rasterio.open(in_class_data)
                        is_class_data_raster = True

                        # Verify CRS match
                        if class_src.crs != src.crs:
                            if not quiet:
                                print(
                                    f"Warning: CRS mismatch between image ({src.crs}) and mask ({class_src.crs})"
                                )
                    except Exception as e:
                        is_class_data_raster = False
                        if not quiet:
                            print(f"Could not open as raster, trying vector: {e}")

                # If not raster or raster open failed, try vector
                if not is_class_data_raster:
                    gdf = gpd.read_file(in_class_data)

                    # Reproject if needed
                    if gdf.crs != src.crs:
                        if not quiet:
                            print(f"Reprojecting mask from {gdf.crs} to {src.crs}")
                        gdf = gdf.to_crs(src.crs)

                    # Apply buffer if requested
                    if buffer_radius > 0:
                        gdf.geometry = gdf.geometry.buffer(buffer_radius)

                    # For vector data, rasterize entire mask up front for efficiency
                    shapes = [
                        (geom, value)
                        for geom, value in zip(gdf.geometry, gdf[class_value_field])
                    ]
                    mask_array = features.rasterize(
                        shapes,
                        out_shape=(height, width),
                        transform=src.transform,
                        all_touched=all_touched,
                        fill=0,
                        dtype=np.uint8,
                    )
            else:
                # Assume GeoDataFrame passed directly
                gdf = in_class_data

                # Reproject if needed
                if gdf.crs != src.crs:
                    if not quiet:
                        print(f"Reprojecting mask from {gdf.crs} to {src.crs}")
                    gdf = gdf.to_crs(src.crs)

                # Apply buffer if requested
                if buffer_radius > 0:
                    gdf.geometry = gdf.geometry.buffer(buffer_radius)

                # Rasterize entire mask up front
                shapes = [
                    (geom, value)
                    for geom, value in zip(gdf.geometry, gdf[class_value_field])
                ]
                mask_array = features.rasterize(
                    shapes,
                    out_shape=(height, width),
                    transform=src.transform,
                    all_touched=all_touched,
                    fill=0,
                    dtype=np.uint8,
                )

        # Calculate tile positions
        tile_positions = []
        for y in range(0, height - tile_size + 1, stride):
            for x in range(0, width - tile_size + 1, stride):
                tile_positions.append((x, y))

        if max_tiles:
            tile_positions = tile_positions[:max_tiles]

        # Process tiles
        pbar = tqdm(tile_positions, desc="Exporting tiles", disable=quiet)

        for tile_idx, (x, y) in enumerate(pbar):
            window = Window(x, y, tile_size, tile_size)

            # Read image tile
            image_tile = src.read(window=window)

            # Read mask tile based on data type
            mask_tile = None
            has_features = False

            if is_class_data_raster and class_src is not None:
                # For raster masks, read directly from the raster source
                # Get window transform and bounds
                window_transform = src.window_transform(window)
                minx = window_transform[2]
                maxy = window_transform[5]
                maxx = minx + tile_size * window_transform[0]
                miny = maxy + tile_size * window_transform[4]

                # Get corresponding window in class raster
                window_class = rasterio.windows.from_bounds(
                    minx, miny, maxx, maxy, class_src.transform
                )

                try:
                    # Read label data from raster
                    mask_tile = class_src.read(
                        1,
                        window=window_class,
                        boundless=True,
                        out_shape=(tile_size, tile_size),
                    )

                    # Check if tile has features
                    has_features = np.any(mask_tile > 0)
                except Exception as e:
                    if not quiet:
                        pbar.write(f"Error reading mask tile at ({x}, {y}): {e}")
                    continue

            elif mask_array is not None:
                # For vector masks (pre-rasterized)
                mask_tile = mask_array[y : y + tile_size, x : x + tile_size]
                has_features = np.any(mask_tile > 0)

            # Skip empty tiles if requested
            if skip_empty_tiles and not has_features:
                stats["tiles_skipped_empty"] += 1
                continue

            # Apply min_feature_ratio filtering if enabled
            if skip_empty_tiles and has_features and min_feature_ratio is not False:
                # Calculate ratio of non-background pixels
                total_pixels = mask_tile.size
                feature_pixels = np.sum(mask_tile > 0)
                feature_ratio = feature_pixels / total_pixels

                # Skip tile if below threshold
                if feature_ratio < min_feature_ratio:
                    stats["tiles_skipped_ratio"] += 1
                    continue

            # Save image tile
            tile_name = f"tile_{tile_idx:06d}.tif"
            image_path = images_dir / tile_name

            # Get transform for this tile
            tile_transform = src.window_transform(window)

            # Write image
            with rasterio.open(
                image_path,
                "w",
                driver="GTiff",
                height=tile_size,
                width=tile_size,
                count=src.count,
                dtype=src.dtypes[0],
                crs=src.crs,
                transform=tile_transform,
                compress="lzw",
            ) as dst:
                dst.write(image_tile)

            # Save mask tile if available
            if mask_tile is not None:
                mask_path = labels_dir / tile_name
                with rasterio.open(
                    mask_path,
                    "w",
                    driver="GTiff",
                    height=tile_size,
                    width=tile_size,
                    count=1,
                    dtype=np.uint8,
                    crs=src.crs,
                    transform=tile_transform,
                    compress="lzw",
                ) as dst:
                    dst.write(mask_tile, 1)

            stats["tiles_exported"] += 1

            # Update progress bar description with selection count
            if not quiet:
                pbar.set_description(
                    f"Exporting tiles ({stats['tiles_exported']}/{tile_idx + 1})"
                )

    # Close raster class source if opened
    if class_src is not None:
        class_src.close()

    # Print summary
    if not quiet:
        print(f"\n{'='*60}")
        print("TILE EXPORT SUMMARY")
        print(f"{'='*60}")
        print(f"Tiles exported: {stats['tiles_exported']}/{len(tile_positions)}")
        if skip_empty_tiles:
            print(f"Tiles skipped (empty): {stats['tiles_skipped_empty']}")
        if min_feature_ratio is not False:
            print(
                f"Tiles skipped (low feature ratio < {min_feature_ratio}): {stats['tiles_skipped_ratio']}"
            )
        print(f"\nOutput directories:")
        print(f"  Images: {stats['output_dirs']['images']}")
        print(f"  Labels: {stats['output_dirs']['labels']}")
        print(f"{'='*60}\n")

    return stats

export_tiles_to_geojson(tile_coordinates, src, output_path, tile_size=None, stride=None)

Export tile rectangles directly to GeoJSON without creating an overview image.

Parameters:

Name	Type	Description	Default
tile_coordinates	list	A list of dictionaries containing tile information.	required
src	DatasetReader	The source raster dataset.	required
output_path	str	The path where the GeoJSON will be saved.	required
tile_size	int	The size of each tile in pixels. Only needed if not in tile_coordinates.	None
stride	int	The stride between tiles in pixels. Used to calculate overlaps between tiles.	None

Returns:

Name	Type	Description
str	str	Path to the saved GeoJSON file.

Source code in geoai/utils/vector.py

def export_tiles_to_geojson(
    tile_coordinates, src, output_path, tile_size=None, stride=None
) -> str:
    """
    Export tile rectangles directly to GeoJSON without creating an overview image.

    Args:
        tile_coordinates (list): A list of dictionaries containing tile information.
        src (rasterio.io.DatasetReader): The source raster dataset.
        output_path (str): The path where the GeoJSON will be saved.
        tile_size (int, optional): The size of each tile in pixels. Only needed if not in tile_coordinates.
        stride (int, optional): The stride between tiles in pixels. Used to calculate overlaps between tiles.

    Returns:
        str: Path to the saved GeoJSON file.
    """
    features = []

    for tile in tile_coordinates:
        # Get the size from the tile or use the provided parameter
        tile_width = tile.get("width", tile.get("size", tile_size))
        tile_height = tile.get("height", tile.get("size", tile_size))

        if tile_width is None or tile_height is None:
            raise ValueError(
                "Tile size not found in tile data and no tile_size parameter provided"
            )

        # Get bounds from the tile
        if "bounds" in tile:
            # If bounds are already in geo coordinates
            minx, miny, maxx, maxy = tile["bounds"]
        else:
            # Try to calculate bounds from transform if available
            if hasattr(src, "transform"):
                # Convert pixel coordinates to geo coordinates
                window_transform = src.transform
                x, y = tile["x"], tile["y"]
                minx = window_transform[2] + x * window_transform[0]
                maxy = window_transform[5] + y * window_transform[4]
                maxx = minx + tile_width * window_transform[0]
                miny = maxy + tile_height * window_transform[4]
            else:
                raise ValueError(
                    "Cannot determine bounds. Neither 'bounds' in tile nor transform in src."
                )

        # Calculate overlap with neighboring tiles if stride is provided
        overlap = 0
        if stride is not None and stride < tile_width:
            overlap = tile_width - stride

        # Create a polygon from the bounds
        polygon = box(minx, miny, maxx, maxy)

        # Create a GeoJSON feature
        feature = {
            "type": "Feature",
            "geometry": mapping(polygon),
            "properties": {
                "index": tile["index"],
                "has_features": tile.get("has_features", False),
                "tile_width_px": tile_width,
                "tile_height_px": tile_height,
            },
        }

        # Add overlap information if stride is provided
        if stride is not None:
            feature["properties"]["stride_px"] = stride
            feature["properties"]["overlap_px"] = overlap

        # Add additional properties from the tile
        for key, value in tile.items():
            if key not in ["bounds", "geometry"]:
                feature["properties"][key] = value

        features.append(feature)

    # Create the GeoJSON collection
    geojson_collection = {
        "type": "FeatureCollection",
        "features": features,
        "properties": {
            "crs": (
                src.crs.to_string() if hasattr(src.crs, "to_string") else str(src.crs)
            ),
            "total_tiles": len(features),
            "source_raster_dimensions": (
                [src.width, src.height] if hasattr(src, "width") else None
            ),
        },
    }

    # Create directory if it doesn't exist
    os.makedirs(os.path.dirname(os.path.abspath(output_path)) or ".", exist_ok=True)

    # Save to file
    with open(output_path, "w") as f:
        json.dump(geojson_collection, f)

    print(f"GeoJSON saved to {output_path}")
    return output_path

geojson_to_coords(geojson, src_crs='epsg:4326', dst_crs='epsg:4326')

Converts a geojson file or a dictionary of feature collection to a list of centroid coordinates.

Parameters:

Name	Type	Description	Default
geojson	str | dict	The geojson file path or a dictionary of feature collection.	required
src_crs	str	The source CRS. Defaults to "epsg:4326".	'epsg:4326'
dst_crs	str	The destination CRS. Defaults to "epsg:4326".	'epsg:4326'

Returns:

Name	Type	Description
list	list	A list of centroid coordinates in the format of [[x1, y1], [x2, y2], ...]

Source code in geoai/utils/vector.py

def geojson_to_coords(
    geojson: str, src_crs: str = "epsg:4326", dst_crs: str = "epsg:4326"
) -> list:
    """Converts a geojson file or a dictionary of feature collection to a list of centroid coordinates.

    Args:
        geojson (str | dict): The geojson file path or a dictionary of feature collection.
        src_crs (str, optional): The source CRS. Defaults to "epsg:4326".
        dst_crs (str, optional): The destination CRS. Defaults to "epsg:4326".

    Returns:
        list: A list of centroid coordinates in the format of [[x1, y1], [x2, y2], ...]
    """

    import json
    import warnings

    warnings.filterwarnings("ignore")

    if isinstance(geojson, dict):
        geojson = json.dumps(geojson)
    gdf = gpd.read_file(geojson, driver="GeoJSON")
    centroids = gdf.geometry.centroid
    centroid_list = [[point.x, point.y] for point in centroids]
    if src_crs != dst_crs:
        centroid_list = transform_coords(
            [x[0] for x in centroid_list],
            [x[1] for x in centroid_list],
            src_crs,
            dst_crs,
        )
        centroid_list = [[x, y] for x, y in zip(centroid_list[0], centroid_list[1])]
    return centroid_list

geojson_to_xy(src_fp, geojson, coord_crs='epsg:4326', **kwargs)

Converts a geojson file or a dictionary of feature collection to a list of pixel coordinates.

Parameters:

Name	Type	Description	Default
src_fp	str	The source raster file path.	required
geojson	str	The geojson file path or a dictionary of feature collection.	required
coord_crs	str	The coordinate CRS of the input coordinates. Defaults to "epsg:4326".	'epsg:4326'
**kwargs	Any	Additional keyword arguments to pass to rasterio.transform.rowcol.	{}

Returns:

Type	Description
List[List[float]]	A list of pixel coordinates in the format of [[x1, y1], [x2, y2], ...]

Source code in geoai/utils/vector.py

def geojson_to_xy(
    src_fp: str, geojson: str, coord_crs: str = "epsg:4326", **kwargs: Any
) -> List[List[float]]:
    """Converts a geojson file or a dictionary of feature collection to a list of pixel coordinates.

    Args:
        src_fp: The source raster file path.
        geojson: The geojson file path or a dictionary of feature collection.
        coord_crs: The coordinate CRS of the input coordinates. Defaults to "epsg:4326".
        **kwargs: Additional keyword arguments to pass to rasterio.transform.rowcol.

    Returns:
        A list of pixel coordinates in the format of [[x1, y1], [x2, y2], ...]
    """
    with rasterio.open(src_fp) as src:
        src_crs = src.crs
    coords = geojson_to_coords(geojson, coord_crs, src_crs)
    return coords_to_xy(src_fp, coords, src_crs, **kwargs)

get_device()

Returns the best available device for deep learning in the order: CUDA (NVIDIA GPU) > MPS (Apple Silicon GPU) > CPU

Source code in geoai/utils/device.py

def get_device():
    """
    Returns the best available device for deep learning in the order:
    CUDA (NVIDIA GPU) > MPS (Apple Silicon GPU) > CPU
    """
    import torch

    if torch.cuda.is_available():
        return torch.device("cuda")
    elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
        return torch.device("mps")
    else:
        return torch.device("cpu")

get_landcover_loss_function(loss_name='crossentropy', num_classes=2, ignore_index=-100, class_weights=None, use_class_weights=False, focal_alpha=1.0, focal_gamma=2.0, device=None)

Get loss function configured for landcover classification.

Parameters:

Name	Type	Description	Default
loss_name	str	Name of loss function ("crossentropy", "focal", "dice", "combo")	'crossentropy'
num_classes	int	Number of classes	2
ignore_index	Union[int, bool]	Class index to ignore, or False to not ignore any class. - If int: pixels with this label value will be ignored during training - If False: no pixels will be ignored (all pixels contribute to loss)	-100
class_weights	Optional[Tensor]	Manual class weights tensor	None
use_class_weights	bool	Whether to use class weights	False
focal_alpha	float	Alpha parameter for focal loss	1.0
focal_gamma	float	Gamma parameter for focal loss	2.0
device	Optional[device]	Device to place loss function on	None

Returns:

Type	Description
Module	Configured loss function

Source code in geoai/landcover_train.py

def get_landcover_loss_function(
    loss_name: str = "crossentropy",
    num_classes: int = 2,
    ignore_index: Union[int, bool] = -100,
    class_weights: Optional[torch.Tensor] = None,
    use_class_weights: bool = False,
    focal_alpha: float = 1.0,
    focal_gamma: float = 2.0,
    device: Optional[torch.device] = None,
) -> nn.Module:
    """
    Get loss function configured for landcover classification.

    Args:
        loss_name: Name of loss function ("crossentropy", "focal", "dice", "combo")
        num_classes: Number of classes
        ignore_index: Class index to ignore, or False to not ignore any class.
            - If int: pixels with this label value will be ignored during training
            - If False: no pixels will be ignored (all pixels contribute to loss)
        class_weights: Manual class weights tensor
        use_class_weights: Whether to use class weights
        focal_alpha: Alpha parameter for focal loss
        focal_gamma: Gamma parameter for focal loss
        device: Device to place loss function on

    Returns:
        Configured loss function
    """

    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    loss_name = loss_name.lower()

    if loss_name == "crossentropy":
        weights = class_weights if use_class_weights else None
        if weights is not None:
            weights = weights.to(device)

        # Convert ignore_index: int stays as-is, False becomes -100 (PyTorch default)
        idx = ignore_index if isinstance(ignore_index, int) else -100

        return LandcoverCrossEntropyLoss(
            weight=weights,
            ignore_index=idx,
            reduction="mean",
        )

    elif loss_name == "focal":
        weights = class_weights if use_class_weights else None
        if weights is not None:
            weights = weights.to(device)

        # Convert ignore_index: int stays as-is, False becomes -100 (PyTorch default)
        idx = ignore_index if isinstance(ignore_index, int) else -100

        return FocalLoss(
            alpha=focal_alpha,
            gamma=focal_gamma,
            ignore_index=idx,
            reduction="mean",
            weight=weights,
        )

    else:
        # Fall back to standard PyTorch loss
        weights = class_weights if use_class_weights else None
        if weights is not None:
            weights = weights.to(device)

        # Convert ignore_index: int stays as-is, False becomes -100 (PyTorch default)
        idx = ignore_index if isinstance(ignore_index, int) else -100

        return nn.CrossEntropyLoss(
            weight=weights,
            ignore_index=idx,
            reduction="mean",
        )

get_model_config(model_id)

Get the model configuration for a Hugging Face model.

Parameters:

Name	Type	Description	Default
model_id	str	The Hugging Face model ID.	required

Returns:

Type	Description
PretrainedConfig	transformers.configuration_utils.PretrainedConfig: The model configuration.

Source code in geoai/hf.py

def get_model_config(
    model_id: str,
) -> "transformers.configuration_utils.PretrainedConfig":
    """
    Get the model configuration for a Hugging Face model.

    Args:
        model_id (str): The Hugging Face model ID.

    Returns:
        transformers.configuration_utils.PretrainedConfig: The model configuration.
    """
    return AutoConfig.from_pretrained(model_id)

get_model_input_channels(model_id)

Check the number of input channels supported by a Hugging Face model.

Parameters:

Name	Type	Description	Default
model_id	str	The Hugging Face model ID.	required

Returns:

Name	Type	Description
int	int	The number of input channels the model accepts.

Raises:

Type	Description
ValueError	If unable to determine the number of input channels.

Source code in geoai/hf.py

def get_model_input_channels(model_id: str) -> int:
    """
    Check the number of input channels supported by a Hugging Face model.

    Args:
        model_id (str): The Hugging Face model ID.

    Returns:
        int: The number of input channels the model accepts.

    Raises:
        ValueError: If unable to determine the number of input channels.
    """
    # Load the model configuration
    config = AutoConfig.from_pretrained(model_id)

    # For Mask2Former models
    if hasattr(config, "backbone_config"):
        if hasattr(config.backbone_config, "num_channels"):
            return config.backbone_config.num_channels

    # Try to load the model and inspect its architecture
    try:
        model = AutoModelForMaskedImageModeling.from_pretrained(model_id)

        # For Swin Transformer-based models like Mask2Former
        if hasattr(model, "backbone") and hasattr(model.backbone, "embeddings"):
            if hasattr(model.backbone.embeddings, "patch_embeddings"):
                # Swin models typically have patch embeddings that indicate channel count
                return model.backbone.embeddings.patch_embeddings.in_channels
    except Exception as e:
        print(f"Couldn't inspect model architecture: {e}")

    # Default for most vision models
    return 3

get_raster_info(raster_path)

Display basic information about a raster dataset.

Parameters:

Name	Type	Description	Default
raster_path	str	Path to the raster file	required

Returns:

Name	Type	Description
dict	Dict[str, Any]	Dictionary containing the basic information about the raster

Source code in geoai/utils/raster.py

def get_raster_info(raster_path: str) -> Dict[str, Any]:
    """Display basic information about a raster dataset.

    Args:
        raster_path (str): Path to the raster file

    Returns:
        dict: Dictionary containing the basic information about the raster
    """
    # Open the raster dataset
    with rasterio.open(raster_path) as src:
        # Get basic metadata

        info = {
            "driver": src.driver,
            "width": src.width,
            "height": src.height,
            "count": src.count,
            "dtype": src.dtypes[0],
            "crs": src.crs.to_string() if src.crs else "No CRS defined",
            "transform": src.transform,
            "bounds": src.bounds,
            "resolution": (round(src.transform[0], 2), round(-src.transform[4], 2)),
            "nodata": src.nodata,
        }

        # Calculate statistics for each band
        stats = []
        for i in range(1, src.count + 1):
            band = src.read(i, masked=True)
            band_stats = {
                "band": i,
                "min": float(band.min()),
                "max": float(band.max()),
                "mean": float(band.mean()),
                "std": float(band.std()),
            }
            stats.append(band_stats)

        info["band_stats"] = stats

    return info

get_raster_info_gdal(raster_path)

Get basic information about a raster dataset using GDAL.

Parameters:

Name	Type	Description	Default
raster_path	str	Path to the raster file	required

Returns:

Name	Type	Description
dict	Optional[Dict[str, Any]]	Dictionary containing the basic information about the raster, or None if the file cannot be opened

Source code in geoai/utils/raster.py

def get_raster_info_gdal(raster_path: str) -> Optional[Dict[str, Any]]:
    """Get basic information about a raster dataset using GDAL.

    Args:
        raster_path (str): Path to the raster file

    Returns:
        dict: Dictionary containing the basic information about the raster,
            or None if the file cannot be opened
    """

    from osgeo import gdal

    # Open the dataset
    ds = gdal.Open(raster_path)
    if ds is None:
        print(f"Error: Could not open {raster_path}")
        return None

    # Get basic information
    info = {
        "driver": ds.GetDriver().ShortName,
        "width": ds.RasterXSize,
        "height": ds.RasterYSize,
        "count": ds.RasterCount,
        "projection": ds.GetProjection(),
        "geotransform": ds.GetGeoTransform(),
    }

    # Calculate resolution
    gt = ds.GetGeoTransform()
    if gt:
        info["resolution"] = (abs(gt[1]), abs(gt[5]))
        info["origin"] = (gt[0], gt[3])

    # Get band information
    bands_info = []
    for i in range(1, ds.RasterCount + 1):
        band = ds.GetRasterBand(i)
        stats = band.GetStatistics(True, True)
        band_info = {
            "band": i,
            "datatype": gdal.GetDataTypeName(band.DataType),
            "min": stats[0],
            "max": stats[1],
            "mean": stats[2],
            "std": stats[3],
            "nodata": band.GetNoDataValue(),
        }
        bands_info.append(band_info)

    info["bands"] = bands_info

    # Close the dataset
    ds = None

    return info

get_raster_resolution(image_path)

Get pixel resolution from the raster using rasterio.

Parameters:

Name	Type	Description	Default
image_path	str	The path to the raster image.	required

Returns:

Type	Description
Tuple[float, float]	A tuple of (x resolution, y resolution).

Source code in geoai/utils/raster.py

def get_raster_resolution(image_path: str) -> Tuple[float, float]:
    """Get pixel resolution from the raster using rasterio.

    Args:
        image_path: The path to the raster image.

    Returns:
        A tuple of (x resolution, y resolution).
    """
    with rasterio.open(image_path) as src:
        res = src.res
    return res

get_raster_stats(raster_path, divide_by=1.0)

Calculate statistics for each band in a raster dataset.

This function computes min, max, mean, and standard deviation values for each band in the provided raster, returning results in a dictionary with lists for each statistic type.

Parameters:

Name	Type	Description	Default
raster_path	str	Path to the raster file	required
divide_by	float	Value to divide pixel values by. Defaults to 1.0, which keeps the original pixel	1.0

Returns:

Name	Type	Description
dict	Dict[str, Any]	Dictionary containing lists of statistics with keys: - 'min': List of minimum values for each band - 'max': List of maximum values for each band - 'mean': List of mean values for each band - 'std': List of standard deviation values for each band

Source code in geoai/utils/raster.py

def get_raster_stats(raster_path: str, divide_by: float = 1.0) -> Dict[str, Any]:
    """Calculate statistics for each band in a raster dataset.

    This function computes min, max, mean, and standard deviation values
    for each band in the provided raster, returning results in a dictionary
    with lists for each statistic type.

    Args:
        raster_path (str): Path to the raster file
        divide_by (float, optional): Value to divide pixel values by.
            Defaults to 1.0, which keeps the original pixel

    Returns:
        dict: Dictionary containing lists of statistics with keys:
            - 'min': List of minimum values for each band
            - 'max': List of maximum values for each band
            - 'mean': List of mean values for each band
            - 'std': List of standard deviation values for each band
    """
    # Initialize the results dictionary with empty lists
    stats = {"min": [], "max": [], "mean": [], "std": []}

    # Open the raster dataset
    with rasterio.open(raster_path) as src:
        # Calculate statistics for each band
        for i in range(1, src.count + 1):
            band = src.read(i, masked=True)

            # Append statistics for this band to each list
            stats["min"].append(float(band.min()) / divide_by)
            stats["max"].append(float(band.max()) / divide_by)
            stats["mean"].append(float(band.mean()) / divide_by)
            stats["std"].append(float(band.std()) / divide_by)

    return stats

get_vector_info(vector_path)

Display basic information about a vector dataset using GeoPandas.

Parameters:

Name	Type	Description	Default
vector_path	str	Path to the vector file	required

Returns:

Name	Type	Description
dict	Optional[Dict[str, Any]]	Dictionary containing the basic information about the vector dataset

Source code in geoai/utils/vector.py

def get_vector_info(vector_path: str) -> Optional[Dict[str, Any]]:
    """Display basic information about a vector dataset using GeoPandas.

    Args:
        vector_path (str): Path to the vector file

    Returns:
        dict: Dictionary containing the basic information about the vector dataset
    """
    # Open the vector dataset
    gdf = (
        gpd.read_parquet(vector_path)
        if vector_path.endswith(".parquet")
        else gpd.read_file(vector_path)
    )

    # Get basic metadata
    info = {
        "file_path": vector_path,
        "driver": os.path.splitext(vector_path)[1][1:].upper(),  # Format from extension
        "feature_count": len(gdf),
        "crs": str(gdf.crs),
        "geometry_type": str(gdf.geom_type.value_counts().to_dict()),
        "attribute_count": len(gdf.columns) - 1,  # Subtract the geometry column
        "attribute_names": list(gdf.columns[gdf.columns != "geometry"]),
        "bounds": gdf.total_bounds.tolist(),
    }

    # Add statistics about numeric attributes
    numeric_columns = gdf.select_dtypes(include=["number"]).columns
    attribute_stats = {}
    for col in numeric_columns:
        if col != "geometry":
            attribute_stats[col] = {
                "min": gdf[col].min(),
                "max": gdf[col].max(),
                "mean": gdf[col].mean(),
                "std": gdf[col].std(),
                "null_count": gdf[col].isna().sum(),
            }

    info["attribute_stats"] = attribute_stats

    return info

get_vector_info_ogr(vector_path)

Get basic information about a vector dataset using OGR.

Parameters:

Name	Type	Description	Default
vector_path	str	Path to the vector file	required

Returns:

Name	Type	Description
dict	Optional[Dict[str, Any]]	Dictionary containing the basic information about the vector dataset, or None if the file cannot be opened

Source code in geoai/utils/vector.py

def get_vector_info_ogr(vector_path: str) -> Optional[Dict[str, Any]]:
    """Get basic information about a vector dataset using OGR.

    Args:
        vector_path (str): Path to the vector file

    Returns:
        dict: Dictionary containing the basic information about the vector dataset,
            or None if the file cannot be opened
    """
    from osgeo import ogr

    # Register all OGR drivers
    ogr.RegisterAll()

    # Open the dataset
    ds = ogr.Open(vector_path)
    if ds is None:
        print(f"Error: Could not open {vector_path}")
        return None

    # Basic dataset information
    info = {
        "file_path": vector_path,
        "driver": ds.GetDriver().GetName(),
        "layer_count": ds.GetLayerCount(),
        "layers": [],
    }

    # Extract information for each layer
    for i in range(ds.GetLayerCount()):
        layer = ds.GetLayer(i)
        layer_info = {
            "name": layer.GetName(),
            "feature_count": layer.GetFeatureCount(),
            "geometry_type": ogr.GeometryTypeToName(layer.GetGeomType()),
            "spatial_ref": (
                layer.GetSpatialRef().ExportToWkt() if layer.GetSpatialRef() else "None"
            ),
            "extent": layer.GetExtent(),
            "fields": [],
        }

        # Get field information
        defn = layer.GetLayerDefn()
        for j in range(defn.GetFieldCount()):
            field_defn = defn.GetFieldDefn(j)
            field_info = {
                "name": field_defn.GetName(),
                "type": field_defn.GetTypeName(),
                "width": field_defn.GetWidth(),
                "precision": field_defn.GetPrecision(),
            }
            layer_info["fields"].append(field_info)

        info["layers"].append(layer_info)

    # Close the dataset
    ds = None

    return info

hybrid_regularization(building_polygons)

A comprehensive hybrid approach to building footprint regularization.

Applies different strategies based on building characteristics.

Parameters:

Name	Type	Description	Default
building_polygons	Union[GeoDataFrame, List[Polygon]]	GeoDataFrame or list of shapely Polygons containing building footprints	required

Returns:

Type	Description
Union[GeoDataFrame, List[Polygon]]	GeoDataFrame or list of shapely Polygons with regularized building footprints

Source code in geoai/utils/geometry.py

def hybrid_regularization(
    building_polygons: Union[gpd.GeoDataFrame, List[Polygon]],
) -> Union[gpd.GeoDataFrame, List[Polygon]]:
    """
    A comprehensive hybrid approach to building footprint regularization.

    Applies different strategies based on building characteristics.

    Args:
        building_polygons: GeoDataFrame or list of shapely Polygons containing building footprints

    Returns:
        GeoDataFrame or list of shapely Polygons with regularized building footprints
    """
    from shapely.affinity import rotate
    from shapely.geometry import Polygon

    # Use minimum_rotated_rectangle instead of oriented_envelope
    try:
        from shapely.minimum_rotated_rectangle import minimum_rotated_rectangle
    except ImportError:
        # For older Shapely versions
        def minimum_rotated_rectangle(geom):
            """Calculate the minimum rotated rectangle for a geometry"""
            # For older Shapely versions, implement a simple version
            return geom.minimum_rotated_rectangle

    # Determine input type for correct return
    is_gdf = isinstance(building_polygons, gpd.GeoDataFrame)

    # Extract geometries if GeoDataFrame
    if is_gdf:
        geom_objects = building_polygons.geometry
    else:
        geom_objects = building_polygons

    results = []

    for building in geom_objects:
        # 1. Analyze building characteristics
        if not hasattr(building, "exterior") or building.is_empty:
            results.append(building)
            continue

        # Calculate shape complexity metrics
        complexity = building.length / (4 * np.sqrt(building.area))

        # Calculate dominant angle
        coords = np.array(building.exterior.coords)[:-1]
        segments = np.diff(np.vstack([coords, coords[0]]), axis=0)
        segment_lengths = np.sqrt(segments[:, 0] ** 2 + segments[:, 1] ** 2)
        segment_angles = np.arctan2(segments[:, 1], segments[:, 0]) * 180 / np.pi

        # Weight angles by segment length
        hist, bins = np.histogram(
            segment_angles % 180, bins=36, range=(0, 180), weights=segment_lengths
        )
        bin_centers = (bins[:-1] + bins[1:]) / 2
        dominant_angle = bin_centers[np.argmax(hist)]

        # Check if building is close to orthogonal
        is_orthogonal = min(dominant_angle % 45, 45 - (dominant_angle % 45)) < 5

        # 2. Apply appropriate regularization strategy
        if complexity > 1.5:
            # Complex buildings: use minimum rotated rectangle
            result = minimum_rotated_rectangle(building)
        elif is_orthogonal:
            # Near-orthogonal buildings: orthogonalize in place
            rotated = rotate(building, -dominant_angle, origin="centroid")

            # Create orthogonal hull in rotated space
            bounds = rotated.bounds
            ortho_hull = Polygon(
                [
                    (bounds[0], bounds[1]),
                    (bounds[2], bounds[1]),
                    (bounds[2], bounds[3]),
                    (bounds[0], bounds[3]),
                ]
            )

            result = rotate(ortho_hull, dominant_angle, origin="centroid")
        else:
            # Diagonal buildings: use custom approach for diagonal buildings
            # Rotate to align with axes
            rotated = rotate(building, -dominant_angle, origin="centroid")

            # Simplify in rotated space
            simplified = rotated.simplify(0.3, preserve_topology=True)

            # Get the bounds in rotated space
            bounds = simplified.bounds
            min_x, min_y, max_x, max_y = bounds

            # Create a rectangular hull in rotated space
            rect_poly = Polygon(
                [(min_x, min_y), (max_x, min_y), (max_x, max_y), (min_x, max_y)]
            )

            # Rotate back to original orientation
            result = rotate(rect_poly, dominant_angle, origin="centroid")

        results.append(result)

    # Return in same format as input
    if is_gdf:
        return gpd.GeoDataFrame(geometry=results, crs=building_polygons.crs)
    else:
        return results

image_segmentation(tif_path, output_path, labels_to_extract=None, dtype='uint8', model_name=None, segmenter_args=None, **kwargs)

Segments an image with a Hugging Face segmentation model and saves the results as a single georeferenced image where each class has a unique integer value.

Parameters:

Name	Type	Description	Default
tif_path	str	Path to the input georeferenced TIF file.	required
output_path	str	Path where the output georeferenced segmentation will be saved.	required
labels_to_extract	list	List of labels to extract. If None, extracts all labels.	None
dtype	str	Data type to use for the output mask. Defaults to "uint8".	'uint8'
model_name	str	Name of the Hugging Face model to use for segmentation, such as "facebook/mask2former-swin-large-cityscapes-semantic". Defaults to None. See https://huggingface.co/models?pipeline_tag=image-segmentation&sort=trending for options.	None
segmenter_args	dict	Additional arguments to pass to the segmenter. Defaults to None.	None
**kwargs	Any	Additional keyword arguments to pass to the segmentation pipeline	{}

Returns:

Name	Type	Description
tuple	str	(Path to saved image, dictionary mapping label names to their assigned values, dictionary mapping label names to confidence scores)

Source code in geoai/hf.py

def image_segmentation(
    tif_path: str,
    output_path: str,
    labels_to_extract: Optional[List[str]] = None,
    dtype: str = "uint8",
    model_name: Optional[str] = None,
    segmenter_args: Optional[Dict] = None,
    **kwargs: Any,
) -> str:
    """
    Segments an image with a Hugging Face segmentation model and saves the results
    as a single georeferenced image where each class has a unique integer value.

    Args:
        tif_path (str): Path to the input georeferenced TIF file.
        output_path (str): Path where the output georeferenced segmentation will be saved.
        labels_to_extract (list, optional): List of labels to extract. If None, extracts all labels.
        dtype (str, optional): Data type to use for the output mask. Defaults to "uint8".
        model_name (str, optional): Name of the Hugging Face model to use for segmentation,
            such as "facebook/mask2former-swin-large-cityscapes-semantic". Defaults to None.
            See https://huggingface.co/models?pipeline_tag=image-segmentation&sort=trending for options.
        segmenter_args (dict, optional): Additional arguments to pass to the segmenter.
            Defaults to None.
        **kwargs: Additional keyword arguments to pass to the segmentation pipeline

    Returns:
        tuple: (Path to saved image, dictionary mapping label names to their assigned values,
            dictionary mapping label names to confidence scores)
    """
    # Load the original georeferenced image to extract metadata
    with rasterio.open(tif_path) as src:
        # Save the metadata for later use
        meta = src.meta.copy()
        # Get the dimensions
        height = src.height
        width = src.width
        # Get the transform and CRS for georeferencing
        # transform = src.transform
        # crs = src.crs

    # Initialize the segmentation pipeline
    if model_name is None:
        model_name = "facebook/mask2former-swin-large-cityscapes-semantic"

    kwargs["task"] = "image-segmentation"

    segmenter = pipeline(model=model_name, **kwargs)

    # Run the segmentation on the GeoTIFF
    if segmenter_args is None:
        segmenter_args = {}

    segments = segmenter(tif_path, **segmenter_args)

    # If no specific labels are requested, extract all available ones
    if labels_to_extract is None:
        labels_to_extract = [segment["label"] for segment in segments]

    # Create an empty mask to hold all the labels
    # Using uint8 for up to 255 classes, switch to uint16 for more
    combined_mask = np.zeros((height, width), dtype=np.uint8)

    # Create a dictionary to map labels to values and store scores
    label_to_value = {}
    label_to_score = {}

    # Process each segment we want to keep
    for i, segment in enumerate(
        [s for s in segments if s["label"] in labels_to_extract]
    ):
        # Assign a unique value to each label (starting from 1)
        value = i + 1
        label = segment["label"]
        score = segment["score"]

        label_to_value[label] = value
        label_to_score[label] = score

        # Convert PIL image to numpy array
        mask = np.array(segment["mask"])

        # Apply a threshold if it's a probability mask (not binary)
        if mask.dtype == float:
            mask = (mask > 0.5).astype(np.uint8)

        # Resize if needed to match original dimensions
        if mask.shape != (height, width):
            mask_img = Image.fromarray(mask)
            mask_img = mask_img.resize((width, height))
            mask = np.array(mask_img)

        # Add this class to the combined mask
        # Only overwrite if the pixel isn't already assigned to another class
        # This handles overlapping segments by giving priority to earlier segments
        combined_mask = np.where(
            (mask > 0) & (combined_mask == 0), value, combined_mask
        )

    # Update metadata for the output raster
    meta.update(
        {
            "count": 1,  # One band for the mask
            "dtype": dtype,  # Use uint8 for up to 255 classes
            "nodata": 0,  # 0 represents no class
        }
    )

    # Save the mask as a new georeferenced GeoTIFF
    with rasterio.open(output_path, "w", **meta) as dst:
        dst.write(combined_mask[np.newaxis, :, :])  # Add channel dimension

    # Create a CSV colormap file with scores included
    csv_path = os.path.splitext(output_path)[0] + "_colormap.csv"
    with open(csv_path, "w", newline="") as csvfile:
        fieldnames = ["ClassValue", "ClassName", "ConfidenceScore"]
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

        writer.writeheader()
        for label, value in label_to_value.items():
            writer.writerow(
                {
                    "ClassValue": value,
                    "ClassName": label,
                    "ConfidenceScore": f"{label_to_score[label]:.4f}",
                }
            )

    return output_path, label_to_value, label_to_score

install_package(package)

Install a Python package.

Parameters:

Name	Type	Description	Default
package	str | list	The package name or a GitHub URL or a list of package names or GitHub URLs.	required

Source code in geoai/utils/device.py

def install_package(package: Union[str, List[str]]) -> None:
    """Install a Python package.

    Args:
        package (str | list): The package name or a GitHub URL or a list of package names or GitHub URLs.
    """
    import subprocess

    if isinstance(package, str):
        packages = [package]
    elif isinstance(package, list):
        packages = package
    else:
        raise ValueError("The package argument must be a string or a list of strings.")

    for package in packages:
        if package.startswith("https"):
            package = f"git+{package}"

        # Execute pip install command and show output in real-time
        command = f"pip install {package}"
        process = subprocess.Popen(command.split(), stdout=subprocess.PIPE)

        # Print output in real-time
        while True:
            output = process.stdout.readline()
            if output == b"" and process.poll() is not None:
                break
            if output:
                print(output.decode("utf-8").strip())

        # Wait for process to complete
        process.wait()

landcover_iou(pred, target, num_classes, ignore_index=False, smooth=1e-06, mode='mean', boundary_weight_map=None, background_class=None)

Calculate IoU for landcover classification with multiple weighting options.

Supports four IoU calculation modes: 1. "mean": Simple mean IoU across all classes 2. "perclass_frequency": Weight by per-class pixel frequency 3. "boundary_weighted": Weight by distance to class boundaries 4. "sparse_labels": For incomplete ground truth - only penalize FP where GT is positive (does NOT penalize predictions in unlabeled/background areas)

Parameters:

Name	Type	Description	Default
pred	Tensor	Predicted classes (N, H, W) or logits (N, C, H, W)	required
target	Tensor	Ground truth (N, H, W)	required
num_classes	int	Number of classes	required
ignore_index	Union[int, bool]	Class index to ignore (default: None)	False
smooth	float	Smoothing factor to avoid division by zero	1e-06
mode	str	IoU calculation mode ("mean", "perclass_frequency", "boundary_weighted", "sparse_labels")	'mean'
boundary_weight_map	Optional[Tensor]	Optional boundary weights (N, H, W)	None
background_class	Optional[int]	Background/unlabeled class for sparse_labels mode (default: 0)	None

Returns:

Type	Description
Union[float, Tuple[float, List[float], List[int]]]	If mode == "mean": float (mean IoU)
Union[float, Tuple[float, List[float], List[int]]]	If mode == "perclass_frequency": tuple (weighted IoU, per-class IoUs, class counts)
Union[float, Tuple[float, List[float], List[int]]]	If mode == "boundary_weighted": float (boundary-weighted IoU)
Union[float, Tuple[float, List[float], List[int]]]	If mode == "sparse_labels": tuple (sparse IoU, per-class IoUs, per-class recall, per-class precision)

Source code in geoai/landcover_train.py

def landcover_iou(
    pred: torch.Tensor,
    target: torch.Tensor,
    num_classes: int,
    ignore_index: Union[int, bool] = False,
    smooth: float = 1e-6,
    mode: str = "mean",
    boundary_weight_map: Optional[torch.Tensor] = None,
    background_class: Optional[int] = None,
) -> Union[float, Tuple[float, List[float], List[int]]]:
    """
    Calculate IoU for landcover classification with multiple weighting options.

    Supports four IoU calculation modes:
    1. "mean": Simple mean IoU across all classes
    2. "perclass_frequency": Weight by per-class pixel frequency
    3. "boundary_weighted": Weight by distance to class boundaries
    4. "sparse_labels": For incomplete ground truth - only penalize FP where GT is positive
       (does NOT penalize predictions in unlabeled/background areas)

    Args:
        pred: Predicted classes (N, H, W) or logits (N, C, H, W)
        target: Ground truth (N, H, W)
        num_classes: Number of classes
        ignore_index: Class index to ignore (default: None)
        smooth: Smoothing factor to avoid division by zero
        mode: IoU calculation mode ("mean", "perclass_frequency", "boundary_weighted", "sparse_labels")
        boundary_weight_map: Optional boundary weights (N, H, W)
        background_class: Background/unlabeled class for sparse_labels mode (default: 0)

    Returns:
        If mode == "mean": float (mean IoU)
        If mode == "perclass_frequency": tuple (weighted IoU, per-class IoUs, class counts)
        If mode == "boundary_weighted": float (boundary-weighted IoU)
        If mode == "sparse_labels": tuple (sparse IoU, per-class IoUs, per-class recall, per-class precision)
    """

    # Convert logits to class predictions if needed
    if pred.dim() == 4:
        pred = torch.argmax(pred, dim=1)

    # Ensure correct shape
    if pred.shape != target.shape:
        raise ValueError(f"Shape mismatch: pred {pred.shape}, target {target.shape}")

    # Create mask for valid pixels
    # Handle ignore_index: int means specific class, False means don't ignore
    if isinstance(ignore_index, int):
        valid_mask = target != ignore_index
    else:
        # ignore_index is False or any other non-int value - don't ignore anything
        valid_mask = torch.ones_like(target, dtype=torch.bool)

    # Simple mean IoU
    if mode == "mean":
        ious = []
        for cls in range(num_classes):
            if isinstance(ignore_index, int) and cls == ignore_index:
                continue

            pred_cls = (pred == cls) & valid_mask
            target_cls = (target == cls) & valid_mask

            intersection = (pred_cls & target_cls).sum().float()
            union = (pred_cls | target_cls).sum().float()

            if union > 0:
                iou = (intersection + smooth) / (union + smooth)
                ious.append(iou.item())

        return sum(ious) / len(ious) if ious else 0.0

    # Per-class frequency weighted IoU
    elif mode == "perclass_frequency":
        ious = []
        class_counts = []

        # Filter out ignore_index from target
        if isinstance(ignore_index, int):
            target_filtered = target[valid_mask]
            pred_filtered = pred[valid_mask]
        else:
            target_filtered = target.view(-1)
            pred_filtered = pred.view(-1)

        total_valid_pixels = target_filtered.numel()

        for cls in range(num_classes):
            if isinstance(ignore_index, int) and cls == ignore_index:
                continue

            pred_cls = pred_filtered == cls
            target_cls = target_filtered == cls

            intersection = (pred_cls & target_cls).sum().float()
            union = (pred_cls | target_cls).sum().float()

            class_pixel_count = target_cls.sum().item()

            if union > 0:
                iou = (intersection + smooth) / (union + smooth)
                ious.append(iou.item())
                class_counts.append(class_pixel_count)
            else:
                ious.append(0.0)
                class_counts.append(0)

        # Calculate frequency-weighted IoU
        if sum(class_counts) > 0:
            weights = [count / total_valid_pixels for count in class_counts]
            weighted_iou = sum(iou * weight for iou, weight in zip(ious, weights))
        else:
            weighted_iou = 0.0

        return weighted_iou, ious, class_counts

    # Boundary-weighted IoU
    elif mode == "boundary_weighted":
        if boundary_weight_map is None:
            raise ValueError("boundary_weight_map required for boundary_weighted mode")

        ious = []
        weights = []

        for cls in range(num_classes):
            if isinstance(ignore_index, int) and cls == ignore_index:
                continue

            pred_cls = (pred == cls) & valid_mask
            target_cls = (target == cls) & valid_mask

            # Weight by boundary map
            weighted_intersection = (
                pred_cls & target_cls
            ).float() * boundary_weight_map
            weighted_union = (pred_cls | target_cls).float() * boundary_weight_map

            intersection_sum = weighted_intersection.sum()
            union_sum = weighted_union.sum()

            if union_sum > 0:
                iou = (intersection_sum + smooth) / (union_sum + smooth)
                weight = union_sum.item()
                ious.append(iou.item())
                weights.append(weight)

        if sum(weights) > 0:
            weighted_iou = sum(iou * w for iou, w in zip(ious, weights)) / sum(weights)
        else:
            weighted_iou = 0.0

        return weighted_iou

    # Sparse labels IoU - for incomplete ground truth
    # Key insight: background (0) means "unlabeled", not "definitely not this class"
    # So we DON'T penalize predictions in background areas
    elif mode == "sparse_labels":
        # Default background class is 0 if not specified
        bg_class = background_class if background_class is not None else 0

        ious = []
        recalls = []
        precisions = []
        per_class_ious = []

        # Mask for labeled pixels (ground truth is NOT background)
        labeled_mask = target != bg_class
        if isinstance(ignore_index, int):
            labeled_mask = labeled_mask & (target != ignore_index)

        for cls in range(num_classes):
            # Skip background class and ignore_index
            if cls == bg_class:
                per_class_ious.append(0.0)
                recalls.append(0.0)
                precisions.append(0.0)
                continue
            if isinstance(ignore_index, int) and cls == ignore_index:
                per_class_ious.append(0.0)
                recalls.append(0.0)
                precisions.append(0.0)
                continue

            # Where prediction says this class
            pred_cls = pred == cls
            # Where ground truth says this class
            target_cls = target == cls

            # TRUE POSITIVE: Prediction matches target (both say this class)
            tp = (pred_cls & target_cls).sum().float()

            # FALSE NEGATIVE: Target says this class but prediction doesn't
            fn = (target_cls & ~pred_cls).sum().float()

            # FALSE POSITIVE (SPARSE VERSION):
            # Prediction says this class, target says DIFFERENT class (but NOT background)
            # Key: We don't count predictions in background as FP!
            fp_sparse = (pred_cls & ~target_cls & labeled_mask).sum().float()

            # Standard IoU but with sparse FP definition
            # Union = TP + FN + FP_sparse
            union_sparse = tp + fn + fp_sparse

            if union_sparse > 0:
                iou = (tp + smooth) / (union_sparse + smooth)
                ious.append(iou.item())
                per_class_ious.append(iou.item())
            else:
                per_class_ious.append(0.0)

            # Also compute recall and precision for diagnostic purposes
            # Recall: Of all true positives, how many did we find?
            if (tp + fn) > 0:
                recall = tp / (tp + fn)
                recalls.append(recall.item())
            else:
                recalls.append(0.0)

            # Precision (sparse): Of predictions in labeled areas, how many are correct?
            if (tp + fp_sparse) > 0:
                precision = tp / (tp + fp_sparse)
                precisions.append(precision.item())
            else:
                precisions.append(0.0)

        # Mean IoU across classes (excluding background)
        mean_sparse_iou = sum(ious) / len(ious) if ious else 0.0

        return mean_sparse_iou, per_class_ious, recalls, precisions

    else:
        raise ValueError(
            f"Unknown mode: {mode}. Use 'mean', 'perclass_frequency', 'boundary_weighted', or 'sparse_labels'"
        )

mask_generation(input_path, output_mask_path, output_csv_path, model='facebook/sam-vit-base', confidence_threshold=0.5, points_per_side=32, crop_size=None, batch_size=1, band_indices=None, min_object_size=0, generator_kwargs=None, **kwargs)

Process a GeoTIFF using SAM mask generation and save results as a GeoTIFF and CSV.

The function reads a GeoTIFF image, applies the SAM mask generator from the Hugging Face transformers pipeline, rasterizes the resulting masks to create a labeled mask GeoTIFF, and saves mask scores and geometries to a CSV file.

Parameters:

Name	Type	Description	Default
input_path	str	Path to the input GeoTIFF image.	required
output_mask_path	str	Path where the output mask GeoTIFF will be saved.	required
output_csv_path	str	Path where the mask scores CSV will be saved.	required
model	str	HuggingFace model checkpoint for the SAM model.	'facebook/sam-vit-base'
confidence_threshold	float	Minimum confidence score for masks to be included.	0.5
points_per_side	int	Number of points to sample along each side of the image.	32
crop_size	Optional[int]	Size of image crops for processing. If None, process the full image.	None
band_indices	Optional[List[int]]	List of band indices to use. If None, use all bands.	None
batch_size	int	Batch size for inference.	1
min_object_size	int	Minimum size in pixels for objects to be included. Smaller masks will be filtered out.	0
generator_kwargs	Optional[Dict]	Additional keyword arguments to pass to the mask generator.	None

Returns:

Type	Description
Tuple[str, str]	Tuple containing the paths to the saved mask GeoTIFF and CSV file.

Raises:

Type	Description
ValueError	If the input file cannot be opened or processed.
RuntimeError	If mask generation fails.

Source code in geoai/hf.py

def mask_generation(
    input_path: str,
    output_mask_path: str,
    output_csv_path: str,
    model: str = "facebook/sam-vit-base",
    confidence_threshold: float = 0.5,
    points_per_side: int = 32,
    crop_size: Optional[int] = None,
    batch_size: int = 1,
    band_indices: Optional[List[int]] = None,
    min_object_size: int = 0,
    generator_kwargs: Optional[Dict] = None,
    **kwargs: Any,
) -> Tuple[str, str]:
    """
    Process a GeoTIFF using SAM mask generation and save results as a GeoTIFF and CSV.

    The function reads a GeoTIFF image, applies the SAM mask generator from the
    Hugging Face transformers pipeline, rasterizes the resulting masks to create
    a labeled mask GeoTIFF, and saves mask scores and geometries to a CSV file.

    Args:
        input_path: Path to the input GeoTIFF image.
        output_mask_path: Path where the output mask GeoTIFF will be saved.
        output_csv_path: Path where the mask scores CSV will be saved.
        model: HuggingFace model checkpoint for the SAM model.
        confidence_threshold: Minimum confidence score for masks to be included.
        points_per_side: Number of points to sample along each side of the image.
        crop_size: Size of image crops for processing. If None, process the full image.
        band_indices: List of band indices to use. If None, use all bands.
        batch_size: Batch size for inference.
        min_object_size: Minimum size in pixels for objects to be included. Smaller masks will be filtered out.
        generator_kwargs: Additional keyword arguments to pass to the mask generator.

    Returns:
        Tuple containing the paths to the saved mask GeoTIFF and CSV file.

    Raises:
        ValueError: If the input file cannot be opened or processed.
        RuntimeError: If mask generation fails.
    """
    # Set up the mask generator
    print("Setting up mask generator...")
    mask_generator = pipeline(model=model, task="mask-generation", **kwargs)

    # Open the GeoTIFF file
    try:
        print(f"Reading input GeoTIFF: {input_path}")
        with rasterio.open(input_path) as src:
            # Read metadata
            profile = src.profile
            # transform = src.transform
            # crs = src.crs

            # Read the image data
            if band_indices is not None:
                print(f"Using specified bands: {band_indices}")
                image_data = np.stack([src.read(i + 1) for i in band_indices])
            else:
                print("Using all bands")
                image_data = src.read()

            # Handle image with more than 3 bands (convert to RGB for visualization)
            if image_data.shape[0] > 3:
                print(
                    f"Converting {image_data.shape[0]} bands to RGB (using first 3 bands)"
                )
                # Select first three bands or perform other band combination
                image_data = image_data[:3]
            elif image_data.shape[0] == 1:
                print("Duplicating single band to create 3-band image")
                # Duplicate single band to create a 3-band image
                image_data = np.vstack([image_data] * 3)

            # Transpose to HWC format for the model
            image_data = np.transpose(image_data, (1, 2, 0))

            # Normalize the image if needed
            if image_data.dtype != np.uint8:
                print(f"Normalizing image from {image_data.dtype} to uint8")
                image_data = (image_data / image_data.max() * 255).astype(np.uint8)
    except Exception as e:
        raise ValueError(f"Failed to open or process input GeoTIFF: {e}")

    # Process the image with the mask generator
    try:
        # Convert numpy array to PIL Image for the pipeline
        # Ensure the array is in the right format (HWC and uint8)
        if image_data.dtype != np.uint8:
            image_data = (image_data / image_data.max() * 255).astype(np.uint8)

        # Create a PIL Image from the numpy array
        print("Converting to PIL Image for mask generation")
        pil_image = Image.fromarray(image_data)

        # Use the SAM pipeline for mask generation
        if generator_kwargs is None:
            generator_kwargs = {}

        print("Running mask generation...")
        mask_results = mask_generator(
            pil_image,
            points_per_side=points_per_side,
            crop_n_points_downscale_factor=1 if crop_size is None else 2,
            point_grids=None,
            pred_iou_thresh=confidence_threshold,
            stability_score_thresh=confidence_threshold,
            crops_n_layers=0 if crop_size is None else 1,
            crop_overlap_ratio=0.5,
            batch_size=batch_size,
            **generator_kwargs,
        )

        print(
            f"Number of initial masks: {len(mask_results['masks']) if isinstance(mask_results, dict) and 'masks' in mask_results else len(mask_results)}"
        )

    except Exception as e:
        raise RuntimeError(f"Mask generation failed: {e}")

    # Create a mask raster with unique IDs for each mask
    mask_raster = np.zeros((image_data.shape[0], image_data.shape[1]), dtype=np.uint32)
    mask_records = []

    # Process each mask based on the structure of mask_results
    if (
        isinstance(mask_results, dict)
        and "masks" in mask_results
        and "scores" in mask_results
    ):
        # Handle dictionary with 'masks' and 'scores' lists
        print("Processing masks...")
        total_masks = len(mask_results["masks"])

        # Create progress bar
        for i, (mask_data, score) in enumerate(
            tqdm(
                zip(mask_results["masks"], mask_results["scores"]),
                total=total_masks,
                desc="Processing masks",
            )
        ):
            mask_id = i + 1  # Start IDs at 1

            # Convert to numpy if not already
            if not isinstance(mask_data, np.ndarray):
                # Try to convert from tensor or other format if needed
                try:
                    mask_data = np.array(mask_data)
                except Exception:
                    print(f"Could not convert mask at index {i} to numpy array")
                    continue

            mask_binary = mask_data.astype(bool)
            area_pixels = np.sum(mask_binary)

            # Skip if mask is smaller than the minimum size
            if area_pixels < min_object_size:
                continue

            # Add the mask to the raster with a unique ID
            mask_raster[mask_binary] = mask_id

            # Create a record for the CSV - without geometry calculation
            mask_records.append(
                {"mask_id": mask_id, "score": float(score), "area_pixels": area_pixels}
            )
    elif isinstance(mask_results, list):
        # Handle list of dictionaries format (SAM original format)
        print("Processing masks...")
        total_masks = len(mask_results)

        # Create progress bar
        for i, mask_result in enumerate(tqdm(mask_results, desc="Processing masks")):
            mask_id = i + 1  # Start IDs at 1

            # Try different possible key names for masks and scores
            mask_data = None
            score = None

            if isinstance(mask_result, dict):
                # Try to find mask data
                if "segmentation" in mask_result:
                    mask_data = mask_result["segmentation"]
                elif "mask" in mask_result:
                    mask_data = mask_result["mask"]

                # Try to find score
                if "score" in mask_result:
                    score = mask_result["score"]
                elif "predicted_iou" in mask_result:
                    score = mask_result["predicted_iou"]
                elif "stability_score" in mask_result:
                    score = mask_result["stability_score"]
                else:
                    score = 1.0  # Default score if none found
            else:
                # If mask_result is not a dict, it might be the mask directly
                try:
                    mask_data = np.array(mask_result)
                    score = 1.0  # Default score
                except Exception:
                    print(f"Could not process mask at index {i}")
                    continue

            if mask_data is not None:
                # Convert to numpy if not already
                if not isinstance(mask_data, np.ndarray):
                    try:
                        mask_data = np.array(mask_data)
                    except Exception:
                        print(f"Could not convert mask at index {i} to numpy array")
                        continue

                mask_binary = mask_data.astype(bool)
                area_pixels = np.sum(mask_binary)

                # Skip if mask is smaller than the minimum size
                if area_pixels < min_object_size:
                    continue

                # Add the mask to the raster with a unique ID
                mask_raster[mask_binary] = mask_id

                # Create a record for the CSV - without geometry calculation
                mask_records.append(
                    {
                        "mask_id": mask_id,
                        "score": float(score),
                        "area_pixels": area_pixels,
                    }
                )
    else:
        # If we couldn't figure out the format, raise an error
        raise ValueError(f"Unexpected format for mask_results: {type(mask_results)}")

    print(f"Number of final masks (after size filtering): {len(mask_records)}")

    # Save the mask raster as a GeoTIFF
    print(f"Saving mask GeoTIFF to {output_mask_path}")
    output_profile = profile.copy()
    output_profile.update(dtype=rasterio.uint32, count=1, compress="lzw", nodata=0)

    with rasterio.open(output_mask_path, "w", **output_profile) as dst:
        dst.write(mask_raster.astype(rasterio.uint32), 1)

    # Save the mask data as a CSV
    print(f"Saving mask metadata to {output_csv_path}")
    mask_df = pd.DataFrame(mask_records)
    mask_df.to_csv(output_csv_path, index=False)

    print("Processing complete!")
    return output_mask_path, output_csv_path

masks_to_vector(mask_path, output_path=None, simplify_tolerance=1.0, mask_threshold=0.5, min_object_area=100, max_object_area=None, nms_iou_threshold=0.5)

Convert a building mask GeoTIFF to vector polygons and save as a vector dataset.

Parameters:

Name	Type	Description	Default
mask_path	str	Path to the building masks GeoTIFF	required
output_path	Optional[str]	Path to save the output GeoJSON (default: mask_path with .geojson extension)	None
simplify_tolerance	float	Tolerance for polygon simplification (default: self.simplify_tolerance)	1.0
mask_threshold	float	Threshold for mask binarization (default: self.mask_threshold)	0.5
min_object_area	int	Minimum area in pixels to keep a building (default: self.min_object_area)	100
max_object_area	Optional[int]	Maximum area in pixels to keep a building (default: self.max_object_area)	None
nms_iou_threshold	float	IoU threshold for non-maximum suppression (default: self.nms_iou_threshold)	0.5

Returns:

Name	Type	Description
Any	Any	GeoDataFrame with building footprints

Source code in geoai/utils/raster.py

def masks_to_vector(
    mask_path: str,
    output_path: Optional[str] = None,
    simplify_tolerance: float = 1.0,
    mask_threshold: float = 0.5,
    min_object_area: int = 100,
    max_object_area: Optional[int] = None,
    nms_iou_threshold: float = 0.5,
) -> Any:
    """
    Convert a building mask GeoTIFF to vector polygons and save as a vector dataset.

    Args:
        mask_path: Path to the building masks GeoTIFF
        output_path: Path to save the output GeoJSON (default: mask_path with .geojson extension)
        simplify_tolerance: Tolerance for polygon simplification (default: self.simplify_tolerance)
        mask_threshold: Threshold for mask binarization (default: self.mask_threshold)
        min_object_area: Minimum area in pixels to keep a building (default: self.min_object_area)
        max_object_area: Maximum area in pixels to keep a building (default: self.max_object_area)
        nms_iou_threshold: IoU threshold for non-maximum suppression (default: self.nms_iou_threshold)

    Returns:
        Any: GeoDataFrame with building footprints
    """
    import cv2  # Lazy import to avoid QGIS opencv conflicts

    # Set default output path if not provided
    # if output_path is None:
    #     output_path = os.path.splitext(mask_path)[0] + ".geojson"

    print(f"Converting mask to GeoJSON with parameters:")
    print(f"- Mask threshold: {mask_threshold}")
    print(f"- Min building area: {min_object_area}")
    print(f"- Simplify tolerance: {simplify_tolerance}")
    print(f"- NMS IoU threshold: {nms_iou_threshold}")

    # Open the mask raster
    with rasterio.open(mask_path) as src:
        # Read the mask data
        mask_data = src.read(1)
        transform = src.transform
        crs = src.crs

        # Print mask statistics
        print(f"Mask dimensions: {mask_data.shape}")
        print(f"Mask value range: {mask_data.min()} to {mask_data.max()}")

        # Prepare for connected component analysis
        # Binarize the mask based on threshold
        binary_mask = (mask_data > (mask_threshold * 255)).astype(np.uint8)

        # Apply morphological operations for better results (optional)
        kernel = np.ones((3, 3), np.uint8)
        binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, kernel)

        # Find connected components
        num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(
            binary_mask, connectivity=8
        )

        print(f"Found {num_labels-1} potential buildings")  # Subtract 1 for background

        # Create list to store polygons and confidence values
        all_polygons = []
        all_confidences = []

        # Process each component (skip the first one which is background)
        for i in tqdm(range(1, num_labels)):
            # Extract this building
            area = stats[i, cv2.CC_STAT_AREA]

            # Skip if too small
            if area < min_object_area:
                continue

            # Skip if too large
            if max_object_area is not None and area > max_object_area:
                continue

            # Create a mask for this building
            building_mask = (labels == i).astype(np.uint8)

            # Find contours
            contours, _ = cv2.findContours(
                building_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
            )

            # Process each contour
            for contour in contours:
                # Skip if too few points
                if contour.shape[0] < 3:
                    continue

                # Simplify contour if it has many points
                if contour.shape[0] > 50 and simplify_tolerance > 0:
                    epsilon = simplify_tolerance * cv2.arcLength(contour, True)
                    contour = cv2.approxPolyDP(contour, epsilon, True)

                # Convert to list of (x, y) coordinates
                polygon_points = contour.reshape(-1, 2)

                # Convert pixel coordinates to geographic coordinates
                geo_points = []
                for x, y in polygon_points:
                    gx, gy = transform * (x, y)
                    geo_points.append((gx, gy))

                # Create Shapely polygon
                if len(geo_points) >= 3:
                    try:
                        shapely_poly = Polygon(geo_points)
                        if shapely_poly.is_valid and shapely_poly.area > 0:
                            all_polygons.append(shapely_poly)

                            # Calculate "confidence" as normalized size
                            # This is a proxy since we don't have model confidence scores
                            normalized_size = min(1.0, area / 1000)  # Cap at 1.0
                            all_confidences.append(normalized_size)
                    except Exception as e:
                        print(f"Error creating polygon: {e}")

        print(f"Created {len(all_polygons)} valid polygons")

        # Create GeoDataFrame
        if not all_polygons:
            print("No valid polygons found")
            return None

        gdf = gpd.GeoDataFrame(
            {
                "geometry": all_polygons,
                "confidence": all_confidences,
                "class": 1,  # Building class
            },
            crs=crs,
        )

        def filter_overlapping_polygons(gdf, **kwargs):
            """
            Filter overlapping polygons using non-maximum suppression.

            Args:
                gdf: GeoDataFrame with polygons
                **kwargs: Optional parameters:
                    nms_iou_threshold: IoU threshold for filtering

            Returns:
                Filtered GeoDataFrame
            """
            if len(gdf) <= 1:
                return gdf

            # Get parameters from kwargs or use instance defaults
            iou_threshold = kwargs.get("nms_iou_threshold", nms_iou_threshold)

            # Sort by confidence
            gdf = gdf.sort_values("confidence", ascending=False)

            # Fix any invalid geometries
            gdf["geometry"] = gdf["geometry"].apply(
                lambda geom: geom.buffer(0) if not geom.is_valid else geom
            )

            keep_indices = []
            polygons = gdf.geometry.values

            for i in range(len(polygons)):
                if i in keep_indices:
                    continue

                keep = True
                for j in keep_indices:
                    # Skip invalid geometries
                    if not polygons[i].is_valid or not polygons[j].is_valid:
                        continue

                    # Calculate IoU
                    try:
                        intersection = polygons[i].intersection(polygons[j]).area
                        union = polygons[i].area + polygons[j].area - intersection
                        iou = intersection / union if union > 0 else 0

                        if iou > iou_threshold:
                            keep = False
                            break
                    except Exception:
                        # Skip on topology exceptions
                        continue

                if keep:
                    keep_indices.append(i)

            return gdf.iloc[keep_indices]

        # Apply non-maximum suppression to remove overlapping polygons
        gdf = filter_overlapping_polygons(gdf, nms_iou_threshold=nms_iou_threshold)

        print(f"Final building count after filtering: {len(gdf)}")

        # Save to file
        if output_path is not None:
            gdf.to_file(output_path)
            print(f"Saved {len(gdf)} building footprints to {output_path}")

        return gdf

mosaic_geotiffs(input_dir, output_file, mask_file=None)

Create a mosaic from all GeoTIFF files as a Cloud Optimized GeoTIFF (COG).

This function identifies all GeoTIFF files in the specified directory, creates a seamless mosaic with proper handling of nodata values, and saves as a Cloud Optimized GeoTIFF format. If a mask file is provided, the output will be clipped to the extent of the mask.

Parameters:

Name	Type	Description	Default
input_dir	str	Path to the directory containing GeoTIFF files.	required
output_file	str	Path to the output Cloud Optimized GeoTIFF file.	required
mask_file	str	Path to a mask file to clip the output. If provided, the output will be clipped to the extent of this mask. Defaults to None.	None

Returns:

Name	Type	Description
bool	None	True if the mosaic was created successfully, False otherwise.

Examples:

>>> mosaic_geotiffs('naip', 'merged_naip.tif')
True
>>> mosaic_geotiffs('naip', 'merged_naip.tif', 'boundary.tif')
True

Source code in geoai/utils/raster.py

def mosaic_geotiffs(
    input_dir: str, output_file: str, mask_file: Optional[str] = None
) -> None:
    """Create a mosaic from all GeoTIFF files as a Cloud Optimized GeoTIFF (COG).

    This function identifies all GeoTIFF files in the specified directory,
    creates a seamless mosaic with proper handling of nodata values, and saves
    as a Cloud Optimized GeoTIFF format. If a mask file is provided, the output
    will be clipped to the extent of the mask.

    Args:
        input_dir (str): Path to the directory containing GeoTIFF files.
        output_file (str): Path to the output Cloud Optimized GeoTIFF file.
        mask_file (str, optional): Path to a mask file to clip the output.
            If provided, the output will be clipped to the extent of this mask.
            Defaults to None.

    Returns:
        bool: True if the mosaic was created successfully, False otherwise.

    Examples:
        >>> mosaic_geotiffs('naip', 'merged_naip.tif')
        True
        >>> mosaic_geotiffs('naip', 'merged_naip.tif', 'boundary.tif')
        True
    """
    import glob

    from osgeo import gdal

    gdal.UseExceptions()
    # Get all tif files in the directory
    tif_files = glob.glob(os.path.join(input_dir, "*.tif"))

    if not tif_files:
        print("No GeoTIFF files found in the specified directory.")
        return False

    # Analyze the first input file to determine compression and nodata settings
    ds = gdal.Open(tif_files[0])
    if ds is None:
        print(f"Unable to open {tif_files[0]}")
        return False

    # Get driver metadata from the first file
    driver = ds.GetDriver()
    creation_options = []

    # Check compression type
    metadata = ds.GetMetadata("IMAGE_STRUCTURE")
    if "COMPRESSION" in metadata:
        compression = metadata["COMPRESSION"]
        creation_options.append(f"COMPRESS={compression}")
    else:
        # Default compression if none detected
        creation_options.append("COMPRESS=LZW")

    # Add COG-specific creation options
    creation_options.extend(["TILED=YES", "BLOCKXSIZE=512", "BLOCKYSIZE=512"])

    # Check for nodata value in the first band of the first file
    band = ds.GetRasterBand(1)
    has_nodata = band.GetNoDataValue() is not None
    nodata_value = band.GetNoDataValue() if has_nodata else None

    # Close the dataset
    ds = None

    # Create a temporary VRT (Virtual Dataset)
    vrt_path = os.path.join(input_dir, "temp_mosaic.vrt")

    # Build VRT from input files with proper nodata handling
    vrt_options = gdal.BuildVRTOptions(
        resampleAlg="nearest",
        srcNodata=nodata_value if has_nodata else None,
        VRTNodata=nodata_value if has_nodata else None,
    )
    vrt_dataset = gdal.BuildVRT(vrt_path, tif_files, options=vrt_options)

    # Close the VRT dataset to flush it to disk
    vrt_dataset = None

    # Create temp mosaic
    temp_mosaic = output_file + ".temp.tif"

    # Convert VRT to GeoTIFF with the same compression as input
    translate_options = gdal.TranslateOptions(
        format="GTiff",
        creationOptions=creation_options,
        noData=nodata_value if has_nodata else None,
    )
    gdal.Translate(temp_mosaic, vrt_path, options=translate_options)

    # Apply mask if provided
    if mask_file and os.path.exists(mask_file):
        print(f"Clipping mosaic to mask: {mask_file}")

        # Create a temporary clipped file
        clipped_mosaic = output_file + ".clipped.tif"

        # Open mask file
        mask_ds = gdal.Open(mask_file)
        if mask_ds is None:
            print(f"Unable to open mask file: {mask_file}")
            # Continue without clipping
        else:
            # Get mask extent
            mask_geotransform = mask_ds.GetGeoTransform()
            mask_projection = mask_ds.GetProjection()
            mask_ulx = mask_geotransform[0]
            mask_uly = mask_geotransform[3]
            mask_lrx = mask_ulx + (mask_geotransform[1] * mask_ds.RasterXSize)
            mask_lry = mask_uly + (mask_geotransform[5] * mask_ds.RasterYSize)

            # Close mask dataset
            mask_ds = None

            # Use warp options to clip
            warp_options = gdal.WarpOptions(
                format="GTiff",
                outputBounds=[mask_ulx, mask_lry, mask_lrx, mask_uly],
                dstSRS=mask_projection,
                creationOptions=creation_options,
                srcNodata=nodata_value if has_nodata else None,
                dstNodata=nodata_value if has_nodata else None,
            )

            # Apply clipping
            gdal.Warp(clipped_mosaic, temp_mosaic, options=warp_options)

            # Remove the unclipped temp mosaic and use the clipped one
            os.remove(temp_mosaic)
            temp_mosaic = clipped_mosaic

    # Create internal overviews for the temp mosaic
    ds = gdal.Open(temp_mosaic, gdal.GA_Update)
    overview_list = [2, 4, 8, 16, 32]
    ds.BuildOverviews("NEAREST", overview_list)
    ds = None  # Close the dataset to ensure overviews are written

    # Convert the temp mosaic to a proper COG
    cog_options = gdal.TranslateOptions(
        format="GTiff",
        creationOptions=[
            "TILED=YES",
            "COPY_SRC_OVERVIEWS=YES",
            "COMPRESS=DEFLATE",
            "PREDICTOR=2",
            "BLOCKXSIZE=512",
            "BLOCKYSIZE=512",
        ],
        noData=nodata_value if has_nodata else None,
    )
    gdal.Translate(output_file, temp_mosaic, options=cog_options)

    # Clean up temporary files
    if os.path.exists(vrt_path):
        os.remove(vrt_path)
    if os.path.exists(temp_mosaic):
        os.remove(temp_mosaic)

    print(f"Cloud Optimized GeoTIFF mosaic created successfully: {output_file}")
    return True

normalize_radiometric(subject_image, reference_image, output_path=None, method='lirrn', p_n=500, num_quantisation_classes=3, num_sampling_rounds=3, subsample_ratio=0.1, random_state=None)

Normalize subject image radiometry to match a reference image.

Adjusts brightness and contrast of the subject image so that its pixel value distribution matches the reference image. This is essential for multi-temporal analysis where images are acquired under different atmospheric conditions, sensor calibrations, or illumination angles.

Currently supports the LIRRN (Location-Independent Relative Radiometric Normalization) method, which uses multi-Otsu thresholding and linear regression to identify pseudo-invariant features and transform pixel values band-by-band.

Reference: doi:10.3390/s24072272

Parameters:

Name	Type	Description	Default
subject_image	Union[str, ndarray]	Path to the subject GeoTIFF or numpy array with shape (H, W, B). The image to be normalized.	required
reference_image	Union[str, ndarray]	Path to the reference GeoTIFF or numpy array with shape (H, W, B). The target radiometry to match.	required
output_path	Optional[str]	Path to save the normalized image as GeoTIFF. Only applicable when subject_image is a file path (so spatial metadata is available). If None, the array is returned without saving. Default: None.	None
method	str	Normalization method. Currently only "lirrn" is supported. Default: "lirrn".	'lirrn'
p_n	int	Number of pseudo-invariant feature samples per quantization level. Higher values increase accuracy but slow computation. Default: 500.	500
num_quantisation_classes	int	Number of brightness strata for stratified sampling. Default: 3.	3
num_sampling_rounds	int	Number of iterative refinement rounds for sample selection. Default: 3.	3
subsample_ratio	float	Fraction of candidates retained for regression. Default: 0.1.	0.1
random_state	Optional[Union[int, Generator]]	Seed or numpy Generator for reproducible results. Default: None (non-deterministic).	None

Returns:

Type	Description
Tuple[ndarray, Dict[str, ndarray]]	Tuple of (normalized_image, metrics) where: - normalized_image: numpy array (H, W, B) float64. - metrics: dict with keys "rmse" and "r_adj", each a numpy array of length B.

Raises:

Type	Description
ValueError	If method is not "lirrn".
ValueError	If p_n < 1 or num_sampling_rounds < 1.
ValueError	If subject and reference have different band counts.
ValueError	If input arrays are not 3-dimensional.
ValueError	If output_path is set but subject_image is an array.
FileNotFoundError	If file paths do not point to existing files.

Examples:

Normalize a satellite image using file paths:

>>> from geoai import normalize_radiometric
>>> norm_img, metrics = normalize_radiometric(
...     "subject.tif",
...     "reference.tif",
...     output_path="normalized.tif",
... )
>>> print(f"RMSE per band: {metrics['rmse']}")

Normalize using numpy arrays:

>>> import numpy as np
>>> subject = np.random.rand(100, 100, 4)
>>> reference = np.random.rand(120, 120, 4)
>>> norm_img, metrics = normalize_radiometric(subject, reference)
>>> norm_img.shape
(100, 100, 4)

Note

The subject and reference images must have the same number of bands but may have different spatial dimensions (height and width).

Source code in geoai/landcover_utils.py

def normalize_radiometric(
    subject_image: Union[str, np.ndarray],
    reference_image: Union[str, np.ndarray],
    output_path: Optional[str] = None,
    method: str = "lirrn",
    p_n: int = 500,
    num_quantisation_classes: int = 3,
    num_sampling_rounds: int = 3,
    subsample_ratio: float = 0.1,
    random_state: Optional[Union[int, np.random.Generator]] = None,
) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
    """Normalize subject image radiometry to match a reference image.

    Adjusts brightness and contrast of the subject image so that its pixel
    value distribution matches the reference image. This is essential for
    multi-temporal analysis where images are acquired under different
    atmospheric conditions, sensor calibrations, or illumination angles.

    Currently supports the LIRRN (Location-Independent Relative Radiometric
    Normalization) method, which uses multi-Otsu thresholding and linear
    regression to identify pseudo-invariant features and transform pixel
    values band-by-band.

    Reference: doi:10.3390/s24072272

    Args:
        subject_image: Path to the subject GeoTIFF or numpy array with
            shape (H, W, B). The image to be normalized.
        reference_image: Path to the reference GeoTIFF or numpy array with
            shape (H, W, B). The target radiometry to match.
        output_path: Path to save the normalized image as GeoTIFF. Only
            applicable when *subject_image* is a file path (so spatial
            metadata is available). If None, the array is returned without
            saving. Default: None.
        method: Normalization method. Currently only ``"lirrn"`` is
            supported. Default: ``"lirrn"``.
        p_n: Number of pseudo-invariant feature samples per quantization
            level. Higher values increase accuracy but slow computation.
            Default: 500.
        num_quantisation_classes: Number of brightness strata for stratified
            sampling. Default: 3.
        num_sampling_rounds: Number of iterative refinement rounds for
            sample selection. Default: 3.
        subsample_ratio: Fraction of candidates retained for regression.
            Default: 0.1.
        random_state: Seed or numpy Generator for reproducible results.
            Default: None (non-deterministic).

    Returns:
        Tuple of (normalized_image, metrics) where:
            - normalized_image: numpy array (H, W, B) float64.
            - metrics: dict with keys ``"rmse"`` and ``"r_adj"``, each a
              numpy array of length B.

    Raises:
        ValueError: If *method* is not ``"lirrn"``.
        ValueError: If *p_n* < 1 or *num_sampling_rounds* < 1.
        ValueError: If subject and reference have different band counts.
        ValueError: If input arrays are not 3-dimensional.
        ValueError: If *output_path* is set but *subject_image* is an array.
        FileNotFoundError: If file paths do not point to existing files.

    Examples:
        Normalize a satellite image using file paths:

        >>> from geoai import normalize_radiometric
        >>> norm_img, metrics = normalize_radiometric(
        ...     "subject.tif",
        ...     "reference.tif",
        ...     output_path="normalized.tif",
        ... )
        >>> print(f"RMSE per band: {metrics['rmse']}")

        Normalize using numpy arrays:

        >>> import numpy as np
        >>> subject = np.random.rand(100, 100, 4)
        >>> reference = np.random.rand(120, 120, 4)
        >>> norm_img, metrics = normalize_radiometric(subject, reference)
        >>> norm_img.shape
        (100, 100, 4)

    Note:
        The subject and reference images must have the same number of bands
        but may have different spatial dimensions (height and width).
    """
    # --- Validate parameters ---
    if method != "lirrn":
        raise ValueError(
            f"Unsupported normalization method {method!r}. "
            "Currently only 'lirrn' is supported."
        )
    if p_n < 1:
        raise ValueError(f"p_n must be >= 1, got {p_n}")
    if num_sampling_rounds < 1:
        raise ValueError(f"num_sampling_rounds must be >= 1, got {num_sampling_rounds}")
    if subsample_ratio <= 0 or subsample_ratio > 1:
        raise ValueError(f"subsample_ratio must be in (0, 1], got {subsample_ratio}")

    # --- Resolve inputs ---
    profile = None
    if isinstance(subject_image, str):
        sub_arr, profile = _load_raster(subject_image)
    else:
        sub_arr = np.asarray(subject_image, dtype=np.float64)
        if sub_arr.ndim != 3:
            raise ValueError(
                f"subject_image must be 3-D (H, W, B), got {sub_arr.ndim}-D"
            )

    if isinstance(reference_image, str):
        ref_arr, _ = _load_raster(reference_image)
    else:
        ref_arr = np.asarray(reference_image, dtype=np.float64)
        if ref_arr.ndim != 3:
            raise ValueError(
                f"reference_image must be 3-D (H, W, B), got {ref_arr.ndim}-D"
            )

    if output_path is not None and profile is None:
        raise ValueError(
            "output_path requires subject_image to be a file path "
            "(not an array) so that spatial metadata is available."
        )

    # Band count check
    if sub_arr.shape[2] != ref_arr.shape[2]:
        raise ValueError(
            f"Band count mismatch: subject has {sub_arr.shape[2]} bands, "
            f"reference has {ref_arr.shape[2]} bands."
        )

    # Handle NaN / inf
    if np.any(~np.isfinite(sub_arr)):
        warnings.warn(
            "subject_image contains NaN or infinite values; " "replacing with 0.",
            stacklevel=2,
        )
        sub_arr = np.nan_to_num(sub_arr, nan=0.0, posinf=0.0, neginf=0.0)

    if np.any(~np.isfinite(ref_arr)):
        warnings.warn(
            "reference_image contains NaN or infinite values; " "replacing with 0.",
            stacklevel=2,
        )
        ref_arr = np.nan_to_num(ref_arr, nan=0.0, posinf=0.0, neginf=0.0)

    # --- Build RNG ---
    if isinstance(random_state, np.random.Generator):
        rng = random_state
    else:
        rng = np.random.default_rng(random_state)

    # --- Run normalization ---
    norm_img, rmse, r_adj = _lirrn(
        p_n,
        sub_arr,
        ref_arr,
        num_quantisation_classes=num_quantisation_classes,
        num_sampling_rounds=num_sampling_rounds,
        subsample_ratio=subsample_ratio,
        rng=rng,
    )

    metrics = {"rmse": rmse, "r_adj": r_adj}

    if output_path is not None:
        _save_raster(output_path, norm_img, profile)

    return norm_img, metrics

orthogonalize(input_path, output_path=None, epsilon=0.2, min_area=10, min_segments=4, area_tolerance=0.7, detect_triangles=True)

Orthogonalizes object masks in a GeoTIFF file.

This function reads a GeoTIFF containing object masks (binary or labeled regions), converts the raster masks to vector polygons, applies orthogonalization to each polygon, and optionally writes the result to a GeoJSON file. The source code is adapted from the Solar Panel Detection algorithm by Esri. See https://www.arcgis.com/home/item.html?id=c2508d72f2614104bfcfd5ccf1429284. Credits to Esri for the original code.

Parameters:

Name	Type	Description	Default
input_path	str	Path to the input GeoTIFF file.	required
output_path	str	Path to save the output GeoJSON file. If None, no file is saved.	None
epsilon	float	Simplification tolerance for the Douglas-Peucker algorithm. Higher values result in more simplification. Default is 0.2.	0.2
min_area	float	Minimum area of polygons to process (smaller ones are kept as-is).	10
min_segments	int	Minimum number of segments to keep after simplification. Default is 4 (for rectangular shapes).	4
area_tolerance	float	Allowed ratio of area change. Values less than 1.0 restrict area change. Default is 0.7 (allows reduction to 70% of original area).	0.7
detect_triangles	bool	If True, performs additional check to avoid creating triangular shapes.	True

Returns:

Name	Type	Description
Any	Any	A GeoDataFrame containing the orthogonalized features.

Source code in geoai/utils/geometry.py

def orthogonalize(
    input_path,
    output_path=None,
    epsilon=0.2,
    min_area=10,
    min_segments=4,
    area_tolerance=0.7,
    detect_triangles=True,
) -> Any:
    """
    Orthogonalizes object masks in a GeoTIFF file.

    This function reads a GeoTIFF containing object masks (binary or labeled regions),
    converts the raster masks to vector polygons, applies orthogonalization to each polygon,
    and optionally writes the result to a GeoJSON file.
    The source code is adapted from the Solar Panel Detection algorithm by Esri.
    See https://www.arcgis.com/home/item.html?id=c2508d72f2614104bfcfd5ccf1429284.
    Credits to Esri for the original code.

    Args:
        input_path (str): Path to the input GeoTIFF file.
        output_path (str, optional): Path to save the output GeoJSON file. If None, no file is saved.
        epsilon (float, optional): Simplification tolerance for the Douglas-Peucker algorithm.
            Higher values result in more simplification. Default is 0.2.
        min_area (float, optional): Minimum area of polygons to process (smaller ones are kept as-is).
        min_segments (int, optional): Minimum number of segments to keep after simplification.
            Default is 4 (for rectangular shapes).
        area_tolerance (float, optional): Allowed ratio of area change. Values less than 1.0 restrict
            area change. Default is 0.7 (allows reduction to 70% of original area).
        detect_triangles (bool, optional): If True, performs additional check to avoid creating triangular shapes.

    Returns:
        Any: A GeoDataFrame containing the orthogonalized features.
    """
    import cv2  # Lazy import to avoid QGIS opencv conflicts

    from functools import partial

    def orthogonalize_ring(ring, epsilon=0.2, min_segments=4):
        """
        Orthogonalizes a ring (list of coordinates).

        Args:
            ring (list): List of [x, y] coordinates forming a ring
            epsilon (float, optional): Simplification tolerance
            min_segments (int, optional): Minimum number of segments to keep

        Returns:
            list: Orthogonalized list of coordinates
        """
        if len(ring) <= 3:
            return ring

        # Convert to numpy array
        ring_arr = np.array(ring)

        # Get orientation
        angle = math.degrees(get_orientation(ring_arr))

        # Simplify using Ramer-Douglas-Peucker algorithm
        ring_arr = simplify(ring_arr, eps=epsilon)

        # If simplified too much, adjust epsilon to maintain minimum segments
        if len(ring_arr) < min_segments:
            # Try with smaller epsilon until we get at least min_segments points
            for adjust_factor in [0.75, 0.5, 0.25, 0.1]:
                test_arr = simplify(np.array(ring), eps=epsilon * adjust_factor)
                if len(test_arr) >= min_segments:
                    ring_arr = test_arr
                    break

        # Convert to dataframe for processing
        df = to_dataframe(ring_arr)

        # Add orientation information
        add_orientation(df, angle)

        # Align segments to orthogonal directions
        df = align(df)

        # Merge collinear line segments
        df = merge_lines(df)

        if len(df) == 0:
            return ring

        # If we have a triangle-like result (3 segments or less), return the original shape
        if len(df) <= 3:
            return ring

        # Join the orthogonalized segments back into a ring
        joined_ring = join_ring(df)

        # If the join operation didn't produce a valid ring, return the original
        if len(joined_ring) == 0 or len(joined_ring[0]) < 3:
            return ring

        # Enhanced validation: check for triangular result and geometric validity
        result_coords = joined_ring[0]

        # If result has 3 or fewer points (triangle), use original
        if len(result_coords) <= 3:  # 2 points + closing point (degenerate)
            return ring

        # Additional validation: check for degenerate geometry
        # Calculate area ratio to detect if the shape got severely distorted
        def calculate_polygon_area(coords):
            if len(coords) < 3:
                return 0
            area = 0
            n = len(coords)
            for i in range(n):
                j = (i + 1) % n
                area += coords[i][0] * coords[j][1]
                area -= coords[j][0] * coords[i][1]
            return abs(area) / 2

        original_area = calculate_polygon_area(ring)
        result_area = calculate_polygon_area(result_coords)

        # If the area changed dramatically (more than 30% shrinkage or 300% growth), use original
        if original_area > 0 and result_area > 0:
            area_ratio = result_area / original_area
            if area_ratio < 0.3 or area_ratio > 3.0:
                return ring

        # Check for triangular spikes and problematic artifacts
        very_acute_angle_count = 0
        triangular_spike_detected = False

        for i in range(len(result_coords) - 1):  # -1 to exclude closing point
            p1 = result_coords[i - 1]
            p2 = result_coords[i]
            p3 = result_coords[(i + 1) % (len(result_coords) - 1)]

            # Calculate angle at p2
            v1 = np.array([p1[0] - p2[0], p1[1] - p2[1]])
            v2 = np.array([p3[0] - p2[0], p3[1] - p2[1]])

            v1_norm = np.linalg.norm(v1)
            v2_norm = np.linalg.norm(v2)

            if v1_norm > 0 and v2_norm > 0:
                cos_angle = np.dot(v1, v2) / (v1_norm * v2_norm)
                cos_angle = np.clip(cos_angle, -1, 1)
                angle = np.arccos(cos_angle)

                # Count very acute angles (< 20 degrees) - these are likely spikes
                if angle < np.pi / 9:  # 20 degrees
                    very_acute_angle_count += 1
                    # If it's very acute with short sides, it's definitely a spike
                    if v1_norm < 5 or v2_norm < 5:
                        triangular_spike_detected = True

        # Check for excessively long edges that might be artifacts
        edge_lengths = []
        for i in range(len(result_coords) - 1):
            edge_len = np.sqrt(
                (result_coords[i + 1][0] - result_coords[i][0]) ** 2
                + (result_coords[i + 1][1] - result_coords[i][1]) ** 2
            )
            edge_lengths.append(edge_len)

        excessive_edge_detected = False
        if len(edge_lengths) > 0:
            avg_edge_length = np.mean(edge_lengths)
            max_edge_length = np.max(edge_lengths)
            # Only reject if edge is extremely disproportionate (8x average)
            if max_edge_length > avg_edge_length * 8:
                excessive_edge_detected = True

        # Check for triangular artifacts by detecting spikes that extend beyond bounds
        # Calculate original bounds
        orig_xs = [p[0] for p in ring]
        orig_ys = [p[1] for p in ring]
        orig_min_x, orig_max_x = min(orig_xs), max(orig_xs)
        orig_min_y, orig_max_y = min(orig_ys), max(orig_ys)
        orig_width = orig_max_x - orig_min_x
        orig_height = orig_max_y - orig_min_y

        # Calculate result bounds
        result_xs = [p[0] for p in result_coords]
        result_ys = [p[1] for p in result_coords]
        result_min_x, result_max_x = min(result_xs), max(result_xs)
        result_min_y, result_max_y = min(result_ys), max(result_ys)

        # Stricter bounds checking to catch triangular artifacts
        bounds_extension_detected = False
        # More conservative: only allow 10% extension
        tolerance_x = max(orig_width * 0.1, 1.0)  # 10% tolerance, at least 1 unit
        tolerance_y = max(orig_height * 0.1, 1.0)  # 10% tolerance, at least 1 unit

        if (
            result_min_x < orig_min_x - tolerance_x
            or result_max_x > orig_max_x + tolerance_x
            or result_min_y < orig_min_y - tolerance_y
            or result_max_y > orig_max_y + tolerance_y
        ):
            bounds_extension_detected = True

        # Reject if we detect triangular spikes, excessive edges, or bounds violations
        if (
            triangular_spike_detected
            or very_acute_angle_count > 2  # Multiple very acute angles
            or excessive_edge_detected
            or bounds_extension_detected
        ):  # Any significant bounds extension
            return ring

        # Convert back to a list and ensure it's closed
        result = joined_ring[0].tolist()
        if len(result) > 0 and (result[0] != result[-1]):
            result.append(result[0])

        return result

    def vectorize_mask(mask, transform):
        """
        Converts a binary mask to vector polygons.

        Args:
            mask (numpy.ndarray): Binary mask where non-zero values represent objects
            transform (rasterio.transform.Affine): Affine transformation matrix

        Returns:
            list: List of GeoJSON features
        """
        shapes = features.shapes(mask, transform=transform)
        features_list = []

        for shape, value in shapes:
            if value > 0:  # Only process non-zero values (actual objects)
                features_list.append(
                    {
                        "type": "Feature",
                        "properties": {"value": int(value)},
                        "geometry": shape,
                    }
                )

        return features_list

    def rasterize_features(features, shape, transform, dtype=np.uint8):
        """
        Converts vector features back to a raster mask.

        Args:
            features (list): List of GeoJSON features
            shape (tuple): Shape of the output raster (height, width)
            transform (rasterio.transform.Affine): Affine transformation matrix
            dtype (numpy.dtype, optional): Data type of the output raster

        Returns:
            numpy.ndarray: Rasterized mask
        """
        mask = features.rasterize(
            [
                (feature["geometry"], feature["properties"]["value"])
                for feature in features
            ],
            out_shape=shape,
            transform=transform,
            fill=0,
            dtype=dtype,
        )

        return mask

    # The following helper functions are from the original code
    def get_orientation(contour):
        """
        Calculate the orientation angle of a contour.

        Args:
            contour (numpy.ndarray): Array of shape (n, 2) containing point coordinates

        Returns:
            float: Orientation angle in radians
        """
        box = cv2.minAreaRect(contour.astype(int))
        (cx, cy), (w, h), angle = box
        return math.radians(angle)

    def simplify(contour, eps=0.2):
        """
        Simplify a contour using the Ramer-Douglas-Peucker algorithm.

        Args:
            contour (numpy.ndarray): Array of shape (n, 2) containing point coordinates
            eps (float, optional): Epsilon value for simplification

        Returns:
            numpy.ndarray: Simplified contour
        """
        return rdp(contour, epsilon=eps)

    def to_dataframe(ring):
        """
        Convert a ring to a pandas DataFrame with line segment information.

        Args:
            ring (numpy.ndarray): Array of shape (n, 2) containing point coordinates

        Returns:
            pandas.DataFrame: DataFrame with line segment information
        """
        df = pd.DataFrame(ring, columns=["x1", "y1"])
        df["x2"] = df["x1"].shift(-1)
        df["y2"] = df["y1"].shift(-1)
        df.dropna(inplace=True)
        df["angle_atan"] = np.arctan2((df["y2"] - df["y1"]), (df["x2"] - df["x1"]))
        df["angle_atan_deg"] = df["angle_atan"] * 57.2958
        df["len"] = np.sqrt((df["y2"] - df["y1"]) ** 2 + (df["x2"] - df["x1"]) ** 2)
        df["cx"] = (df["x2"] + df["x1"]) / 2.0
        df["cy"] = (df["y2"] + df["y1"]) / 2.0
        return df

    def add_orientation(df, angle):
        """
        Add orientation information to the DataFrame.

        Args:
            df (pandas.DataFrame): DataFrame with line segment information
            angle (float): Orientation angle in degrees

        Returns:
            None: Modifies the DataFrame in-place
        """
        rtangle = angle + 90
        is_parallel = (
            (df["angle_atan_deg"] > (angle - 45))
            & (df["angle_atan_deg"] < (angle + 45))
        ) | (
            (df["angle_atan_deg"] + 180 > (angle - 45))
            & (df["angle_atan_deg"] + 180 < (angle + 45))
        )
        df["angle"] = math.radians(angle)
        df["angle"] = df["angle"].where(is_parallel, math.radians(rtangle))

    def align(df):
        """
        Align line segments to their nearest orthogonal direction.

        Args:
            df (pandas.DataFrame): DataFrame with line segment information

        Returns:
            pandas.DataFrame: DataFrame with aligned line segments
        """
        # Handle edge case with empty dataframe
        if len(df) == 0:
            return df.copy()

        df_clone = df.copy()

        # Ensure angle column exists and has valid values
        if "angle" not in df_clone.columns or df_clone["angle"].isna().any():
            # If angle data is missing, add default angles based on atan2
            df_clone["angle"] = df_clone["angle_atan"]

        # Ensure length and center point data is valid
        if "len" not in df_clone.columns or df_clone["len"].isna().any():
            # Recalculate lengths if missing
            df_clone["len"] = np.sqrt(
                (df_clone["x2"] - df_clone["x1"]) ** 2
                + (df_clone["y2"] - df_clone["y1"]) ** 2
            )

        if "cx" not in df_clone.columns or df_clone["cx"].isna().any():
            df_clone["cx"] = (df_clone["x1"] + df_clone["x2"]) / 2.0

        if "cy" not in df_clone.columns or df_clone["cy"].isna().any():
            df_clone["cy"] = (df_clone["y1"] + df_clone["y2"]) / 2.0

        # Apply orthogonal alignment
        df_clone["x1"] = df_clone["cx"] - ((df_clone["len"] / 2) * np.cos(df["angle"]))
        df_clone["x2"] = df_clone["cx"] + ((df_clone["len"] / 2) * np.cos(df["angle"]))
        df_clone["y1"] = df_clone["cy"] - ((df_clone["len"] / 2) * np.sin(df["angle"]))
        df_clone["y2"] = df_clone["cy"] + ((df_clone["len"] / 2) * np.sin(df["angle"]))

        return df_clone

    def merge_lines(df_aligned):
        """
        Merge collinear line segments.

        Args:
            df_aligned (pandas.DataFrame): DataFrame with aligned line segments

        Returns:
            pandas.DataFrame: DataFrame with merged line segments
        """
        ortho_lines = []
        groups = df_aligned.groupby(
            (df_aligned["angle"].shift() != df_aligned["angle"]).cumsum()
        )
        for x, y in groups:
            group_cx = (y["cx"] * y["len"]).sum() / y["len"].sum()
            group_cy = (y["cy"] * y["len"]).sum() / y["len"].sum()
            cumlen = y["len"].sum()

            ortho_lines.append((group_cx, group_cy, cumlen, y["angle"].iloc[0]))

        ortho_list = []
        for cx, cy, length, rot_angle in ortho_lines:
            X1 = cx - (length / 2) * math.cos(rot_angle)
            X2 = cx + (length / 2) * math.cos(rot_angle)
            Y1 = cy - (length / 2) * math.sin(rot_angle)
            Y2 = cy + (length / 2) * math.sin(rot_angle)

            ortho_list.append(
                {
                    "x1": X1,
                    "y1": Y1,
                    "x2": X2,
                    "y2": Y2,
                    "len": length,
                    "cx": cx,
                    "cy": cy,
                    "angle": rot_angle,
                }
            )

        # Improved fix: Prevent merging that would create triangular or problematic shapes
        if (
            len(ortho_list) > 3 and ortho_list[0]["angle"] == ortho_list[-1]["angle"]
        ):  # join first and last segment if they're in same direction
            # Check if merging would result in 3 or 4 segments (potentially triangular)
            resulting_segments = len(ortho_list) - 1
            if resulting_segments <= 4:
                # For very small polygons, be extra cautious about merging
                # Calculate the spatial relationship between first and last segments
                first_center = np.array([ortho_list[0]["cx"], ortho_list[0]["cy"]])
                last_center = np.array([ortho_list[-1]["cx"], ortho_list[-1]["cy"]])
                center_distance = np.linalg.norm(first_center - last_center)

                # Get average segment length for comparison
                avg_length = sum(seg["len"] for seg in ortho_list) / len(ortho_list)

                # Only merge if segments are close enough and it won't create degenerate shapes
                if center_distance > avg_length * 1.5:
                    # Skip merging - segments are too far apart
                    pass
                else:
                    # Proceed with merging only for well-connected segments
                    totlen = ortho_list[0]["len"] + ortho_list[-1]["len"]
                    merge_cx = (
                        (ortho_list[0]["cx"] * ortho_list[0]["len"])
                        + (ortho_list[-1]["cx"] * ortho_list[-1]["len"])
                    ) / totlen

                    merge_cy = (
                        (ortho_list[0]["cy"] * ortho_list[0]["len"])
                        + (ortho_list[-1]["cy"] * ortho_list[-1]["len"])
                    ) / totlen

                    rot_angle = ortho_list[0]["angle"]
                    X1 = merge_cx - (totlen / 2) * math.cos(rot_angle)
                    X2 = merge_cx + (totlen / 2) * math.cos(rot_angle)
                    Y1 = merge_cy - (totlen / 2) * math.sin(rot_angle)
                    Y2 = merge_cy + (totlen / 2) * math.sin(rot_angle)

                    ortho_list[-1] = {
                        "x1": X1,
                        "y1": Y1,
                        "x2": X2,
                        "y2": Y2,
                        "len": totlen,
                        "cx": merge_cx,
                        "cy": merge_cy,
                        "angle": rot_angle,
                    }
                    ortho_list = ortho_list[1:]
            else:
                # For larger polygons, proceed with standard merging
                totlen = ortho_list[0]["len"] + ortho_list[-1]["len"]
                merge_cx = (
                    (ortho_list[0]["cx"] * ortho_list[0]["len"])
                    + (ortho_list[-1]["cx"] * ortho_list[-1]["len"])
                ) / totlen

                merge_cy = (
                    (ortho_list[0]["cy"] * ortho_list[0]["len"])
                    + (ortho_list[-1]["cy"] * ortho_list[-1]["len"])
                ) / totlen

                rot_angle = ortho_list[0]["angle"]
                X1 = merge_cx - (totlen / 2) * math.cos(rot_angle)
                X2 = merge_cx + (totlen / 2) * math.cos(rot_angle)
                Y1 = merge_cy - (totlen / 2) * math.sin(rot_angle)
                Y2 = merge_cy + (totlen / 2) * math.sin(rot_angle)

                ortho_list[-1] = {
                    "x1": X1,
                    "y1": Y1,
                    "x2": X2,
                    "y2": Y2,
                    "len": totlen,
                    "cx": merge_cx,
                    "cy": merge_cy,
                    "angle": rot_angle,
                }
                ortho_list = ortho_list[1:]
        ortho_df = pd.DataFrame(ortho_list)
        return ortho_df

    def find_intersection(x1, y1, x2, y2, x3, y3, x4, y4):
        """
        Find the intersection point of two line segments.

        Args:
            x1, y1, x2, y2: Coordinates of the first line segment
            x3, y3, x4, y4: Coordinates of the second line segment

        Returns:
            list: [x, y] coordinates of the intersection point

        Raises:
            ZeroDivisionError: If the lines are parallel or collinear
        """
        # Calculate the denominator of the intersection formula
        denominator = (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4)

        # Check if lines are parallel or collinear (denominator close to zero)
        if abs(denominator) < 1e-10:
            raise ZeroDivisionError("Lines are parallel or collinear")

        px = (
            (x1 * y2 - y1 * x2) * (x3 - x4) - (x1 - x2) * (x3 * y4 - y3 * x4)
        ) / denominator
        py = (
            (x1 * y2 - y1 * x2) * (y3 - y4) - (y1 - y2) * (x3 * y4 - y3 * x4)
        ) / denominator

        # Check if the intersection point is within a reasonable distance
        # from both line segments to avoid extreme extrapolation
        def point_on_segment(x, y, x1, y1, x2, y2, tolerance=2.0):
            # Check if point (x,y) is near the line segment from (x1,y1) to (x2,y2)
            # First check if it's near the infinite line
            line_len = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
            if line_len < 1e-10:
                return np.sqrt((x - x1) ** 2 + (y - y1) ** 2) <= tolerance

            t = ((x - x1) * (x2 - x1) + (y - y1) * (y2 - y1)) / (line_len**2)

            # Check distance to the infinite line
            proj_x = x1 + t * (x2 - x1)
            proj_y = y1 + t * (y2 - y1)
            dist_to_line = np.sqrt((x - proj_x) ** 2 + (y - proj_y) ** 2)

            # Check if the projection is near the segment, not just the infinite line
            if t < -tolerance or t > 1 + tolerance:
                # If far from the segment, compute distance to the nearest endpoint
                dist_to_start = np.sqrt((x - x1) ** 2 + (y - y1) ** 2)
                dist_to_end = np.sqrt((x - x2) ** 2 + (y - y2) ** 2)
                return min(dist_to_start, dist_to_end) <= tolerance * 2

            return dist_to_line <= tolerance

        # Check if intersection is reasonably close to both line segments
        if not (
            point_on_segment(px, py, x1, y1, x2, y2)
            and point_on_segment(px, py, x3, y3, x4, y4)
        ):
            # If intersection is far from segments, it's probably extrapolating too much
            raise ValueError("Intersection point too far from line segments")

        return [px, py]

    def join_ring(merged_df):
        """
        Join line segments to form a closed ring.

        Args:
            merged_df (pandas.DataFrame): DataFrame with merged line segments

        Returns:
            numpy.ndarray: Array of shape (1, n, 2) containing the ring coordinates
        """
        # Handle edge cases
        if len(merged_df) < 3:
            # Not enough segments to form a valid polygon
            return np.array([[]])

        ring = []

        # Find intersections between adjacent line segments
        for i in range(len(merged_df) - 1):
            x1, y1, x2, y2, *_ = merged_df.iloc[i]
            x3, y3, x4, y4, *_ = merged_df.iloc[i + 1]

            try:
                intersection = find_intersection(x1, y1, x2, y2, x3, y3, x4, y4)

                # Check if the intersection point is too far from either line segment
                # This helps prevent extending edges beyond reasonable bounds
                dist_to_seg1 = min(
                    np.sqrt((intersection[0] - x1) ** 2 + (intersection[1] - y1) ** 2),
                    np.sqrt((intersection[0] - x2) ** 2 + (intersection[1] - y2) ** 2),
                )
                dist_to_seg2 = min(
                    np.sqrt((intersection[0] - x3) ** 2 + (intersection[1] - y3) ** 2),
                    np.sqrt((intersection[0] - x4) ** 2 + (intersection[1] - y4) ** 2),
                )

                # Use the maximum of line segment lengths as a reference
                max_len = max(
                    np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2),
                    np.sqrt((x4 - x3) ** 2 + (y4 - y3) ** 2),
                )

                # Improved intersection validation
                # Calculate angle between segments to detect sharp corners
                v1 = np.array([x2 - x1, y2 - y1])
                v2 = np.array([x4 - x3, y4 - y3])
                v1_norm = np.linalg.norm(v1)
                v2_norm = np.linalg.norm(v2)

                if v1_norm > 0 and v2_norm > 0:
                    cos_angle = np.dot(v1, v2) / (v1_norm * v2_norm)
                    cos_angle = np.clip(cos_angle, -1, 1)
                    angle = np.arccos(cos_angle)

                    # Check for very sharp angles that could create triangular artifacts
                    is_sharp_angle = (
                        angle < np.pi / 6 or angle > 5 * np.pi / 6
                    )  # <30° or >150°
                else:
                    is_sharp_angle = False

                # Determine whether to use intersection or segment endpoint
                if (
                    dist_to_seg1 > max_len * 0.5
                    or dist_to_seg2 > max_len * 0.5
                    or is_sharp_angle
                ):
                    # Use a more conservative approach for problematic intersections
                    # Use the closer endpoint between segments
                    dist_x2_to_seg2 = min(
                        np.sqrt((x2 - x3) ** 2 + (y2 - y3) ** 2),
                        np.sqrt((x2 - x4) ** 2 + (y2 - y4) ** 2),
                    )
                    dist_x3_to_seg1 = min(
                        np.sqrt((x3 - x1) ** 2 + (y3 - y1) ** 2),
                        np.sqrt((x3 - x2) ** 2 + (y3 - y2) ** 2),
                    )

                    if dist_x2_to_seg2 <= dist_x3_to_seg1:
                        ring.append([x2, y2])
                    else:
                        ring.append([x3, y3])
                else:
                    ring.append(intersection)
            except Exception:
                # If intersection calculation fails, use the endpoint of the first segment
                ring.append([x2, y2])

        # Connect last segment with first segment
        x1, y1, x2, y2, *_ = merged_df.iloc[-1]
        x3, y3, x4, y4, *_ = merged_df.iloc[0]

        try:
            intersection = find_intersection(x1, y1, x2, y2, x3, y3, x4, y4)

            # Check if the intersection point is too far from either line segment
            dist_to_seg1 = min(
                np.sqrt((intersection[0] - x1) ** 2 + (intersection[1] - y1) ** 2),
                np.sqrt((intersection[0] - x2) ** 2 + (intersection[1] - y2) ** 2),
            )
            dist_to_seg2 = min(
                np.sqrt((intersection[0] - x3) ** 2 + (intersection[1] - y3) ** 2),
                np.sqrt((intersection[0] - x4) ** 2 + (intersection[1] - y4) ** 2),
            )

            max_len = max(
                np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2),
                np.sqrt((x4 - x3) ** 2 + (y4 - y3) ** 2),
            )

            # Apply same sharp angle detection for closing segment
            v1 = np.array([x2 - x1, y2 - y1])
            v2 = np.array([x4 - x3, y4 - y3])
            v1_norm = np.linalg.norm(v1)
            v2_norm = np.linalg.norm(v2)

            if v1_norm > 0 and v2_norm > 0:
                cos_angle = np.dot(v1, v2) / (v1_norm * v2_norm)
                cos_angle = np.clip(cos_angle, -1, 1)
                angle = np.arccos(cos_angle)
                is_sharp_angle = angle < np.pi / 6 or angle > 5 * np.pi / 6
            else:
                is_sharp_angle = False

            if (
                dist_to_seg1 > max_len * 0.5
                or dist_to_seg2 > max_len * 0.5
                or is_sharp_angle
            ):
                # Use conservative approach for closing segment
                dist_x2_to_seg2 = min(
                    np.sqrt((x2 - x3) ** 2 + (y2 - y3) ** 2),
                    np.sqrt((x2 - x4) ** 2 + (y2 - y4) ** 2),
                )
                dist_x3_to_seg1 = min(
                    np.sqrt((x3 - x1) ** 2 + (y3 - y1) ** 2),
                    np.sqrt((x3 - x2) ** 2 + (y3 - y2) ** 2),
                )

                if dist_x2_to_seg2 <= dist_x3_to_seg1:
                    ring.append([x2, y2])
                else:
                    ring.append([x3, y3])
            else:
                ring.append(intersection)
        except Exception:
            # If intersection calculation fails, use the endpoint of the last segment
            ring.append([x2, y2])

        # Ensure the ring is closed
        if len(ring) > 0 and (ring[0][0] != ring[-1][0] or ring[0][1] != ring[-1][1]):
            ring.append(ring[0])

        return np.array([ring])

    def rdp(M, epsilon=0, dist=None, algo="iter", return_mask=False):
        """
        Simplifies a given array of points using the Ramer-Douglas-Peucker algorithm.

        Args:
            M (numpy.ndarray): Array of shape (n, d) containing point coordinates
            epsilon (float, optional): Epsilon value for simplification
            dist (callable, optional): Distance function
            algo (str, optional): Algorithm to use ('iter' or 'rec')
            return_mask (bool, optional): Whether to return a mask instead of the simplified array

        Returns:
            numpy.ndarray or list: Simplified points or mask
        """
        if dist is None:
            dist = pldist

        if algo == "iter":
            algo = partial(rdp_iter, return_mask=return_mask)
        elif algo == "rec":
            if return_mask:
                raise NotImplementedError(
                    'return_mask=True not supported with algo="rec"'
                )
            algo = rdp_rec

        if "numpy" in str(type(M)):
            return algo(M, epsilon, dist)

        return algo(np.array(M), epsilon, dist).tolist()

    def pldist(point, start, end):
        """
        Calculates the distance from 'point' to the line given by 'start' and 'end'.

        Args:
            point (numpy.ndarray): Point coordinates
            start (numpy.ndarray): Start point of the line
            end (numpy.ndarray): End point of the line

        Returns:
            float: Distance from point to line
        """
        if np.all(np.equal(start, end)):
            return np.linalg.norm(point - start)

        # Fix for NumPy 2.0 deprecation warning - handle 2D vectors properly
        # Instead of using cross product directly, calculate the area of the
        # parallelogram formed by the vectors and divide by the length of the line
        line_vec = end - start
        point_vec = point - start

        # Area of parallelogram = |a|*|b|*sin(θ)
        # For 2D vectors: |a×b| = |a|*|b|*sin(θ) = determinant([ax, ay], [bx, by])
        area = abs(line_vec[0] * point_vec[1] - line_vec[1] * point_vec[0])

        # Distance = Area / |line_vec|
        return area / np.linalg.norm(line_vec)

    def rdp_rec(M, epsilon, dist=pldist):
        """
        Recursive implementation of the Ramer-Douglas-Peucker algorithm.

        Args:
            M (numpy.ndarray): Array of shape (n, d) containing point coordinates
            epsilon (float): Epsilon value for simplification
            dist (callable, optional): Distance function

        Returns:
            numpy.ndarray: Simplified points
        """
        dmax = 0.0
        index = -1

        for i in range(1, M.shape[0]):
            d = dist(M[i], M[0], M[-1])

            if d > dmax:
                index = i
                dmax = d

        if dmax > epsilon:
            r1 = rdp_rec(M[: index + 1], epsilon, dist)
            r2 = rdp_rec(M[index:], epsilon, dist)

            return np.vstack((r1[:-1], r2))
        else:
            return np.vstack((M[0], M[-1]))

    def _rdp_iter(M, start_index, last_index, epsilon, dist=pldist):
        """
        Internal iterative implementation of the Ramer-Douglas-Peucker algorithm.

        Args:
            M (numpy.ndarray): Array of shape (n, d) containing point coordinates
            start_index (int): Start index
            last_index (int): Last index
            epsilon (float): Epsilon value for simplification
            dist (callable, optional): Distance function

        Returns:
            numpy.ndarray: Boolean mask of points to keep
        """
        stk = []
        stk.append([start_index, last_index])
        global_start_index = start_index
        indices = np.ones(last_index - start_index + 1, dtype=bool)

        while stk:
            start_index, last_index = stk.pop()

            dmax = 0.0
            index = start_index

            for i in range(index + 1, last_index):
                if indices[i - global_start_index]:
                    d = dist(M[i], M[start_index], M[last_index])
                    if d > dmax:
                        index = i
                        dmax = d

            if dmax > epsilon:
                stk.append([start_index, index])
                stk.append([index, last_index])
            else:
                for i in range(start_index + 1, last_index):
                    indices[i - global_start_index] = False

        return indices

    def rdp_iter(M, epsilon, dist=pldist, return_mask=False):
        """
        Iterative implementation of the Ramer-Douglas-Peucker algorithm.

        Args:
            M (numpy.ndarray): Array of shape (n, d) containing point coordinates
            epsilon (float): Epsilon value for simplification
            dist (callable, optional): Distance function
            return_mask (bool, optional): Whether to return a mask instead of the simplified array

        Returns:
            numpy.ndarray: Simplified points or boolean mask
        """
        mask = _rdp_iter(M, 0, len(M) - 1, epsilon, dist)

        if return_mask:
            return mask

        return M[mask]

    # Read the raster data
    with rasterio.open(input_path) as src:
        # Read the first band (assuming it contains the mask)
        mask = src.read(1)
        transform = src.transform
        crs = src.crs

        # Extract shapes from the raster mask
        shapes = list(features.shapes(mask, transform=transform))

        # Initialize progress bar
        print(f"Processing {len(shapes)} features...")

        # Convert shapes to GeoJSON features
        features_list = []
        for shape, value in tqdm(shapes, desc="Converting features", unit="shape"):
            if value > 0:  # Only process non-zero values (actual objects)
                # Convert GeoJSON geometry to Shapely polygon
                polygon = Polygon(shape["coordinates"][0])

                # Skip tiny polygons
                if polygon.area < min_area:
                    features_list.append(
                        {
                            "type": "Feature",
                            "properties": {"value": int(value)},
                            "geometry": shape,
                        }
                    )
                    continue

                # Check if shape is triangular and if we want to avoid triangular shapes
                if detect_triangles:
                    # Create a simplified version to check number of vertices
                    simple_polygon = polygon.simplify(epsilon)
                    if (
                        len(simple_polygon.exterior.coords) <= 4
                    ):  # 3 points + closing point
                        # Likely a triangular shape - skip orthogonalization
                        features_list.append(
                            {
                                "type": "Feature",
                                "properties": {"value": int(value)},
                                "geometry": shape,
                            }
                        )
                        continue

                # Process larger, non-triangular polygons
                try:
                    # Convert shapely polygon to a ring format for orthogonalization
                    exterior_ring = list(polygon.exterior.coords)
                    interior_rings = [
                        list(interior.coords) for interior in polygon.interiors
                    ]

                    # Calculate bounding box aspect ratio to help with parameter tuning
                    minx, miny, maxx, maxy = polygon.bounds
                    width = maxx - minx
                    height = maxy - miny
                    aspect_ratio = max(width, height) / max(1.0, min(width, height))

                    # Determine if this shape is likely to be a building/rectangular object
                    # Long thin objects might require different treatment
                    is_rectangular = aspect_ratio < 3.0

                    # Rectangular objects usually need more careful orthogonalization
                    epsilon_adjusted = epsilon
                    min_segments_adjusted = min_segments

                    if is_rectangular:
                        # For rectangular objects, use more conservative epsilon
                        epsilon_adjusted = epsilon * 0.75
                        # Ensure we get at least 4 points for a proper rectangle
                        min_segments_adjusted = max(4, min_segments)

                    # Orthogonalize the exterior and interior rings
                    orthogonalized_exterior = orthogonalize_ring(
                        exterior_ring,
                        epsilon=epsilon_adjusted,
                        min_segments=min_segments_adjusted,
                    )

                    orthogonalized_interiors = [
                        orthogonalize_ring(
                            ring,
                            epsilon=epsilon_adjusted,
                            min_segments=min_segments_adjusted,
                        )
                        for ring in interior_rings
                    ]

                    # Validate the result - calculate area change
                    original_area = polygon.area
                    orthogonalized_poly = Polygon(orthogonalized_exterior)

                    if orthogonalized_poly.is_valid:
                        area_ratio = (
                            orthogonalized_poly.area / original_area
                            if original_area > 0
                            else 0
                        )

                        # If area changed too much, revert to original
                        if area_ratio < area_tolerance or area_ratio > (
                            1.0 / area_tolerance
                        ):
                            # Use original polygon instead
                            geometry = shape
                        else:
                            # Create a new geometry with orthogonalized rings
                            geometry = {
                                "type": "Polygon",
                                "coordinates": [orthogonalized_exterior],
                            }

                            # Add interior rings if they exist
                            if orthogonalized_interiors:
                                geometry["coordinates"].extend(
                                    [ring for ring in orthogonalized_interiors]
                                )
                    else:
                        # If resulting polygon is invalid, use original
                        geometry = shape

                    # Add the feature to the list
                    features_list.append(
                        {
                            "type": "Feature",
                            "properties": {"value": int(value)},
                            "geometry": geometry,
                        }
                    )
                except Exception as e:
                    # Keep the original shape if orthogonalization fails
                    features_list.append(
                        {
                            "type": "Feature",
                            "properties": {"value": int(value)},
                            "geometry": shape,
                        }
                    )

        # Create the final GeoJSON structure
        geojson = {
            "type": "FeatureCollection",
            "crs": {"type": "name", "properties": {"name": str(crs)}},
            "features": features_list,
        }

        # Convert to GeoDataFrame and set the CRS
        gdf = gpd.GeoDataFrame.from_features(geojson["features"], crs=crs)

        # Save to file if output_path is provided
        if output_path:
            print(f"Saving to {output_path}...")
            gdf.to_file(output_path)
            print("Done!")

        return gdf

print_raster_info(raster_path, show_preview=True, figsize=(10, 8))

Print formatted information about a raster dataset and optionally show a preview.

Parameters:

Name	Type	Description	Default
raster_path	str	Path to the raster file	required
show_preview	bool	Whether to display a visual preview of the raster. Defaults to True.	True
figsize	tuple	Figure size as (width, height). Defaults to (10, 8).	(10, 8)

Returns:

Name	Type	Description
dict	Optional[Dict[str, Any]]	Dictionary containing raster information if successful, None otherwise

Source code in geoai/utils/raster.py

def print_raster_info(
    raster_path: str, show_preview: bool = True, figsize: Tuple[int, int] = (10, 8)
) -> Optional[Dict[str, Any]]:
    """Print formatted information about a raster dataset and optionally show a preview.

    Args:
        raster_path (str): Path to the raster file
        show_preview (bool, optional): Whether to display a visual preview of the raster.
            Defaults to True.
        figsize (tuple, optional): Figure size as (width, height). Defaults to (10, 8).

    Returns:
        dict: Dictionary containing raster information if successful, None otherwise
    """
    try:
        info = get_raster_info(raster_path)

        # Print basic information
        print(f"===== RASTER INFORMATION: {raster_path} =====")
        print(f"Driver: {info['driver']}")
        print(f"Dimensions: {info['width']} x {info['height']} pixels")
        print(f"Number of bands: {info['count']}")
        print(f"Data type: {info['dtype']}")
        print(f"Coordinate Reference System: {info['crs']}")
        print(f"Georeferenced Bounds: {info['bounds']}")
        print(f"Pixel Resolution: {info['resolution'][0]}, {info['resolution'][1]}")
        print(f"NoData Value: {info['nodata']}")

        # Print band statistics
        print("\n----- Band Statistics -----")
        for band_stat in info["band_stats"]:
            print(f"Band {band_stat['band']}:")
            print(f"  Min: {band_stat['min']:.2f}")
            print(f"  Max: {band_stat['max']:.2f}")
            print(f"  Mean: {band_stat['mean']:.2f}")
            print(f"  Std Dev: {band_stat['std']:.2f}")

        # Show a preview if requested
        if show_preview:
            with rasterio.open(raster_path) as src:
                # For multi-band images, show RGB composite or first band
                if src.count >= 3:
                    # Try to show RGB composite
                    rgb = np.dstack([src.read(i) for i in range(1, 4)])
                    plt.figure(figsize=figsize)
                    plt.imshow(rgb)
                    plt.title(f"RGB Preview: {raster_path}")
                else:
                    # Show first band for single-band images
                    plt.figure(figsize=figsize)
                    show(
                        src.read(1),
                        cmap="viridis",
                        title=f"Band 1 Preview: {raster_path}",
                    )
                    plt.colorbar(label="Pixel Value")
                plt.show()

    except Exception as e:
        print(f"Error reading raster: {str(e)}")

print_vector_info(vector_path, show_preview=True, figsize=(10, 8))

Print formatted information about a vector dataset and optionally show a preview.

Parameters:

Name	Type	Description	Default
vector_path	str	Path to the vector file	required
show_preview	bool	Whether to display a visual preview of the vector data. Defaults to True.	True
figsize	tuple	Figure size as (width, height). Defaults to (10, 8).	(10, 8)

Returns:

Name	Type	Description
dict	Optional[Dict[str, Any]]	Dictionary containing vector information if successful, None otherwise

Source code in geoai/utils/vector.py

def print_vector_info(
    vector_path: str, show_preview: bool = True, figsize: Tuple[int, int] = (10, 8)
) -> Optional[Dict[str, Any]]:
    """Print formatted information about a vector dataset and optionally show a preview.

    Args:
        vector_path (str): Path to the vector file
        show_preview (bool, optional): Whether to display a visual preview of the vector data.
            Defaults to True.
        figsize (tuple, optional): Figure size as (width, height). Defaults to (10, 8).

    Returns:
        dict: Dictionary containing vector information if successful, None otherwise
    """
    try:
        info = get_vector_info(vector_path)

        # Print basic information
        print(f"===== VECTOR INFORMATION: {vector_path} =====")
        print(f"Driver: {info['driver']}")
        print(f"Feature count: {info['feature_count']}")
        print(f"Geometry types: {info['geometry_type']}")
        print(f"Coordinate Reference System: {info['crs']}")
        print(f"Bounds: {info['bounds']}")
        print(f"Number of attributes: {info['attribute_count']}")
        print(f"Attribute names: {', '.join(info['attribute_names'])}")

        # Print attribute statistics
        if info["attribute_stats"]:
            print("\n----- Attribute Statistics -----")
            for attr, stats in info["attribute_stats"].items():
                print(f"Attribute: {attr}")
                for stat_name, stat_value in stats.items():
                    print(
                        f"  {stat_name}: {stat_value:.4f}"
                        if isinstance(stat_value, float)
                        else f"  {stat_name}: {stat_value}"
                    )

        # Show a preview if requested
        if show_preview:
            gdf = (
                gpd.read_parquet(vector_path)
                if vector_path.endswith(".parquet")
                else gpd.read_file(vector_path)
            )
            fig, ax = plt.subplots(figsize=figsize)
            gdf.plot(ax=ax, cmap="viridis")
            ax.set_title(f"Preview: {vector_path}")
            plt.tight_layout()
            plt.show()

            # # Show a sample of the attribute table
            # if not gdf.empty:
            #     print("\n----- Sample of attribute table (first 5 rows) -----")
            #     print(gdf.head().to_string())

    except Exception as e:
        print(f"Error reading vector data: {str(e)}")

raster_to_vector(raster_path, output_path=None, threshold=0, min_area=10, simplify_tolerance=None, class_values=None, attribute_name='class', unique_attribute_value=False, output_format='geojson', plot_result=False)

Convert a raster label mask to vector polygons.

Parameters:

Name	Type	Description	Default
raster_path	str	Path to the input raster file (e.g., GeoTIFF).	required
output_path	str	Path to save the output vector file. If None, returns GeoDataFrame without saving.	None
threshold	int / float	Pixel values greater than this threshold will be vectorized.	0
min_area	float	Minimum polygon area in square map units to keep.	10
simplify_tolerance	float	Tolerance for geometry simplification. None for no simplification.	None
class_values	list	Specific pixel values to vectorize. If None, all values > threshold are vectorized.	None
attribute_name	str	Name of the attribute field for the class values.	'class'
unique_attribute_value	bool	Whether to generate unique values for each shape within a class.	False
output_format	str	Format for output file - 'geojson', 'shapefile', 'gpkg'. Auto-detected from the file extension of output_path when possible.	'geojson'
plot_result	bool	Whether to plot the resulting polygons overlaid on the raster.	False

Returns:

Type	Description
GeoDataFrame	geopandas.GeoDataFrame: A GeoDataFrame containing the vectorized polygons.

Source code in geoai/utils/raster.py

def raster_to_vector(
    raster_path: str,
    output_path: Optional[str] = None,
    threshold: float = 0,
    min_area: float = 10,
    simplify_tolerance: Optional[float] = None,
    class_values: Optional[List[int]] = None,
    attribute_name: str = "class",
    unique_attribute_value: bool = False,
    output_format: str = "geojson",
    plot_result: bool = False,
) -> gpd.GeoDataFrame:
    """
    Convert a raster label mask to vector polygons.

    Args:
        raster_path (str): Path to the input raster file (e.g., GeoTIFF).
        output_path (str): Path to save the output vector file. If None, returns GeoDataFrame without saving.
        threshold (int/float): Pixel values greater than this threshold will be vectorized.
        min_area (float): Minimum polygon area in square map units to keep.
        simplify_tolerance (float): Tolerance for geometry simplification. None for no simplification.
        class_values (list): Specific pixel values to vectorize. If None, all values > threshold are vectorized.
        attribute_name (str): Name of the attribute field for the class values.
        unique_attribute_value (bool): Whether to generate unique values for each shape within a class.
        output_format (str): Format for output file - 'geojson', 'shapefile', 'gpkg'.
            Auto-detected from the file extension of output_path when possible.
        plot_result (bool): Whether to plot the resulting polygons overlaid on the raster.

    Returns:
        geopandas.GeoDataFrame: A GeoDataFrame containing the vectorized polygons.
    """
    # Open the raster file
    with rasterio.open(raster_path) as src:
        # Read the data
        data = src.read(1)

        # Get metadata
        transform = src.transform
        crs = src.crs

        # Check if CRS is geographic (area in sq degrees, not sq meters)
        is_geographic = crs is not None and crs.is_geographic

        # Initialize list to store features
        all_features = []

        if class_values is not None:
            # Create a mask for each specified class value and vectorize per class
            masks = {val: (data == val) for val in class_values}
            for class_val in class_values:
                mask = masks[class_val]
                shape_count = 1
                for geom, value in features.shapes(
                    mask.astype(np.uint8), mask=mask, transform=transform
                ):
                    geom = shape(geom)
                    if not is_geographic and geom.area < min_area:
                        continue
                    if simplify_tolerance is not None:
                        geom = geom.simplify(simplify_tolerance)
                    if unique_attribute_value:
                        all_features.append(
                            {
                                "geometry": geom,
                                attribute_name: class_val * shape_count,
                            }
                        )
                    else:
                        all_features.append(
                            {"geometry": geom, attribute_name: class_val}
                        )
                    shape_count += 1
        else:
            # No class_values specified: vectorize all pixels above threshold.
            # Pass the raw data array (not a binary mask) so that adjacent regions
            # with different pixel values (e.g. SAM unique-value masks) are kept as
            # separate polygons rather than merged into one connected blob.
            binary_mask = (data > threshold).astype(np.uint8)
            shape_count = 1
            for geom, value in features.shapes(
                data.astype(np.int32), mask=binary_mask, transform=transform
            ):
                class_val = int(value)
                if class_val == 0:
                    continue
                geom = shape(geom)
                if not is_geographic and geom.area < min_area:
                    continue
                if simplify_tolerance is not None:
                    geom = geom.simplify(simplify_tolerance)
                if unique_attribute_value:
                    all_features.append(
                        {"geometry": geom, attribute_name: class_val * shape_count}
                    )
                else:
                    all_features.append({"geometry": geom, attribute_name: class_val})
                shape_count += 1

        # Create GeoDataFrame
        if all_features:
            gdf = gpd.GeoDataFrame(all_features, crs=crs)
        else:
            print("Warning: No features were extracted from the raster.")
            # Return empty GeoDataFrame with correct CRS
            gdf = gpd.GeoDataFrame([], geometry=[], crs=crs)

        # For geographic CRS, filter by area using a projected CRS
        if is_geographic and min_area > 0 and len(gdf) > 0:
            utm_crs = gdf.estimate_utm_crs()
            projected_area = gdf.to_crs(utm_crs).geometry.area
            gdf = gdf[projected_area >= min_area].reset_index(drop=True)

        # Save to file if requested
        if output_path is not None:
            # Auto-detect output_format from file extension when not explicitly set
            ext = os.path.splitext(output_path)[1].lower()
            if ext == ".gpkg":
                output_format = "gpkg"
            elif ext == ".shp":
                output_format = "shapefile"
            elif ext in (".geojson", ".json"):
                output_format = "geojson"

            # Create directory if it doesn't exist
            os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)

            # Save to file based on format
            if output_format.lower() == "geojson":
                gdf.to_file(output_path, driver="GeoJSON")
            elif output_format.lower() == "shapefile":
                gdf.to_file(output_path)
            elif output_format.lower() == "gpkg":
                gdf.to_file(output_path, driver="GPKG")
            else:
                raise ValueError(f"Unsupported output format: {output_format}")

            print(f"Vectorized data saved to {output_path}")

        # Plot result if requested
        if plot_result:
            fig, ax = plt.subplots(figsize=(12, 12))

            # Plot raster
            raster_img = src.read()
            if raster_img.shape[0] == 1:
                plt.imshow(raster_img[0], cmap="viridis", alpha=0.7)
            else:
                # Use first 3 bands for RGB display
                rgb = raster_img[:3].transpose(1, 2, 0)
                # Normalize for display
                rgb = np.clip(rgb / rgb.max(), 0, 1)
                plt.imshow(rgb)

            # Plot vector boundaries
            if not gdf.empty:
                gdf.plot(ax=ax, facecolor="none", edgecolor="red", linewidth=2)

            plt.title("Raster with Vectorized Boundaries")
            plt.axis("off")
            plt.tight_layout()
            plt.show()

        return gdf

raster_to_vector_batch(input_dir, output_dir, pattern='*.tif', threshold=0, min_area=10, simplify_tolerance=None, class_values=None, attribute_name='class', output_format='geojson', merge_output=False, merge_filename='merged_vectors')

Batch convert multiple raster files to vector polygons.

Parameters:

Name	Type	Description	Default
input_dir	str	Directory containing input raster files.	required
output_dir	str	Directory to save output vector files.	required
pattern	str	Pattern to match raster files (e.g., '*.tif').	'*.tif'
threshold	int / float	Pixel values greater than this threshold will be vectorized.	0
min_area	float	Minimum polygon area in square map units to keep.	10
simplify_tolerance	float	Tolerance for geometry simplification. None for no simplification.	None
class_values	list	Specific pixel values to vectorize. If None, all values > threshold are vectorized.	None
attribute_name	str	Name of the attribute field for the class values.	'class'
output_format	str	Format for output files - 'geojson', 'shapefile', 'gpkg'.	'geojson'
merge_output	bool	Whether to merge all output vectors into a single file.	False
merge_filename	str	Filename for the merged output (without extension).	'merged_vectors'

Returns:

Type	Description
Optional[GeoDataFrame]	geopandas.GeoDataFrame or None: If merge_output is True, returns the merged GeoDataFrame.

Source code in geoai/utils/raster.py

def raster_to_vector_batch(
    input_dir: str,
    output_dir: str,
    pattern: str = "*.tif",
    threshold: float = 0,
    min_area: float = 10,
    simplify_tolerance: Optional[float] = None,
    class_values: Optional[List[int]] = None,
    attribute_name: str = "class",
    output_format: str = "geojson",
    merge_output: bool = False,
    merge_filename: str = "merged_vectors",
) -> Optional[gpd.GeoDataFrame]:
    """
    Batch convert multiple raster files to vector polygons.

    Args:
        input_dir (str): Directory containing input raster files.
        output_dir (str): Directory to save output vector files.
        pattern (str): Pattern to match raster files (e.g., '*.tif').
        threshold (int/float): Pixel values greater than this threshold will be vectorized.
        min_area (float): Minimum polygon area in square map units to keep.
        simplify_tolerance (float): Tolerance for geometry simplification. None for no simplification.
        class_values (list): Specific pixel values to vectorize. If None, all values > threshold are vectorized.
        attribute_name (str): Name of the attribute field for the class values.
        output_format (str): Format for output files - 'geojson', 'shapefile', 'gpkg'.
        merge_output (bool): Whether to merge all output vectors into a single file.
        merge_filename (str): Filename for the merged output (without extension).

    Returns:
        geopandas.GeoDataFrame or None: If merge_output is True, returns the merged GeoDataFrame.
    """
    import glob

    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)

    # Get list of raster files
    raster_files = glob.glob(os.path.join(input_dir, pattern))

    if not raster_files:
        print(f"No files matching pattern '{pattern}' found in {input_dir}")
        return None

    print(f"Found {len(raster_files)} raster files to process")

    # Process each raster file
    gdfs = []
    for raster_file in tqdm(raster_files, desc="Processing rasters"):
        # Get output filename
        base_name = os.path.splitext(os.path.basename(raster_file))[0]
        if output_format.lower() == "geojson":
            out_file = os.path.join(output_dir, f"{base_name}.geojson")
        elif output_format.lower() == "shapefile":
            out_file = os.path.join(output_dir, f"{base_name}.shp")
        elif output_format.lower() == "gpkg":
            out_file = os.path.join(output_dir, f"{base_name}.gpkg")
        else:
            raise ValueError(f"Unsupported output format: {output_format}")

        # Convert raster to vector
        if merge_output:
            # Don't save individual files if merging
            gdf = raster_to_vector(
                raster_file,
                output_path=None,
                threshold=threshold,
                min_area=min_area,
                simplify_tolerance=simplify_tolerance,
                class_values=class_values,
                attribute_name=attribute_name,
            )

            # Add filename as attribute
            if not gdf.empty:
                gdf["source_file"] = base_name
                gdfs.append(gdf)
        else:
            # Save individual files
            raster_to_vector(
                raster_file,
                output_path=out_file,
                threshold=threshold,
                min_area=min_area,
                simplify_tolerance=simplify_tolerance,
                class_values=class_values,
                attribute_name=attribute_name,
                output_format=output_format,
            )

    # Merge output if requested
    if merge_output and gdfs:
        merged_gdf = gpd.GeoDataFrame(pd.concat(gdfs, ignore_index=True))

        # Set CRS to the CRS of the first GeoDataFrame
        if merged_gdf.crs is None and gdfs:
            merged_gdf.crs = gdfs[0].crs

        # Save merged output
        if output_format.lower() == "geojson":
            merged_file = os.path.join(output_dir, f"{merge_filename}.geojson")
            merged_gdf.to_file(merged_file, driver="GeoJSON")
        elif output_format.lower() == "shapefile":
            merged_file = os.path.join(output_dir, f"{merge_filename}.shp")
            merged_gdf.to_file(merged_file)
        elif output_format.lower() == "gpkg":
            merged_file = os.path.join(output_dir, f"{merge_filename}.gpkg")
            merged_gdf.to_file(merged_file, driver="GPKG")

        print(f"Merged vector data saved to {merged_file}")
        return merged_gdf

    return None

read_raster(source, band=None, masked=True, **kwargs)

Reads raster data from various formats using rioxarray.

This function reads raster data from local files or URLs into a rioxarray data structure with preserved geospatial metadata.

Parameters:

Name	Type	Description	Default
source	str	String path to the raster file or URL.	required
band	Optional[Union[int, List[int]]]	Integer or list of integers specifying which band(s) to read. Defaults to None (all bands).	None
masked	bool	Boolean indicating whether to mask nodata values. Defaults to True.	True
**kwargs	Any	Additional keyword arguments to pass to rioxarray.open_rasterio.	{}

Returns:

Type	Description
DataArray	xarray.DataArray: A DataArray containing the raster data with geospatial metadata preserved.

Raises:

Type	Description
ValueError	If the file format is not supported or source cannot be accessed.

Examples:

Read a local GeoTIFF

>>> raster = read_raster("path/to/data.tif")
>>>
Read only band 1 from a remote GeoTIFF
>>> raster = read_raster("https://example.com/data.tif", band=1)
>>>
Read a raster without masking nodata values
>>> raster = read_raster("path/to/data.tif", masked=False)

Source code in geoai/utils/raster.py

def read_raster(
    source: str,
    band: Optional[Union[int, List[int]]] = None,
    masked: bool = True,
    **kwargs: Any,
) -> xr.DataArray:
    """Reads raster data from various formats using rioxarray.

    This function reads raster data from local files or URLs into a rioxarray
    data structure with preserved geospatial metadata.

    Args:
        source: String path to the raster file or URL.
        band: Integer or list of integers specifying which band(s) to read.
            Defaults to None (all bands).
        masked: Boolean indicating whether to mask nodata values.
            Defaults to True.
        **kwargs: Additional keyword arguments to pass to rioxarray.open_rasterio.

    Returns:
        xarray.DataArray: A DataArray containing the raster data with geospatial
            metadata preserved.

    Raises:
        ValueError: If the file format is not supported or source cannot be accessed.

    Examples:
        Read a local GeoTIFF
        >>> raster = read_raster("path/to/data.tif")
        >>>
        Read only band 1 from a remote GeoTIFF
        >>> raster = read_raster("https://example.com/data.tif", band=1)
        >>>
        Read a raster without masking nodata values
        >>> raster = read_raster("path/to/data.tif", masked=False)
    """
    import urllib.parse

    from rasterio.errors import RasterioIOError

    # Determine if source is a URL or local file
    parsed_url = urllib.parse.urlparse(source)
    is_url = parsed_url.scheme in ["http", "https"]

    # If it's a local file, check if it exists
    if not is_url and not os.path.exists(source):
        raise ValueError(f"Raster file does not exist: {source}")

    try:
        # Open the raster with rioxarray
        raster = rxr.open_rasterio(source, masked=masked, **kwargs)

        # Handle band selection if specified
        if band is not None:
            if isinstance(band, (list, tuple)):
                # Convert from 1-based indexing to 0-based indexing
                band_indices = [b - 1 for b in band]
                raster = raster.isel(band=band_indices)
            else:
                # Single band selection (convert from 1-based to 0-based indexing)
                raster = raster.isel(band=band - 1)

        return raster

    except RasterioIOError as e:
        raise ValueError(f"Could not read raster from source '{source}': {str(e)}")
    except Exception as e:
        raise ValueError(f"Error reading raster data: {str(e)}")

read_vector(source, layer=None, **kwargs)

Reads vector data from various formats including GeoParquet.

This function dynamically determines the file type based on extension and reads it into a GeoDataFrame. It supports both local files and HTTP/HTTPS URLs.

Parameters:

Name	Type	Description	Default
source	str	String path to the vector file or URL.	required
layer	Optional[str]	String or integer specifying which layer to read from multi-layer files (only applicable for formats like GPKG, GeoJSON, etc.). Defaults to None.	None
**kwargs	Any	Additional keyword arguments to pass to the underlying reader.	{}

Returns:

Type	Description
GeoDataFrame	geopandas.GeoDataFrame: A GeoDataFrame containing the vector data.

Raises:

Type	Description
ValueError	If the file format is not supported or source cannot be accessed.

Examples:

Read a local shapefile

>>> gdf = read_vector("path/to/data.shp")
>>>
Read a GeoParquet file from URL
>>> gdf = read_vector("https://example.com/data.parquet")
>>>
Read a specific layer from a GeoPackage
>>> gdf = read_vector("path/to/data.gpkg", layer="layer_name")

Source code in geoai/utils/raster.py

def read_vector(
    source: str, layer: Optional[str] = None, **kwargs: Any
) -> gpd.GeoDataFrame:
    """Reads vector data from various formats including GeoParquet.

    This function dynamically determines the file type based on extension
    and reads it into a GeoDataFrame. It supports both local files and HTTP/HTTPS URLs.

    Args:
        source: String path to the vector file or URL.
        layer: String or integer specifying which layer to read from multi-layer
            files (only applicable for formats like GPKG, GeoJSON, etc.).
            Defaults to None.
        **kwargs: Additional keyword arguments to pass to the underlying reader.

    Returns:
        geopandas.GeoDataFrame: A GeoDataFrame containing the vector data.

    Raises:
        ValueError: If the file format is not supported or source cannot be accessed.

    Examples:
        Read a local shapefile
        >>> gdf = read_vector("path/to/data.shp")
        >>>
        Read a GeoParquet file from URL
        >>> gdf = read_vector("https://example.com/data.parquet")
        >>>
        Read a specific layer from a GeoPackage
        >>> gdf = read_vector("path/to/data.gpkg", layer="layer_name")
    """

    import urllib.parse

    # Determine if source is a URL or local file
    parsed_url = urllib.parse.urlparse(source)
    is_url = parsed_url.scheme in ["http", "https"]

    # If it's a local file, check if it exists
    if not is_url and not os.path.exists(source):
        raise ValueError(f"File does not exist: {source}")

    # Get file extension
    _, ext = os.path.splitext(source)
    ext = ext.lower()

    # Handle GeoParquet files
    if ext in [".parquet", ".pq", ".geoparquet"]:
        return gpd.read_parquet(source, **kwargs)

    # Handle common vector formats
    if ext in [".shp", ".geojson", ".json", ".gpkg", ".gml", ".kml", ".gpx"]:
        # For formats that might have multiple layers
        if ext in [".gpkg", ".gml"] and layer is not None:
            return gpd.read_file(source, layer=layer, **kwargs)
        return gpd.read_file(source, **kwargs)

    # Try to identify valid layers for formats that might have them
    # Only attempt this for local files
    if layer is None and ext in [".gpkg", ".gml"] and not is_url:
        try:
            try:
                import pyogrio

                layers = pyogrio.list_layers(source)[:, 0].tolist()
            except ImportError:
                import fiona

                layers = fiona.listlayers(source)
            if layers:
                return gpd.read_file(source, layer=layers[0], **kwargs)
        except Exception:
            # If listing layers fails, we'll fall through to the generic read attempt
            pass

    # For other formats or when layer listing fails, attempt to read using GeoPandas
    try:
        return gpd.read_file(source, **kwargs)
    except Exception as e:
        raise ValueError(f"Could not read from source '{source}': {str(e)}")

region_groups(image, connectivity=1, min_size=10, max_size=None, threshold=None, properties=None, intensity_image=None, out_csv=None, out_vector=None, out_image=None, **kwargs)

Segment regions in an image and filter them based on size.

Parameters:

Name	Type	Description	Default
image	Union[str, DataArray, ndarray]	Input image, can be a file path, xarray DataArray, or numpy array.	required
connectivity	int	Connectivity for labeling. Defaults to 1 for 4-connectivity. Use 2 for 8-connectivity.	1
min_size	int	Minimum size of regions to keep. Defaults to 10.	10
max_size	Optional[int]	Maximum size of regions to keep. Defaults to None.	None
threshold	Optional[int]	Threshold for filling holes. Defaults to None, which is equal to min_size.	None
properties	Optional[List[str]]	List of properties to measure. See https://scikit-image.org/docs/stable/api/skimage.measure.html#skimage.measure.regionprops Defaults to None.	None
intensity_image	Optional[Union[str, DataArray, ndarray]]	Intensity image to measure properties. Defaults to None.	None
out_csv	Optional[str]	Path to save the properties as a CSV file. Defaults to None.	None
out_vector	Optional[str]	Path to save the vector file. Defaults to None.	None
out_image	Optional[str]	Path to save the output image. Defaults to None.	None

Returns:

Type	Description
Union[Tuple[ndarray, DataFrame], Tuple[DataArray, DataFrame]]	Union[Tuple[np.ndarray, pd.DataFrame], Tuple[xr.DataArray, pd.DataFrame]]: Labeled image and properties DataFrame.

Source code in geoai/utils/geometry.py

def region_groups(
    image: Union[str, "xr.DataArray", np.ndarray],
    connectivity: int = 1,
    min_size: int = 10,
    max_size: Optional[int] = None,
    threshold: Optional[int] = None,
    properties: Optional[List[str]] = None,
    intensity_image: Optional[Union[str, "xr.DataArray", np.ndarray]] = None,
    out_csv: Optional[str] = None,
    out_vector: Optional[str] = None,
    out_image: Optional[str] = None,
    **kwargs: Any,
) -> Union[Tuple[np.ndarray, "pd.DataFrame"], Tuple["xr.DataArray", "pd.DataFrame"]]:
    """
    Segment regions in an image and filter them based on size.

    Args:
        image (Union[str, xr.DataArray, np.ndarray]): Input image, can be a file
            path, xarray DataArray, or numpy array.
        connectivity (int, optional): Connectivity for labeling. Defaults to 1
            for 4-connectivity. Use 2 for 8-connectivity.
        min_size (int, optional): Minimum size of regions to keep. Defaults to 10.
        max_size (Optional[int], optional): Maximum size of regions to keep.
            Defaults to None.
        threshold (Optional[int], optional): Threshold for filling holes.
            Defaults to None, which is equal to min_size.
        properties (Optional[List[str]], optional): List of properties to measure.
            See https://scikit-image.org/docs/stable/api/skimage.measure.html#skimage.measure.regionprops
            Defaults to None.
        intensity_image (Optional[Union[str, xr.DataArray, np.ndarray]], optional):
            Intensity image to measure properties. Defaults to None.
        out_csv (Optional[str], optional): Path to save the properties as a CSV file.
            Defaults to None.
        out_vector (Optional[str], optional): Path to save the vector file.
            Defaults to None.
        out_image (Optional[str], optional): Path to save the output image.
            Defaults to None.

    Returns:
        Union[Tuple[np.ndarray, pd.DataFrame], Tuple[xr.DataArray, pd.DataFrame]]: Labeled image and properties DataFrame.
    """
    import scipy.ndimage as ndi
    from skimage import measure

    # Import raster_to_vector lazily to avoid circular imports
    from geoai.utils import raster_to_vector

    if isinstance(image, str):
        ds = rxr.open_rasterio(image)
        da = ds.sel(band=1)
        array = da.values.squeeze()
    elif isinstance(image, xr.DataArray):
        da = image
        array = image.values.squeeze()
    elif isinstance(image, np.ndarray):
        array = image
    else:
        raise ValueError(
            "The input image must be a file path, xarray DataArray, or numpy array."
        )

    if threshold is None:
        threshold = min_size

    # Define a custom function to calculate median intensity
    def intensity_median(region, intensity_image):
        # Extract the intensity values for the region
        return np.median(intensity_image[region])

    # Add your custom function to the list of extra properties
    if intensity_image is not None:
        extra_props = (intensity_median,)
    else:
        extra_props = None

    if properties is None:
        properties = [
            "label",
            "area",
            "area_bbox",
            "area_convex",
            "area_filled",
            "axis_major_length",
            "axis_minor_length",
            "eccentricity",
            "diameter_areagth",
            "extent",
            "orientation",
            "perimeter",
            "solidity",
        ]

        if intensity_image is not None:

            properties += [
                "intensity_max",
                "intensity_mean",
                "intensity_min",
                "intensity_std",
            ]

    if intensity_image is not None:
        if isinstance(intensity_image, str):
            ds = rxr.open_rasterio(intensity_image)
            intensity_da = ds.sel(band=1)
            intensity_image = intensity_da.values.squeeze()
        elif isinstance(intensity_image, xr.DataArray):
            intensity_image = intensity_image.values.squeeze()
        elif isinstance(intensity_image, np.ndarray):
            pass
        else:
            raise ValueError(
                "The intensity_image must be a file path, xarray DataArray, or numpy array."
            )

    label_image = measure.label(array, connectivity=connectivity)
    props = measure.regionprops_table(
        label_image, properties=properties, intensity_image=intensity_image, **kwargs
    )

    df = pd.DataFrame(props)

    # Get the labels of regions with area smaller than the threshold
    small_regions = df[df["area"] < min_size]["label"].values
    # Set the corresponding labels in the label_image to zero
    for region_label in small_regions:
        label_image[label_image == region_label] = 0

    if max_size is not None:
        large_regions = df[df["area"] > max_size]["label"].values
        for region_label in large_regions:
            label_image[label_image == region_label] = 0

    # Find the background (holes) which are zeros
    holes = label_image == 0

    # Label the holes (connected components in the background)
    labeled_holes, _ = ndi.label(holes)

    # Measure properties of the labeled holes, including area and bounding box
    hole_props = measure.regionprops(labeled_holes)

    # Loop through each hole and fill it if it is smaller than the threshold
    for prop in hole_props:
        if prop.area < threshold:
            # Get the coordinates of the small hole
            coords = prop.coords

            # Find the surrounding region's ID (non-zero value near the hole)
            surrounding_region_values = []
            for coord in coords:
                x, y = coord
                # Get a 3x3 neighborhood around the hole pixel
                neighbors = label_image[max(0, x - 1) : x + 2, max(0, y - 1) : y + 2]
                # Exclude the hole pixels (zeros) and get region values
                region_values = neighbors[neighbors != 0]
                if region_values.size > 0:
                    surrounding_region_values.append(
                        region_values[0]
                    )  # Take the first non-zero value

            if surrounding_region_values:
                # Fill the hole with the mode (most frequent) of the surrounding region values
                fill_value = max(
                    set(surrounding_region_values), key=surrounding_region_values.count
                )
                label_image[coords[:, 0], coords[:, 1]] = fill_value

    label_image, num_labels = measure.label(
        label_image, connectivity=connectivity, return_num=True
    )
    props = measure.regionprops_table(
        label_image,
        properties=properties,
        intensity_image=intensity_image,
        extra_properties=extra_props,
        **kwargs,
    )

    df = pd.DataFrame(props)
    df["elongation"] = df["axis_major_length"] / df["axis_minor_length"]

    dtype = "uint8"
    if num_labels > 255 and num_labels <= 65535:
        dtype = "uint16"
    elif num_labels > 65535:
        dtype = "uint32"

    if out_csv is not None:
        df.to_csv(out_csv, index=False)

    if isinstance(image, np.ndarray):
        return label_image, df
    else:
        da.values = label_image
        if out_image is not None:
            da.rio.to_raster(out_image, dtype=dtype)

        if out_vector is not None:
            tmp_raster = None
            tmp_vector = None
            try:
                if out_image is None:
                    tmp_raster = temp_file_path(".tif")
                    da.rio.to_raster(tmp_raster, dtype=dtype)
                    tmp_vector = temp_file_path(".gpkg")
                    raster_to_vector(
                        tmp_raster,
                        tmp_vector,
                        attribute_name="value",
                        unique_attribute_value=True,
                    )
                else:
                    tmp_vector = temp_file_path(".gpkg")
                    raster_to_vector(
                        out_image,
                        tmp_vector,
                        attribute_name="value",
                        unique_attribute_value=True,
                    )
                gdf = gpd.read_file(tmp_vector)
                gdf["label"] = gdf["value"].astype(int)
                gdf.drop(columns=["value"], inplace=True)
                gdf2 = pd.merge(gdf, df, on="label", how="left")
                gdf2.to_file(out_vector)
                gdf2.sort_values("label", inplace=True)
                df = gdf2
            finally:
                try:
                    if tmp_raster is not None and os.path.exists(tmp_raster):
                        os.remove(tmp_raster)
                    if tmp_vector is not None and os.path.exists(tmp_vector):
                        os.remove(tmp_vector)
                except Exception as e:
                    print(f"Warning: Failed to delete temporary files: {str(e)}")

        return da, df

regularization(building_polygons, angle_tolerance=10, simplify_tolerance=0.5, orthogonalize=True, preserve_topology=True)

Regularizes building footprint polygons with multiple techniques beyond minimum rotated rectangles.

Parameters:

Name	Type	Description	Default
building_polygons	Union[GeoDataFrame, List[Polygon]]	GeoDataFrame or list of shapely Polygons containing building footprints	required
angle_tolerance	float	Degrees within which angles will be regularized to 90/180 degrees	10
simplify_tolerance	float	Distance tolerance for Douglas-Peucker simplification	0.5
orthogonalize	bool	Whether to enforce orthogonal angles in the final polygons	True
preserve_topology	bool	Whether to preserve topology during simplification	True

Returns:

Type	Description
Union[GeoDataFrame, List[Polygon]]	GeoDataFrame or list of shapely Polygons with regularized building footprints

Source code in geoai/utils/geometry.py

def regularization(
    building_polygons: Union[gpd.GeoDataFrame, List[Polygon]],
    angle_tolerance: float = 10,
    simplify_tolerance: float = 0.5,
    orthogonalize: bool = True,
    preserve_topology: bool = True,
) -> Union[gpd.GeoDataFrame, List[Polygon]]:
    """
    Regularizes building footprint polygons with multiple techniques beyond minimum
    rotated rectangles.

    Args:
        building_polygons: GeoDataFrame or list of shapely Polygons containing building footprints
        angle_tolerance: Degrees within which angles will be regularized to 90/180 degrees
        simplify_tolerance: Distance tolerance for Douglas-Peucker simplification
        orthogonalize: Whether to enforce orthogonal angles in the final polygons
        preserve_topology: Whether to preserve topology during simplification

    Returns:
        GeoDataFrame or list of shapely Polygons with regularized building footprints
    """
    from shapely import wkt
    from shapely.affinity import rotate, translate
    from shapely.geometry import Polygon, shape

    regularized_buildings = []

    # Check if we're dealing with a GeoDataFrame
    if isinstance(building_polygons, gpd.GeoDataFrame):
        geom_objects = building_polygons.geometry
    else:
        geom_objects = building_polygons

    for building in geom_objects:
        # Handle potential string representations of geometries
        if isinstance(building, str):
            try:
                # Try to parse as WKT
                building = wkt.loads(building)
            except Exception:
                print(f"Failed to parse geometry string: {building[:30]}...")
                continue

        # Ensure we have a valid geometry
        if not hasattr(building, "simplify"):
            print(f"Invalid geometry type: {type(building)}")
            continue

        # Step 1: Simplify to remove noise and small vertices
        simplified = building.simplify(
            simplify_tolerance, preserve_topology=preserve_topology
        )

        if orthogonalize:
            # Make sure we have a valid polygon with an exterior
            if not hasattr(simplified, "exterior") or simplified.exterior is None:
                print(f"Simplified geometry has no exterior: {simplified}")
                regularized_buildings.append(building)  # Use original instead
                continue

            # Step 2: Get the dominant angle to rotate building
            coords = np.array(simplified.exterior.coords)

            # Make sure we have enough coordinates for angle calculation
            if len(coords) < 3:
                print(f"Not enough coordinates for angle calculation: {len(coords)}")
                regularized_buildings.append(building)  # Use original instead
                continue

            segments = np.diff(coords, axis=0)
            angles = np.arctan2(segments[:, 1], segments[:, 0]) * 180 / np.pi

            # Find most common angle classes (0, 90, 180, 270 degrees)
            binned_angles = np.round(angles / 90) * 90
            dominant_angle = np.bincount(binned_angles.astype(int) % 180).argmax()

            # Step 3: Rotate to align with axes, regularize, then rotate back
            rotated = rotate(simplified, -dominant_angle, origin="centroid")

            # Step 4: Rectify coordinates to enforce right angles
            ext_coords = np.array(rotated.exterior.coords)
            rect_coords = []

            # Regularize each vertex to create orthogonal corners
            for i in range(len(ext_coords) - 1):
                rect_coords.append(ext_coords[i])

                # Check if we need to add a right-angle vertex
                angle = (
                    np.arctan2(
                        ext_coords[(i + 1) % (len(ext_coords) - 1), 1]
                        - ext_coords[i, 1],
                        ext_coords[(i + 1) % (len(ext_coords) - 1), 0]
                        - ext_coords[i, 0],
                    )
                    * 180
                    / np.pi
                )

                if abs(angle % 90) > angle_tolerance and abs(angle % 90) < (
                    90 - angle_tolerance
                ):
                    # Add intermediate point to create right angle
                    rect_coords.append(
                        [
                            ext_coords[(i + 1) % (len(ext_coords) - 1), 0],
                            ext_coords[i, 1],
                        ]
                    )

            # Close the polygon by adding the first point again
            rect_coords.append(rect_coords[0])

            # Create regularized polygon and rotate back
            regularized = Polygon(rect_coords)
            final_building = rotate(regularized, dominant_angle, origin="centroid")
        else:
            final_building = simplified

        regularized_buildings.append(final_building)

    # If input was a GeoDataFrame, return a GeoDataFrame
    if isinstance(building_polygons, gpd.GeoDataFrame):
        return gpd.GeoDataFrame(
            geometry=regularized_buildings, crs=building_polygons.crs
        )
    else:
        return regularized_buildings

regularize(data, parallel_threshold=1.0, target_crs=None, simplify=True, simplify_tolerance=0.5, allow_45_degree=True, diagonal_threshold_reduction=15, allow_circles=True, circle_threshold=0.9, num_cores=1, include_metadata=False, output_path=None, **kwargs)

Regularizes polygon geometries in a GeoDataFrame by aligning edges.

Aligns edges to be parallel or perpendicular (optionally also 45 degrees) to their main direction. Handles reprojection, initial simplification, regularization, geometry cleanup, and parallel processing.

This function is a wrapper around the regularize_geodataframe function from the buildingregulariser package. Credits to the original author Nick Wright. Check out the repo at https://github.com/DPIRD-DMA/Building-Regulariser.

Parameters:

Name	Type	Description	Default
data	Union[GeoDataFrame, str]	Input GeoDataFrame with polygon or multipolygon geometries, or a file path to the GeoDataFrame.	required
parallel_threshold	float	Distance threshold for merging nearly parallel adjacent edges during regularization. Defaults to 1.0.	1.0
target_crs	Optional[Union[str, CRS]]	Target Coordinate Reference System for processing. If None, uses the input GeoDataFrame's CRS. Processing is more reliable in a projected CRS. Defaults to None.	None
simplify	bool	If True, applies initial simplification to the geometry before regularization. Defaults to True.	True
simplify_tolerance	float	Tolerance for the initial simplification step (if simplify is True). Also used for geometry cleanup steps. Defaults to 0.5.	0.5
allow_45_degree	bool	If True, allows edges to be oriented at 45-degree angles relative to the main direction during regularization. Defaults to True.	True
diagonal_threshold_reduction	float	Reduction factor in degrees to reduce the likelihood of diagonal edges being created. Larger values reduce the likelihood of diagonal edges. Defaults to 15.	15
allow_circles	bool	If True, attempts to detect polygons that are nearly circular and replaces them with perfect circles. Defaults to True.	True
circle_threshold	float	Intersection over Union (IoU) threshold used for circle detection (if allow_circles is True). Value between 0 and 1. Defaults to 0.9.	0.9
num_cores	int	Number of CPU cores to use for parallel processing. If 1, processing is done sequentially. Defaults to 1.	1
include_metadata	bool	If True, includes metadata about the regularization process in the output GeoDataFrame. Defaults to False.	False
output_path	Optional[str]	Path to save the output GeoDataFrame. If None, the output is not saved. Defaults to None.	None
**kwargs	Any	Additional keyword arguments to pass to the to_file method when saving the output.	{}

Returns:

Type	Description
Any	gpd.GeoDataFrame: A new GeoDataFrame with regularized polygon geometries. Original attributes are
Any	preserved. Geometries that failed processing might be dropped.

Raises:

Type	Description
ValueError	If the input data is not a GeoDataFrame or a file path, or if the input GeoDataFrame is empty.

Source code in geoai/utils/geometry.py

def regularize(
    data: Union[gpd.GeoDataFrame, str],
    parallel_threshold: float = 1.0,
    target_crs: Optional[Union[str, "pyproj.CRS"]] = None,
    simplify: bool = True,
    simplify_tolerance: float = 0.5,
    allow_45_degree: bool = True,
    diagonal_threshold_reduction: float = 15,
    allow_circles: bool = True,
    circle_threshold: float = 0.9,
    num_cores: int = 1,
    include_metadata: bool = False,
    output_path: Optional[str] = None,
    **kwargs: Any,
) -> Any:
    """Regularizes polygon geometries in a GeoDataFrame by aligning edges.

    Aligns edges to be parallel or perpendicular (optionally also 45 degrees)
    to their main direction. Handles reprojection, initial simplification,
    regularization, geometry cleanup, and parallel processing.

    This function is a wrapper around the `regularize_geodataframe` function
    from the `buildingregulariser` package. Credits to the original author
    Nick Wright. Check out the repo at https://github.com/DPIRD-DMA/Building-Regulariser.

    Args:
        data (Union[gpd.GeoDataFrame, str]): Input GeoDataFrame with polygon or multipolygon geometries,
            or a file path to the GeoDataFrame.
        parallel_threshold (float, optional): Distance threshold for merging nearly parallel adjacent edges
            during regularization. Defaults to 1.0.
        target_crs (Optional[Union[str, "pyproj.CRS"]], optional): Target Coordinate Reference System for
            processing. If None, uses the input GeoDataFrame's CRS. Processing is more reliable in a
            projected CRS. Defaults to None.
        simplify (bool, optional): If True, applies initial simplification to the geometry before
            regularization. Defaults to True.
        simplify_tolerance (float, optional): Tolerance for the initial simplification step (if `simplify`
            is True). Also used for geometry cleanup steps. Defaults to 0.5.
        allow_45_degree (bool, optional): If True, allows edges to be oriented at 45-degree angles relative
            to the main direction during regularization. Defaults to True.
        diagonal_threshold_reduction (float, optional): Reduction factor in degrees to reduce the likelihood
            of diagonal edges being created. Larger values reduce the likelihood of diagonal edges.
            Defaults to 15.
        allow_circles (bool, optional): If True, attempts to detect polygons that are nearly circular and
            replaces them with perfect circles. Defaults to True.
        circle_threshold (float, optional): Intersection over Union (IoU) threshold used for circle detection
            (if `allow_circles` is True). Value between 0 and 1. Defaults to 0.9.
        num_cores (int, optional): Number of CPU cores to use for parallel processing. If 1, processing is
            done sequentially. Defaults to 1.
        include_metadata (bool, optional): If True, includes metadata about the regularization process in the
            output GeoDataFrame. Defaults to False.
        output_path (Optional[str], optional): Path to save the output GeoDataFrame. If None, the output is
            not saved. Defaults to None.
        **kwargs: Additional keyword arguments to pass to the `to_file` method when saving the output.

    Returns:
        gpd.GeoDataFrame: A new GeoDataFrame with regularized polygon geometries. Original attributes are
        preserved. Geometries that failed processing might be dropped.

    Raises:
        ValueError: If the input data is not a GeoDataFrame or a file path, or if the input GeoDataFrame is empty.
    """
    try:
        from buildingregulariser import regularize_geodataframe
    except ImportError:
        install_package("buildingregulariser")
        from buildingregulariser import regularize_geodataframe

    if isinstance(data, str):
        data = gpd.read_file(data)
    elif not isinstance(data, gpd.GeoDataFrame):
        raise ValueError("Input data must be a GeoDataFrame or a file path.")

    # Check if the input data is empty
    if data.empty:
        raise ValueError("Input GeoDataFrame is empty.")

    gdf = regularize_geodataframe(
        data,
        parallel_threshold=parallel_threshold,
        target_crs=target_crs,
        simplify=simplify,
        simplify_tolerance=simplify_tolerance,
        allow_45_degree=allow_45_degree,
        diagonal_threshold_reduction=diagonal_threshold_reduction,
        allow_circles=allow_circles,
        circle_threshold=circle_threshold,
        num_cores=num_cores,
        include_metadata=include_metadata,
    )

    if output_path:
        gdf.to_file(output_path, **kwargs)

    return gdf

rowcol_to_xy(src_fp, rows=None, cols=None, boxes=None, zs=None, offset='center', output=None, dst_crs='EPSG:4326', **kwargs)

Converts a list of (row, col) coordinates to (x, y) coordinates.

Parameters:

Name	Type	Description	Default
src_fp	str	The source raster file path.	required
rows	list	A list of row coordinates. Defaults to None.	None
cols	list	A list of col coordinates. Defaults to None.	None
boxes	list	A list of (row, col) coordinates in the format of [[left, top, right, bottom], [left, top, right, bottom], ...]	None
zs	Optional[List[float]]	zs (list or float, optional): Height associated with coordinates. Primarily used for RPC based coordinate transformations.	None
offset	str	Determines if the returned coordinates are for the center of the pixel or for a corner.	'center'
output	str	The output vector file path. Defaults to None.	None
dst_crs	str	The destination CRS. Defaults to "EPSG:4326".	'EPSG:4326'
**kwargs	Any	Additional keyword arguments to pass to rasterio.transform.xy.	{}

Returns:

Type	Description
Tuple[List[float], List[float]]	A list of (x, y) coordinates.

Source code in geoai/utils/conversion.py

def rowcol_to_xy(
    src_fp: str,
    rows: Optional[List[int]] = None,
    cols: Optional[List[int]] = None,
    boxes: Optional[List[List[int]]] = None,
    zs: Optional[List[float]] = None,
    offset: str = "center",
    output: Optional[str] = None,
    dst_crs: str = "EPSG:4326",
    **kwargs: Any,
) -> Tuple[List[float], List[float]]:
    """Converts a list of (row, col) coordinates to (x, y) coordinates.

    Args:
        src_fp (str): The source raster file path.
        rows (list, optional): A list of row coordinates. Defaults to None.
        cols (list, optional): A list of col coordinates. Defaults to None.
        boxes (list, optional): A list of (row, col) coordinates in the format of [[left, top, right, bottom], [left, top, right, bottom], ...]
        zs: zs (list or float, optional): Height associated with coordinates. Primarily used for RPC based coordinate transformations.
        offset (str, optional): Determines if the returned coordinates are for the center of the pixel or for a corner.
        output (str, optional): The output vector file path. Defaults to None.
        dst_crs (str, optional): The destination CRS. Defaults to "EPSG:4326".
        **kwargs: Additional keyword arguments to pass to rasterio.transform.xy.

    Returns:
        A list of (x, y) coordinates.
    """

    if boxes is not None:
        rows = []
        cols = []

        for box in boxes:
            rows.append(box[1])
            rows.append(box[3])
            cols.append(box[0])
            cols.append(box[2])

    if rows is None or cols is None:
        raise ValueError("rows and cols must be provided.")

    with rasterio.open(src_fp) as src:
        xs, ys = rasterio.transform.xy(src.transform, rows, cols, zs, offset, **kwargs)
        src_crs = src.crs

    if boxes is None:
        return [[x, y] for x, y in zip(xs, ys)]
    else:
        result = [[xs[i], ys[i + 1], xs[i + 1], ys[i]] for i in range(0, len(xs), 2)]

        if output is not None:
            from .vector import boxes_to_vector

            boxes_to_vector(result, src_crs, dst_crs, output)
        else:
            return result

smooth_vector(vector_data, output_path=None, segment_length=None, smooth_iterations=3, num_cores=0, merge_collection=True, merge_field=None, merge_multipolygons=True, preserve_area=True, area_tolerance=0.01, **kwargs)

Smooth a vector data using the smoothify library. See https://github.com/DPIRD-DMA/Smoothify for more details.

Parameters:

Name	Type	Description	Default
vector_data	Union[str, GeoDataFrame]	The vector data to smooth.	required
output_path	str	The path to save the smoothed vector data. If None, returns the smoothed vector data.	None
segment_length	float	Resolution of the original raster data in map units. If None (default), automatically detects by finding the minimum segment length (from a data sample). Recommended to specify explicitly when known.	None
smooth_iterations	int	The number of iterations to smooth the vector data.	3
num_cores	int	Number of cores to use for parallel processing. If 0 (default), uses all available cores.	0
merge_collection	bool	Whether to merge/dissolve adjacent geometries in collections before smoothing.	True
merge_field	str	Column name to use for dissolving geometries. Only valid when merge_collection=True. If None, dissolves all geometries together. If specified, dissolves geometries grouped by the column values.	None
merge_multipolygons	bool	Whether to merge adjacent polygons within MultiPolygons before smoothing	True
preserve_area	bool	Whether to restore original area after smoothing via buffering (applies to Polygons only)	True
area_tolerance	float	Percentage of original area allowed as error (e.g., 0.01 = 0.01% error = 99.99% preservation). Only affects Polygons when preserve_area=True	0.01

Returns:

Type	Description
GeoDataFrame	gpd.GeoDataFrame: The smoothed vector data.

Examples:

>>> import geoai
>>> gdf = geoai.read_vector("path/to/vector.geojson")
>>> smoothed_gdf = geoai.smooth_vector(gdf, smooth_iterations=3, output_path="path/to/smoothed_vector.geojson")
>>> smoothed_gdf.head()
>>> smoothed_gdf.explore()

Source code in geoai/utils/vector.py

def smooth_vector(
    vector_data: Union[str, gpd.GeoDataFrame],
    output_path: str = None,
    segment_length: float = None,
    smooth_iterations: int = 3,
    num_cores: int = 0,
    merge_collection: bool = True,
    merge_field: str = None,
    merge_multipolygons: bool = True,
    preserve_area: bool = True,
    area_tolerance: float = 0.01,
    **kwargs: Any,
) -> gpd.GeoDataFrame:
    """Smooth a vector data using the smoothify library.
        See https://github.com/DPIRD-DMA/Smoothify for more details.

    Args:
        vector_data: The vector data to smooth.
        output_path: The path to save the smoothed vector data. If None, returns the smoothed vector data.
        segment_length: Resolution of the original raster data in map units. If None (default), automatically
            detects by finding the minimum segment length (from a data sample). Recommended to specify explicitly when known.
        smooth_iterations: The number of iterations to smooth the vector data.
        num_cores: Number of cores to use for parallel processing. If 0 (default), uses all available cores.
        merge_collection: Whether to merge/dissolve adjacent geometries in collections before smoothing.
        merge_field: Column name to use for dissolving geometries. Only valid when merge_collection=True.
            If None, dissolves all geometries together. If specified, dissolves geometries grouped by the column values.
        merge_multipolygons: Whether to merge adjacent polygons within MultiPolygons before smoothing
        preserve_area: Whether to restore original area after smoothing via buffering (applies to Polygons only)
        area_tolerance: Percentage of original area allowed as error (e.g., 0.01 = 0.01% error = 99.99% preservation).
            Only affects Polygons when preserve_area=True

    Returns:
        gpd.GeoDataFrame: The smoothed vector data.

    Examples:
        >>> import geoai
        >>> gdf = geoai.read_vector("path/to/vector.geojson")
        >>> smoothed_gdf = geoai.smooth_vector(gdf, smooth_iterations=3, output_path="path/to/smoothed_vector.geojson")
        >>> smoothed_gdf.head()
        >>> smoothed_gdf.explore()
    """
    try:
        from smoothify import smoothify
    except ImportError:
        install_package("smoothify")
        from smoothify import smoothify

    if isinstance(vector_data, str):
        vector_data = leafmap.read_vector(vector_data)

    smoothed_vector_data = smoothify(
        geom=vector_data,
        segment_length=segment_length,
        smooth_iterations=smooth_iterations,
        num_cores=num_cores,
        merge_collection=merge_collection,
        merge_field=merge_field,
        merge_multipolygons=merge_multipolygons,
        preserve_area=preserve_area,
        area_tolerance=area_tolerance,
        **kwargs,
    )
    if output_path is not None:
        smoothed_vector_data.to_file(output_path)
    return smoothed_vector_data

stack_bands(input_files, output_file, resolution=None, dtype=None, temp_vrt='stack.vrt', overwrite=False, compress='DEFLATE', output_format='COG', extra_gdal_translate_args=None)

Stack bands from multiple images into a single multi-band GeoTIFF.

Parameters:

Name	Type	Description	Default
input_files	List[str]	List of input image paths.	required
output_file	str	Path to the output stacked image.	required
resolution	float	Output resolution. If None, inferred from first image.	None
dtype	str	Output data type (e.g., "UInt16", "Float32").	None
temp_vrt	str	Temporary VRT filename.	'stack.vrt'
overwrite	bool	Whether to overwrite the output file.	False
compress	str	Compression method.	'DEFLATE'
output_format	str	GDAL output format (default is "COG").	'COG'
extra_gdal_translate_args	List[str]	Extra arguments for gdal_translate.	None

Returns:

Name	Type	Description
str	str	Path to the output file.

Source code in geoai/utils/raster.py

def stack_bands(
    input_files: List[str],
    output_file: str,
    resolution: Optional[float] = None,
    dtype: Optional[str] = None,  # e.g., "UInt16", "Float32"
    temp_vrt: str = "stack.vrt",
    overwrite: bool = False,
    compress: str = "DEFLATE",
    output_format: str = "COG",
    extra_gdal_translate_args: Optional[List[str]] = None,
) -> str:
    """
    Stack bands from multiple images into a single multi-band GeoTIFF.

    Args:
        input_files (List[str]): List of input image paths.
        output_file (str): Path to the output stacked image.
        resolution (float, optional): Output resolution. If None, inferred from first image.
        dtype (str, optional): Output data type (e.g., "UInt16", "Float32").
        temp_vrt (str): Temporary VRT filename.
        overwrite (bool): Whether to overwrite the output file.
        compress (str): Compression method.
        output_format (str): GDAL output format (default is "COG").
        extra_gdal_translate_args (List[str], optional): Extra arguments for gdal_translate.

    Returns:
        str: Path to the output file.
    """
    import leafmap

    if not input_files:
        raise ValueError("No input files provided.")
    elif isinstance(input_files, str):
        input_files = leafmap.find_files(input_files, ".tif")

    if os.path.exists(output_file) and not overwrite:
        print(f"Output file already exists: {output_file}")
        return output_file

    # Infer resolution if not provided
    if resolution is None:
        resolution_x, resolution_y = get_raster_resolution(input_files[0])
    else:
        resolution_x = resolution_y = resolution

    # Step 1: Build VRT
    vrt_cmd = ["gdalbuildvrt", "-separate", temp_vrt] + input_files
    subprocess.run(vrt_cmd, check=True)

    # Step 2: Translate VRT to output GeoTIFF
    translate_cmd = [
        "gdal_translate",
        "-tr",
        str(resolution_x),
        str(resolution_y),
        temp_vrt,
        output_file,
        "-of",
        output_format,
        "-co",
        f"COMPRESS={compress}",
    ]

    if dtype:
        translate_cmd.insert(1, "-ot")
        translate_cmd.insert(2, dtype)

    if extra_gdal_translate_args:
        translate_cmd += extra_gdal_translate_args

    subprocess.run(translate_cmd, check=True)

    # Step 3: Clean up VRT
    if os.path.exists(temp_vrt):
        os.remove(temp_vrt)

    return output_file

temp_file_path(ext)

Returns a temporary file path.

Parameters:

Name	Type	Description	Default
ext	str	The file extension.	required

Returns:

Name	Type	Description
str	str	The temporary file path.

Source code in geoai/utils/device.py

def temp_file_path(ext: str) -> str:
    """Returns a temporary file path.

    Args:
        ext (str): The file extension.

    Returns:
        str: The temporary file path.
    """

    import tempfile
    import uuid

    if not ext.startswith("."):
        ext = "." + ext
    file_id = str(uuid.uuid4())
    file_path = os.path.join(tempfile.gettempdir(), f"{file_id}{ext}")

    return file_path

train_segmentation_landcover(images_dir, labels_dir, output_dir, input_format='directory', architecture='unet', encoder_name='resnet34', encoder_weights='imagenet', num_channels=3, num_classes=2, batch_size=8, num_epochs=50, learning_rate=0.001, weight_decay=0.0001, seed=42, val_split=0.2, print_freq=10, verbose=True, save_best_only=True, plot_curves=False, device=None, checkpoint_path=None, resume_training=False, target_size=None, resize_mode='resize', num_workers=None, loss_function='crossentropy', ignore_index=0, use_class_weights=False, focal_alpha=1.0, focal_gamma=2.0, custom_multipliers=None, max_class_weight=50.0, use_inverse_frequency=True, validation_iou_mode='standard', boundary_alpha=1.0, background_class=0, training_callback=None, **kwargs)

Train a semantic segmentation model with landcover-specific enhancements.

This is a standalone version that wraps geoai.train.train_segmentation_model with landcover-specific loss functions, class weights, and metrics.

Parameters:

Name	Type	Description	Default
images_dir	str	Directory containing training images	required
labels_dir	str	Directory containing training labels	required
output_dir	str	Directory to save model checkpoints and training history	required
input_format	str	Data format ("directory", "COCO", "YOLO")	'directory'
architecture	str	Model architecture (default: "unet")	'unet'
encoder_name	str	Encoder backbone (default: "resnet34")	'resnet34'
encoder_weights	Optional[str]	Pretrained weights ("imagenet" or None)	'imagenet'
num_channels	int	Number of input channels (default: 3)	3
num_classes	int	Number of output classes (default: 2)	2
batch_size	int	Training batch size (default: 8)	8
num_epochs	int	Number of training epochs (default: 50)	50
learning_rate	float	Initial learning rate (default: 0.001)	0.001
weight_decay	float	Weight decay for optimizer (default: 1e-4)	0.0001
seed	int	Random seed for reproducibility (default: 42)	42
val_split	float	Validation split ratio (default: 0.2)	0.2
print_freq	int	Frequency of training progress prints (default: 10)	10
verbose	bool	Enable verbose output (default: True)	True
save_best_only	bool	Only save best model checkpoint (default: True)	True
plot_curves	bool	Plot training curves at end (default: False)	False
device	Optional[device]	Torch device (auto-detected if None)	None
checkpoint_path	Optional[str]	Path to checkpoint for resuming training	None
resume_training	bool	Whether to resume from checkpoint (default: False)	False
target_size	Optional[Tuple[int, int]]	Target size for resizing images (H, W) or None	None
resize_mode	str	How to resize ("resize", "crop", or "pad")	'resize'
num_workers	Optional[int]	Number of dataloader workers (default: auto)	None
loss_function	str	Loss function name ("crossentropy", "focal")	'crossentropy'
ignore_index	Union[int, bool]	Class index to ignore during training. (default: 0) - If int: pixels with this label value will be ignored during training - If False: no pixels will be ignored (all pixels contribute to loss)	0
use_class_weights	bool	Whether to compute and use class weights (default: False)	False
focal_alpha	float	Focal loss alpha parameter (default: 1.0)	1.0
focal_gamma	float	Focal loss gamma parameter (default: 2.0)	2.0
custom_multipliers	Optional[Dict[int, float]]	Custom class weight multipliers {class_id: multiplier}	None
max_class_weight	float	Maximum allowed class weight (default: 50.0)	50.0
use_inverse_frequency	bool	Use inverse frequency for weights (default: True)	True
validation_iou_mode	str	IoU calculation mode for validation (default: "standard") - "standard": Unweighted mean IoU (all classes equal importance) - "perclass_frequency": Frequency-weighted IoU (classes weighted by pixel count) - "boundary_weighted": Boundary-distance weighted IoU (wIoU, focus on edges) - "sparse_labels": For incomplete ground truth - predictions in background areas are NOT penalized. Uses custom training loop with sparse IoU for model selection. BEST FOR INCOMPLETE/SPARSE HABITAT MASKS.	'standard'
boundary_alpha	float	Boundary importance factor for wIoU mode (default: 1.0) Higher values = more focus on boundaries (0.01-100 range)	1.0
background_class	int	Class ID for background/unlabeled pixels in sparse_labels mode (default: 0). Predictions in this class area are NOT counted as false positives.	0
training_callback	Optional[callable]	Optional callback function for automatic metric tracking	None
**kwargs	Any	Additional arguments passed to base training function	{}

Returns:

Type	Description
Module	Trained model

Example

from landcover_train import train_segmentation_landcover

model = train_segmentation_landcover( ... images_dir="tiles/images", ... labels_dir="tiles/labels", ... output_dir="models/landcover_001", ... num_classes=5, ... loss_function="focal", ... ignore_index=0, # Ignore background ... use_class_weights=True, ... custom_multipliers={1: 1.5, 4: 0.8}, # Boost class 1, reduce class 4 ... max_class_weight=50.0, ... use_inverse_frequency=True, # Use inverse frequency weighting ... validation_iou_mode="boundary_weighted", # Focus on boundaries ... boundary_alpha=2.0, # Moderate boundary emphasis ... )

Source code in geoai/landcover_train.py

def train_segmentation_landcover(
    images_dir: str,
    labels_dir: str,
    output_dir: str,
    input_format: str = "directory",
    architecture: str = "unet",
    encoder_name: str = "resnet34",
    encoder_weights: Optional[str] = "imagenet",
    num_channels: int = 3,
    num_classes: int = 2,
    batch_size: int = 8,
    num_epochs: int = 50,
    learning_rate: float = 0.001,
    weight_decay: float = 1e-4,
    seed: int = 42,
    val_split: float = 0.2,
    print_freq: int = 10,
    verbose: bool = True,
    save_best_only: bool = True,
    plot_curves: bool = False,
    device: Optional[torch.device] = None,
    checkpoint_path: Optional[str] = None,
    resume_training: bool = False,
    target_size: Optional[Tuple[int, int]] = None,
    resize_mode: str = "resize",
    num_workers: Optional[int] = None,
    loss_function: str = "crossentropy",
    ignore_index: Union[int, bool] = 0,
    use_class_weights: bool = False,
    focal_alpha: float = 1.0,
    focal_gamma: float = 2.0,
    custom_multipliers: Optional[Dict[int, float]] = None,
    max_class_weight: float = 50.0,
    use_inverse_frequency: bool = True,
    validation_iou_mode: str = "standard",
    boundary_alpha: float = 1.0,
    background_class: int = 0,
    training_callback: Optional[callable] = None,
    **kwargs: Any,
) -> torch.nn.Module:
    """
    Train a semantic segmentation model with landcover-specific enhancements.

    This is a standalone version that wraps geoai.train.train_segmentation_model
    with landcover-specific loss functions, class weights, and metrics.

    Args:
        images_dir: Directory containing training images
        labels_dir: Directory containing training labels
        output_dir: Directory to save model checkpoints and training history
        input_format: Data format ("directory", "COCO", "YOLO")
        architecture: Model architecture (default: "unet")
        encoder_name: Encoder backbone (default: "resnet34")
        encoder_weights: Pretrained weights ("imagenet" or None)
        num_channels: Number of input channels (default: 3)
        num_classes: Number of output classes (default: 2)
        batch_size: Training batch size (default: 8)
        num_epochs: Number of training epochs (default: 50)
        learning_rate: Initial learning rate (default: 0.001)
        weight_decay: Weight decay for optimizer (default: 1e-4)
        seed: Random seed for reproducibility (default: 42)
        val_split: Validation split ratio (default: 0.2)
        print_freq: Frequency of training progress prints (default: 10)
        verbose: Enable verbose output (default: True)
        save_best_only: Only save best model checkpoint (default: True)
        plot_curves: Plot training curves at end (default: False)
        device: Torch device (auto-detected if None)
        checkpoint_path: Path to checkpoint for resuming training
        resume_training: Whether to resume from checkpoint (default: False)
        target_size: Target size for resizing images (H, W) or None
        resize_mode: How to resize ("resize", "crop", or "pad")
        num_workers: Number of dataloader workers (default: auto)
        loss_function: Loss function name ("crossentropy", "focal")
        ignore_index: Class index to ignore during training. (default: 0)
            - If int: pixels with this label value will be ignored during training
            - If False: no pixels will be ignored (all pixels contribute to loss)
        use_class_weights: Whether to compute and use class weights (default: False)
        focal_alpha: Focal loss alpha parameter (default: 1.0)
        focal_gamma: Focal loss gamma parameter (default: 2.0)
        custom_multipliers: Custom class weight multipliers {class_id: multiplier}
        max_class_weight: Maximum allowed class weight (default: 50.0)
        use_inverse_frequency: Use inverse frequency for weights (default: True)
        validation_iou_mode: IoU calculation mode for validation (default: "standard")
            - "standard": Unweighted mean IoU (all classes equal importance)
            - "perclass_frequency": Frequency-weighted IoU (classes weighted by pixel count)
            - "boundary_weighted": Boundary-distance weighted IoU (wIoU, focus on edges)
            - "sparse_labels": For incomplete ground truth - predictions in background
              areas are NOT penalized. Uses custom training loop with sparse IoU for
              model selection. BEST FOR INCOMPLETE/SPARSE HABITAT MASKS.
        boundary_alpha: Boundary importance factor for wIoU mode (default: 1.0)
            Higher values = more focus on boundaries (0.01-100 range)
        background_class: Class ID for background/unlabeled pixels in sparse_labels mode
            (default: 0). Predictions in this class area are NOT counted as false positives.
        training_callback: Optional callback function for automatic metric tracking
        **kwargs: Additional arguments passed to base training function

    Returns:
        Trained model

    Example:
        >>> from landcover_train import train_segmentation_landcover
        >>>
        >>> model = train_segmentation_landcover(
        ...     images_dir="tiles/images",
        ...     labels_dir="tiles/labels",
        ...     output_dir="models/landcover_001",
        ...     num_classes=5,
        ...     loss_function="focal",
        ...     ignore_index=0,  # Ignore background
        ...     use_class_weights=True,
        ...     custom_multipliers={1: 1.5, 4: 0.8},  # Boost class 1, reduce class 4
        ...     max_class_weight=50.0,
        ...     use_inverse_frequency=True,  # Use inverse frequency weighting
        ...     validation_iou_mode="boundary_weighted",  # Focus on boundaries
        ...     boundary_alpha=2.0,  # Moderate boundary emphasis
        ... )
    """

    # Convert ignore_index to format expected by PyTorch loss functions
    if isinstance(ignore_index, bool) and ignore_index is False:
        ignore_idx_for_loss = -100  # PyTorch default (effectively no ignoring)
    elif isinstance(ignore_index, int):
        ignore_idx_for_loss = ignore_index
    else:
        ignore_idx_for_loss = -100

    # Compute class weights if requested
    class_weights_tensor = None
    if use_class_weights:
        if verbose:
            print("\n" + "=" * 60)
            print("COMPUTING CLASS WEIGHTS")
            print("=" * 60)

        class_weights_tensor = compute_class_weights(
            labels_dir=labels_dir,
            num_classes=num_classes,
            ignore_index=ignore_index,
            custom_multipliers=custom_multipliers,
            max_weight=max_class_weight,
            use_inverse_frequency=use_inverse_frequency,
        )

        if verbose:
            print("=" * 60 + "\n")

    # Create custom loss function using landcover-specific implementation
    # This ensures ignore_index and class_weights are properly used
    import torch

    device = (
        device
        if device is not None
        else torch.device("cuda" if torch.cuda.is_available() else "cpu")
    )

    if class_weights_tensor is not None:
        class_weights_tensor = class_weights_tensor.to(device)

    # Create the loss function with proper ignore_index support
    criterion = get_landcover_loss_function(
        loss_name=loss_function,
        num_classes=num_classes,
        ignore_index=ignore_idx_for_loss,
        class_weights=class_weights_tensor,
        use_class_weights=use_class_weights,
        focal_alpha=focal_alpha,
        focal_gamma=focal_gamma,
        device=device,
    )

    if verbose:
        print(
            f"Created {loss_function} loss function with ignore_index={ignore_idx_for_loss}"
        )
        if use_class_weights:
            print(f"Class weights applied: {class_weights_tensor}")

    # ==========================================================================
    # ALL MODES: Use custom training loop with landcover_iou for model selection
    # ==========================================================================
    # This ensures ALL IoU modes work correctly, not just sparse_labels
    # The base geoai training function ignores validation_iou_mode parameter

    if verbose:
        mode_descriptions = {
            "standard": "STANDARD (unweighted mean IoU)",
            "perclass_frequency": "PER-CLASS FREQUENCY-WEIGHTED IoU",
            "boundary_weighted": f"BOUNDARY-WEIGHTED IoU (wIoU, α={boundary_alpha})",
            "sparse_labels": f"SPARSE LABELS IoU (bg={background_class} ignored)",
        }
        print("\n" + "=" * 60)
        print(
            f"CUSTOM TRAINING LOOP: {mode_descriptions.get(validation_iou_mode, validation_iou_mode)}"
        )
        print("=" * 60)
        if validation_iou_mode == "sparse_labels":
            print(
                f"Background class: {background_class} (predictions here NOT penalized)"
            )
        elif validation_iou_mode == "boundary_weighted":
            print(
                f"Boundary alpha: {boundary_alpha} (higher = more focus on boundaries)"
            )
        elif validation_iou_mode == "perclass_frequency":
            print("Classes weighted by pixel frequency in dataset")
        print(f"Using {validation_iou_mode} IoU for model selection during training")
        print("=" * 60 + "\n")

    model = _train_with_custom_iou(
        images_dir=images_dir,
        labels_dir=labels_dir,
        output_dir=output_dir,
        architecture=architecture,
        encoder_name=encoder_name,
        encoder_weights=encoder_weights,
        num_channels=num_channels,
        num_classes=num_classes,
        batch_size=batch_size,
        num_epochs=num_epochs,
        learning_rate=learning_rate,
        weight_decay=weight_decay,
        seed=seed,
        val_split=val_split,
        print_freq=print_freq,
        verbose=verbose,
        save_best_only=save_best_only,
        plot_curves=plot_curves,
        device=device,
        checkpoint_path=checkpoint_path,
        resume_training=resume_training,
        target_size=target_size,
        num_workers=num_workers,
        criterion=criterion,
        validation_iou_mode=validation_iou_mode,
        boundary_alpha=boundary_alpha,
        background_class=background_class,
        ignore_index=(
            ignore_idx_for_loss
            if isinstance(ignore_idx_for_loss, int) and ignore_idx_for_loss != -100
            else False
        ),
        training_callback=training_callback,
        **kwargs,
    )
    return model

vector_to_geojson(filename, output=None, **kwargs)

Converts a vector file to a geojson file.

Parameters:

Name	Type	Description	Default
filename	str	The vector file path.	required
output	str	The output geojson file path. Defaults to None.	None

Returns:

Name	Type	Description
dict	str	The geojson dictionary.

Source code in geoai/utils/vector.py

def vector_to_geojson(
    filename: str, output: Optional[str] = None, **kwargs: Any
) -> str:
    """Converts a vector file to a geojson file.

    Args:
        filename (str): The vector file path.
        output (str, optional): The output geojson file path. Defaults to None.

    Returns:
        dict: The geojson dictionary.
    """

    if filename.startswith("http"):
        filename = download_file(filename)

    gdf = gpd.read_file(filename, **kwargs)
    if output is None:
        return gdf.__geo_interface__
    else:
        gdf.to_file(output, driver="GeoJSON")

vector_to_raster(vector_path, output_path=None, reference_raster=None, attribute_field=None, output_shape=None, transform=None, pixel_size=None, bounds=None, crs=None, all_touched=False, fill_value=0, dtype=np.uint8, nodata=None, plot_result=False)

Convert vector data to a raster.

Parameters:

Name	Type	Description	Default
vector_path	str or GeoDataFrame	Path to the input vector file or a GeoDataFrame.	required
output_path	str	Path to save the output raster file. If None, returns the array without saving.	None
reference_raster	str	Path to a reference raster for dimensions, transform and CRS.	None
attribute_field	str	Field name in the vector data to use for pixel values. If None, all vector features will be burned with value 1.	None
output_shape	tuple	Shape of the output raster as (height, width). Required if reference_raster is not provided.	None
transform	Affine	Affine transformation matrix. Required if reference_raster is not provided.	None
pixel_size	float or tuple	Pixel size (resolution) as single value or (x_res, y_res). Used to calculate transform if transform is not provided.	None
bounds	tuple	Bounds of the output raster as (left, bottom, right, top). Used to calculate transform if transform is not provided.	None
crs	str or CRS	Coordinate reference system of the output raster. Required if reference_raster is not provided.	None
all_touched	bool	If True, all pixels touched by geometries will be burned in. If False, only pixels whose center is within the geometry will be burned in.	False
fill_value	int	Value to fill the raster with before burning in features.	0
dtype	dtype	Data type of the output raster.	uint8
nodata	int	No data value for the output raster.	None
plot_result	bool	Whether to plot the resulting raster.	False

Returns:

Type	Description
ndarray	numpy.ndarray: The rasterized data array if output_path is None, else None.

Source code in geoai/utils/raster.py

def vector_to_raster(
    vector_path: Union[str, gpd.GeoDataFrame],
    output_path: Optional[str] = None,
    reference_raster: Optional[str] = None,
    attribute_field: Optional[str] = None,
    output_shape: Optional[Tuple[int, int]] = None,
    transform: Optional[Any] = None,
    pixel_size: Optional[float] = None,
    bounds: Optional[List[float]] = None,
    crs: Optional[str] = None,
    all_touched: bool = False,
    fill_value: Union[int, float] = 0,
    dtype: Any = np.uint8,
    nodata: Optional[Union[int, float]] = None,
    plot_result: bool = False,
) -> np.ndarray:
    """
    Convert vector data to a raster.

    Args:
        vector_path (str or GeoDataFrame): Path to the input vector file or a GeoDataFrame.
        output_path (str): Path to save the output raster file. If None, returns the array without saving.
        reference_raster (str): Path to a reference raster for dimensions, transform and CRS.
        attribute_field (str): Field name in the vector data to use for pixel values.
            If None, all vector features will be burned with value 1.
        output_shape (tuple): Shape of the output raster as (height, width).
            Required if reference_raster is not provided.
        transform (affine.Affine): Affine transformation matrix.
            Required if reference_raster is not provided.
        pixel_size (float or tuple): Pixel size (resolution) as single value or (x_res, y_res).
            Used to calculate transform if transform is not provided.
        bounds (tuple): Bounds of the output raster as (left, bottom, right, top).
            Used to calculate transform if transform is not provided.
        crs (str or CRS): Coordinate reference system of the output raster.
            Required if reference_raster is not provided.
        all_touched (bool): If True, all pixels touched by geometries will be burned in.
            If False, only pixels whose center is within the geometry will be burned in.
        fill_value (int): Value to fill the raster with before burning in features.
        dtype (numpy.dtype): Data type of the output raster.
        nodata (int): No data value for the output raster.
        plot_result (bool): Whether to plot the resulting raster.

    Returns:
        numpy.ndarray: The rasterized data array if output_path is None, else None.
    """
    # Load vector data
    if isinstance(vector_path, gpd.GeoDataFrame):
        gdf = vector_path
    else:
        gdf = gpd.read_file(vector_path)

    # Check if vector data is empty
    if gdf.empty:
        warnings.warn("The input vector data is empty. Creating an empty raster.")

    # Get CRS from vector data if not provided
    if crs is None and reference_raster is None:
        crs = gdf.crs

    # Get transform and output shape from reference raster if provided
    if reference_raster is not None:
        with rasterio.open(reference_raster) as src:
            transform = src.transform
            output_shape = src.shape
            crs = src.crs
            if nodata is None:
                nodata = src.nodata
    else:
        # Check if we have all required parameters
        if transform is None:
            if pixel_size is None or bounds is None:
                raise ValueError(
                    "Either reference_raster, transform, or both pixel_size and bounds must be provided."
                )

            # Calculate transform from pixel size and bounds
            if isinstance(pixel_size, (int, float)):
                x_res = y_res = float(pixel_size)
            else:
                x_res, y_res = pixel_size
                y_res = abs(y_res) * -1  # Convert to negative for north-up raster

            left, bottom, right, top = bounds
            transform = rasterio.transform.from_bounds(
                left,
                bottom,
                right,
                top,
                int((right - left) / x_res),
                int((top - bottom) / abs(y_res)),
            )

        if output_shape is None:
            # Calculate output shape from bounds and pixel size
            if bounds is None or pixel_size is None:
                raise ValueError(
                    "output_shape must be provided if reference_raster is not provided and "
                    "cannot be calculated from bounds and pixel_size."
                )

            if isinstance(pixel_size, (int, float)):
                x_res = y_res = float(pixel_size)
            else:
                x_res, y_res = pixel_size

            left, bottom, right, top = bounds
            width = int((right - left) / x_res)
            height = int((top - bottom) / abs(y_res))
            output_shape = (height, width)

    # Ensure CRS is set
    if crs is None:
        raise ValueError(
            "CRS must be provided either directly, from reference_raster, or from input vector data."
        )

    # Reproject vector data if its CRS doesn't match the output CRS
    if gdf.crs != crs:
        print(f"Reprojecting vector data from {gdf.crs} to {crs}")
        gdf = gdf.to_crs(crs)

    # Create empty raster filled with fill_value
    raster_data = np.full(output_shape, fill_value, dtype=dtype)

    # Burn vector features into raster
    if not gdf.empty:
        # Prepare shapes for burning
        if attribute_field is not None and attribute_field in gdf.columns:
            # Use attribute field for values
            shapes = [
                (geom, value) for geom, value in zip(gdf.geometry, gdf[attribute_field])
            ]
        else:
            # Burn with value 1
            shapes = [(geom, 1) for geom in gdf.geometry]

        # Burn shapes into raster
        burned = features.rasterize(
            shapes=shapes,
            out_shape=output_shape,
            transform=transform,
            fill=fill_value,
            all_touched=all_touched,
            dtype=dtype,
        )

        # Update raster data
        raster_data = burned

    # Save raster if output path is provided
    if output_path is not None:
        # Create directory if it doesn't exist
        os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)

        # Define metadata
        metadata = {
            "driver": "GTiff",
            "height": output_shape[0],
            "width": output_shape[1],
            "count": 1,
            "dtype": raster_data.dtype,
            "crs": crs,
            "transform": transform,
        }

        # Add nodata value if provided
        if nodata is not None:
            metadata["nodata"] = nodata

        # Write raster
        with rasterio.open(output_path, "w", **metadata) as dst:
            dst.write(raster_data, 1)

        print(f"Rasterized data saved to {output_path}")

    # Plot result if requested
    if plot_result:
        fig, ax = plt.subplots(figsize=(10, 10))

        # Plot raster
        im = ax.imshow(raster_data, cmap="viridis")
        plt.colorbar(im, ax=ax, label=attribute_field if attribute_field else "Value")

        # Plot vector boundaries for reference
        if output_path is not None:
            # Get the extent of the raster
            with rasterio.open(output_path) as src:
                bounds = src.bounds
                raster_bbox = box(*bounds)
        else:
            # Calculate extent from transform and shape
            height, width = output_shape
            left, top = transform * (0, 0)
            right, bottom = transform * (width, height)
            raster_bbox = box(left, bottom, right, top)

        # Clip vector to raster extent for clarity in plot
        if not gdf.empty:
            gdf_clipped = gpd.clip(gdf, raster_bbox)
            if not gdf_clipped.empty:
                gdf_clipped.boundary.plot(ax=ax, color="red", linewidth=1)

        plt.title("Rasterized Vector Data")
        plt.tight_layout()
        plt.show()

    return raster_data

visualize_vector_by_attribute(vector_path, attribute_name, cmap='viridis', figsize=(10, 8))

Create a thematic map visualization of vector data based on an attribute.

Parameters:

Name	Type	Description	Default
vector_path	str	Path to the vector file	required
attribute_name	str	Name of the attribute to visualize	required
cmap	str	Matplotlib colormap name. Defaults to 'viridis'.	'viridis'
figsize	tuple	Figure size as (width, height). Defaults to (10, 8).	(10, 8)

Returns:

Name	Type	Description
bool	bool	True if visualization was successful, False otherwise

Source code in geoai/utils/vector.py

def visualize_vector_by_attribute(
    vector_path: str,
    attribute_name: str,
    cmap: str = "viridis",
    figsize: Tuple[int, int] = (10, 8),
) -> bool:
    """Create a thematic map visualization of vector data based on an attribute.

    Args:
        vector_path (str): Path to the vector file
        attribute_name (str): Name of the attribute to visualize
        cmap (str, optional): Matplotlib colormap name. Defaults to 'viridis'.
        figsize (tuple, optional): Figure size as (width, height). Defaults to (10, 8).

    Returns:
        bool: True if visualization was successful, False otherwise
    """
    try:
        # Read the vector data
        gdf = gpd.read_file(vector_path)

        # Check if attribute exists
        if attribute_name not in gdf.columns:
            print(f"Attribute '{attribute_name}' not found in the dataset")
            return False

        # Create the plot
        fig, ax = plt.subplots(figsize=figsize)

        # Determine plot type based on data type
        if pd.api.types.is_numeric_dtype(gdf[attribute_name]):
            # Continuous data
            gdf.plot(column=attribute_name, cmap=cmap, legend=True, ax=ax)
        else:
            # Categorical data
            gdf.plot(column=attribute_name, categorical=True, legend=True, ax=ax)

        # Add title and labels
        ax.set_title(f"{os.path.basename(vector_path)} - {attribute_name}")
        ax.set_xlabel("Longitude")
        ax.set_ylabel("Latitude")

        # Add basemap or additional elements if available
        # Note: Additional options could be added here for more complex maps

        plt.tight_layout()
        plt.show()

    except Exception as e:
        print(f"Error visualizing data: {str(e)}")

write_colormap(image, colormap, output=None)

Write a colormap to an image.

Parameters:

Name	Type	Description	Default
image	Union[str, ndarray]	The image to write the colormap to.	required
colormap	Union[str, Dict]	The colormap to write to the image.	required
output	Optional[str]	The output file path.	None

Source code in geoai/utils/raster.py

def write_colormap(
    image: Union[str, np.ndarray],
    colormap: Union[str, Dict],
    output: Optional[str] = None,
) -> Optional[str]:
    """Write a colormap to an image.

    Args:
        image: The image to write the colormap to.
        colormap: The colormap to write to the image.
        output: The output file path.
    """
    if isinstance(colormap, str):
        colormap = leafmap.get_image_colormap(colormap)
    leafmap.write_image_colormap(image, colormap, output)