geoai-py 0.3.5__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.5"
+__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
@@ -1887,21 +1912,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 +1960,7 @@ class ObjectDetector:
                 raster_path=raster_path,
                 chip_size=chip_size,
                 overlap=overlap,
+                band_indexes=band_indexes,
                 verbose=verbose,
             )
 
@@ -2012,6 +2047,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
@@ -2164,7 +2214,7 @@ 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:
@@ -2277,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 *

geoai/segment.py

@@ -0,0 +1,305 @@
+"""This module provides functionality for segmenting high-resolution satellite imagery using vision-language models."""
+
+import os
+import numpy as np
+import torch
+from tqdm import tqdm
+from PIL import Image
+import rasterio
+from rasterio.windows import Window
+from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
+
+
+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', 'cpu'). If None, will use CUDA if available.
+        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' or 'cpu').
+        tile_size (int): Size of tiles for processing.
+        overlap (int): Overlap between tiles.
+    """
+
+    def __init__(
+        self,
+        model_name="CIDAS/clipseg-rd64-refined",
+        device=None,
+        tile_size=512,
+        overlap=32,
+    ):
+        """
+        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', 'cpu'). If None, will use CUDA if available.
+            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
+
+        # Set device
+        if device is None:
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        else:
+            self.device = device
+
+        # Load model and processor
+        self.processor = CLIPSegProcessor.from_pretrained(model_name)
+        self.model = CLIPSegForImageSegmentation.from_pretrained(model_name).to(
+            self.device
+        )
+
+        print(f"Model loaded on {self.device}")
+
+    def segment_image(
+        self, 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)
+
+        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 if necessary (handling different satellite bands)
+                            if tile_data.shape[0] > 3:
+                                # Use first three bands for RGB representation
+                                rgb_tile = tile_data[:3].transpose(1, 2, 0)
+                                # Normalize data 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)
+                            elif tile_data.shape[0] == 1:
+                                # Create RGB from grayscale
+                                rgb_tile = np.repeat(
+                                    tile_data[0][:, :, np.newaxis], 3, axis=2
+                                )
+                                # Normalize if needed
+                                if rgb_tile.max() > 0:
+                                    rgb_tile = (
+                                        (rgb_tile - rgb_tile.min())
+                                        / (rgb_tile.max() - rgb_tile.min())
+                                        * 255
+                                    ).astype(np.uint8)
+                            else:
+                                # Already 3-channel, assume RGB
+                                rgb_tile = tile_data.transpose(1, 2, 0)
+                                # Normalize 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)
+
+                            # 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
+                                pil_image.thumbnail(
+                                    (self.tile_size, self.tile_size), Image.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.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:
+                            print(f"Error processing tile at ({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")
+
+            print(f"Segmentation saved to {output_path}")
+            return output_path
+
+    def segment_image_batch(
+        self,
+        input_paths,
+        output_dir,
+        text_prompt,
+        threshold=0.5,
+        smoothing_sigma=1.0,
+        suffix="_segmented",
+    ):
+        """
+        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