geoai-py 0.3.5__py2.py3-none-any.whl → 0.4.0__py2.py3-none-any.whl

geoai/__init__.py

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

geoai/download.py

@@ -1,18 +1,19 @@
 """This module provides functions to download data, including NAIP imagery and building data from Overture Maps."""
 
+import logging
 import os
-from typing import List, Tuple, Optional, Dict, Any
-import rioxarray
-import numpy as np
+import subprocess
+from typing import Any, Dict, List, Optional, Tuple
+
+import geopandas as gpd
 import matplotlib.pyplot as plt
-from pystac_client import Client
+import numpy as np
 import planetary_computer as pc
-import geopandas as gpd
+import requests
+import rioxarray
+from pystac_client import Client
 from shapely.geometry import box
 from tqdm import tqdm
-import requests
-import subprocess
-import logging
 
 # Configure logging
 logging.basicConfig(

geoai/extract.py

@@ -1,19 +1,26 @@
+"""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 torchvision.models.detection import (
+    fasterrcnn_resnet50_fpn_v2,
+    maskrcnn_resnet50_fpn,
+)
 from tqdm import tqdm
 
-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:
@@ -60,6 +67,7 @@ class CustomDataset(NonGeoDataset):
         chip_size=(512, 512),
         overlap=0.5,
         transforms=None,
+        band_indexes=None,
         verbose=False,
     ):
         """
@@ -70,6 +78,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 +91,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 +201,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:
@@ -256,7 +269,9 @@ class ObjectDetector:
     Object extraction using Mask R-CNN with TorchGeo.
     """
 
-    def __init__(self, model_path=None, repo_id=None, model=None, device=None):
+    def __init__(
+        self, model_path=None, repo_id=None, model=None, num_classes=2, device=None
+    ):
         """
         Initialize the object extractor.
 
@@ -264,6 +279,7 @@ class ObjectDetector:
             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
@@ -283,7 +299,7 @@ class ObjectDetector:
         self.simplify_tolerance = 1.0  # Tolerance for polygon simplification
 
         # Initialize model
-        self.model = self.initialize_model(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)):
@@ -328,11 +344,12 @@ class ObjectDetector:
             print("Please specify a local model path or ensure internet connectivity.")
             raise
 
-    def initialize_model(self, model):
+    def initialize_model(self, model, num_classes=2):
         """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.
@@ -347,7 +364,7 @@ class ObjectDetector:
             model = maskrcnn_resnet50_fpn(
                 weights=None,
                 progress=False,
-                num_classes=2,  # Background + object
+                num_classes=num_classes,  # Background + object
                 weights_backbone=None,
                 # These parameters ensure consistent normalization
                 image_mean=image_mean,
@@ -594,7 +611,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 +796,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 +812,7 @@ class ObjectDetector:
         batch_size=4,
         filter_edges=True,
         edge_buffer=20,
+        band_indexes=None,
         **kwargs,
     ):
         """
@@ -799,10 +820,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 +865,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 +1046,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
@@ -1281,13 +1309,14 @@ class ObjectDetector:
         Returns:
             GeoDataFrame with regularized objects
         """
+        import math
+
+        import cv2
+        import geopandas as gpd
         import numpy as np
-        from shapely.geometry import Polygon, MultiPolygon, box
         from shapely.affinity import rotate, translate
-        import geopandas as gpd
-        import math
+        from shapely.geometry import MultiPolygon, Polygon, box
         from tqdm import tqdm
-        import cv2
 
         def get_angle(p1, p2, p3):
             """Calculate angle between three points in degrees (0-180)"""
@@ -1887,21 +1916,30 @@ class ObjectDetector:
         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:
@@ -1926,6 +1964,7 @@ class ObjectDetector:
                 raster_path=raster_path,
                 chip_size=chip_size,
                 overlap=overlap,
+                band_indexes=band_indexes,
                 verbose=verbose,
             )
 
@@ -2012,6 +2051,21 @@ class 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
@@ -2062,7 +2116,7 @@ class ObjectDetector:
         output_path=None,
         confidence_threshold=0.5,
         min_object_area=100,
-        max_object_size=None,
+        max_object_area=None,
         **kwargs,
     ):
         """
@@ -2073,7 +2127,7 @@ class ObjectDetector:
             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
+            max_object_area: Maximum area in pixels to keep an object. Default: None
             **kwargs: Additional parameters
 
         Returns:
@@ -2097,8 +2151,9 @@ class ObjectDetector:
             print(f"Found {num_features} connected components")
 
             # Process each component
-            car_polygons = []
-            car_confidences = []
+            polygons = []
+            confidences = []
+            pixels = []
 
             # Add progress bar
             for label in tqdm(range(1, num_features + 1), desc="Processing components"):
@@ -2129,8 +2184,8 @@ class ObjectDetector:
                     if area < min_object_area:
                         continue
 
-                    if max_object_size is not None:
-                        if area > max_object_size:
+                    if max_object_area is not None:
+                        if area > max_object_area:
                             continue
 
                     # Get minimum area rectangle
@@ -2147,16 +2202,18 @@ class ObjectDetector:
                     poly = Polygon(geo_points)
 
                     # Add to lists
-                    car_polygons.append(poly)
-                    car_confidences.append(confidence)
+                    polygons.append(poly)
+                    confidences.append(confidence)
+                    pixels.append(area)
 
             # Create GeoDataFrame
-            if car_polygons:
+            if polygons:
                 gdf = gpd.GeoDataFrame(
                     {
-                        "geometry": car_polygons,
-                        "confidence": car_confidences,
-                        "class": [1] * len(car_polygons),
+                        "geometry": polygons,
+                        "confidence": confidences,
+                        "class": [1] * len(polygons),
+                        "pixels": pixels,
                     },
                     crs=crs,
                 )
@@ -2164,11 +2221,11 @@ class 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:
-                print("No valid car polygons found")
+                print("No valid polygons found")
                 return None
 
 
@@ -2277,3 +2334,66 @@ 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
+        )
+
+
+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="parking_spot_detection.pth",
+        repo_id=None,
+        model=None,
+        num_classes=3,
+        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.
+            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,
+        )

geoai/geoai.py

@@ -1,4 +1,7 @@
 """Main module."""
 
-from .utils import *
 from .extract import *
+from .hf import *
+from .segment import *
+from .utils import *
+from .train import train_MaskRCNN_model, object_detection