geoai-py 0.3.4__py2.py3-none-any.whl → 0.3.6__py2.py3-none-any.whl

geoai/__init__.py

@@ -2,7 +2,7 @@
 
 __author__ = """Qiusheng Wu"""
 __email__ = "giswqs@gmail.com"
-__version__ = "0.3.4"
+__version__ = "0.3.6"
 
 
 import os

geoai/extract.py

@@ -1,21 +1,29 @@
+"""This module provides a dataset class for object extraction from raster data"""
+
+# Standard Library
 import os
+
+# Third-Party Libraries
+import cv2
+import geopandas as gpd
+import matplotlib.pyplot as plt
 import numpy as np
+import rasterio
+import scipy.ndimage as ndimage
 import torch
-import matplotlib.pyplot as plt
+from huggingface_hub import hf_hub_download
+from rasterio.windows import Window
 from shapely.geometry import Polygon, box
-import geopandas as gpd
 from tqdm import tqdm
+from torchvision.models.detection import (
+    maskrcnn_resnet50_fpn,
+    fasterrcnn_resnet50_fpn_v2,
+)
 
-import cv2
-from torchvision.models.detection import maskrcnn_resnet50_fpn
-import torchvision.transforms as T
-import rasterio
-from rasterio.windows import Window
-from rasterio.features import shapes
-from huggingface_hub import hf_hub_download
-import scipy.ndimage as ndimage
+# Local Imports
 from .utils import get_raster_stats
 
+
 try:
     from torchgeo.datasets import NonGeoDataset
 except ImportError as e:
@@ -60,6 +68,7 @@ class CustomDataset(NonGeoDataset):
         chip_size=(512, 512),
         overlap=0.5,
         transforms=None,
+        band_indexes=None,
         verbose=False,
     ):
         """
@@ -70,6 +79,7 @@ class CustomDataset(NonGeoDataset):
             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:
@@ -82,6 +92,7 @@ class CustomDataset(NonGeoDataset):
         self.chip_size = chip_size
         self.overlap = overlap
         self.transforms = transforms
+        self.band_indexes = band_indexes
         self.verbose = verbose
         self.warned_about_bands = False
 
@@ -191,7 +202,10 @@ class CustomDataset(NonGeoDataset):
                 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
-                image = image[:3]
+                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:
@@ -594,7 +608,7 @@ class ObjectDetector:
 
         Args:
             mask_path: Path to the object masks GeoTIFF
-            output_path: Path to save the output GeoJSON (default: mask_path with .geojson extension)
+            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)
@@ -779,7 +793,10 @@ class ObjectDetector:
 
             # Save to file
             if output_path:
-                gdf.to_file(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
@@ -792,6 +809,7 @@ class ObjectDetector:
         batch_size=4,
         filter_edges=True,
         edge_buffer=20,
+        band_indexes=None,
         **kwargs,
     ):
         """
@@ -799,10 +817,11 @@ class ObjectDetector:
 
         Args:
             raster_path: Path to input raster file
-            output_path: Path to output GeoJSON file (optional)
+            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)
@@ -843,7 +862,10 @@ class ObjectDetector:
 
         # Create dataset
         dataset = CustomDataset(
-            raster_path=raster_path, chip_size=chip_size, overlap=overlap
+            raster_path=raster_path,
+            chip_size=chip_size,
+            overlap=overlap,
+            band_indexes=band_indexes,
         )
         self.raster_stats = dataset.raster_stats
 
@@ -1021,7 +1043,10 @@ class ObjectDetector:
 
         # Save to file if requested
         if output_path:
-            gdf.to_file(output_path, driver="GeoJSON")
+            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
@@ -1881,108 +1906,36 @@ class ObjectDetector:
                     plt.savefig(sample_output, dpi=300, bbox_inches="tight")
                     print(f"Sample visualization saved to {sample_output}")
 
-
-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="building_footprints_usa.pth",
-        repo_id=None,
-        model=None,
-        device=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,
-        min_area=10,
-        angle_threshold=15,
-        orthogonality_threshold=0.3,
-        rectangularity_threshold=0.7,
-    ):
-        """
-        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,
-        )
-
-
-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="car_detection_usa.pth", repo_id=None, model=None, device=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 generate_masks(
         self,
         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.
+
         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:
@@ -1994,6 +1947,8 @@ class CarDetector(ObjectDetector):
         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"
@@ -2003,8 +1958,9 @@ class CarDetector(ObjectDetector):
             # Create dataset with the specified overlap
             dataset = CustomDataset(
                 raster_path=raster_path,
-                chip_size=self.chip_size,
+                chip_size=chip_size,
                 overlap=overlap,
+                band_indexes=band_indexes,
                 verbose=verbose,
             )
 
@@ -2091,6 +2047,21 @@ class CarDetector(ObjectDetector):
                     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
@@ -2135,13 +2106,24 @@ class CarDetector(ObjectDetector):
             print(f"Masks with confidence values saved to {output_path}")
             return output_path
 
-    def vectorize_masks(self, masks_path, output_path=None, **kwargs):
+    def vectorize_masks(
+        self,
+        masks_path,
+        output_path=None,
+        confidence_threshold=0.5,
+        min_object_area=100,
+        max_object_size=None,
+        **kwargs,
+    ):
         """
         Convert masks with confidence to vector polygons.
 
         Args:
-            masks_path: Path to masks GeoTIFF with confidence band
-            output_path: Path for output GeoJSON
+            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_size: Maximum area in pixels to keep an object. Default: None
             **kwargs: Additional parameters
 
         Returns:
@@ -2182,6 +2164,10 @@ class CarDetector(ObjectDetector):
                 else:
                     confidence = 0.0
 
+                # Skip if confidence is below threshold
+                if confidence < confidence_threshold:
+                    continue
+
                 # Find contours
                 contours, _ = cv2.findContours(
                     component_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
@@ -2190,9 +2176,13 @@ class CarDetector(ObjectDetector):
                 for contour in contours:
                     # Filter by size
                     area = cv2.contourArea(contour)
-                    if area < kwargs.get("min_object_area", 100):
+                    if area < min_object_area:
                         continue
 
+                    if max_object_size is not None:
+                        if area > max_object_size:
+                            continue
+
                     # Get minimum area rectangle
                     rect = cv2.minAreaRect(contour)
                     box_points = cv2.boxPoints(rect)
@@ -2224,7 +2214,7 @@ class CarDetector(ObjectDetector):
                 # Save to file if requested
                 if output_path:
                     gdf.to_file(output_path, driver="GeoJSON")
-                    print(f"Saved {len(gdf)} cars with confidence to {output_path}")
+                    print(f"Saved {len(gdf)} objects with confidence to {output_path}")
 
                 return gdf
             else:
@@ -2232,6 +2222,88 @@ class CarDetector(ObjectDetector):
                 return None
 
 
+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="building_footprints_usa.pth",
+        repo_id=None,
+        model=None,
+        device=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,
+        min_area=10,
+        angle_threshold=15,
+        orthogonality_threshold=0.3,
+        rectangularity_threshold=0.7,
+    ):
+        """
+        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,
+        )
+
+
+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="car_detection_usa.pth", repo_id=None, model=None, device=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
+        )
+
+
 class ShipDetector(ObjectDetector):
     """
     Ship detection using a pre-trained Mask R-CNN model.
@@ -2255,3 +2327,32 @@ class ShipDetector(ObjectDetector):
         super().__init__(
             model_path=model_path, repo_id=repo_id, model=model, device=device
         )
+
+
+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="solar_panel_detection.pth",
+        repo_id=None,
+        model=None,
+        device=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
+        )

geoai/geoai.py

@@ -2,3 +2,4 @@
 
 from .utils import *
 from .extract import *
+from .segment import *