geoai-py 0.8.3__py2.py3-none-any.whl → 0.9.1__py2.py3-none-any.whl

geoai/geoai.py

@@ -1,12 +1,14 @@
 """Main module."""
 
 import logging
+from typing import Any, Dict, List, Optional
 
 logging.getLogger("maplibre").setLevel(logging.ERROR)
 
 import leafmap
 import leafmap.maplibregl as maplibregl
 
+from .change_detection import ChangeDetection
 from .classify import classify_image, classify_images, train_classifier
 from .download import (
     download_naip,
@@ -26,17 +28,17 @@ from .extract import *
 from .hf import *
 from .segment import *
 from .train import (
+    get_instance_segmentation_model,
+    instance_segmentation,
+    instance_segmentation_batch,
+    instance_segmentation_inference_on_geotiff,
     object_detection,
     object_detection_batch,
-    train_MaskRCNN_model,
-    train_segmentation_model,
     semantic_segmentation,
     semantic_segmentation_batch,
     train_instance_segmentation_model,
-    instance_segmentation,
-    instance_segmentation_batch,
-    get_instance_segmentation_model,
-    instance_segmentation_inference_on_geotiff,
+    train_MaskRCNN_model,
+    train_segmentation_model,
 )
 from .utils import *
 
@@ -44,7 +46,7 @@ from .utils import *
 class Map(leafmap.Map):
     """A subclass of leafmap.Map for GeoAI applications."""
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         """Initialize the Map class."""
         super().__init__(*args, **kwargs)
 
@@ -52,7 +54,7 @@ class Map(leafmap.Map):
 class MapLibre(maplibregl.Map):
     """A subclass of maplibregl.Map for GeoAI applications."""
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         """Initialize the MapLibre class."""
         super().__init__(*args, **kwargs)
 
@@ -73,7 +75,7 @@ def create_vector_data(
     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
@@ -157,9 +159,9 @@ def edit_vector_data(
     frame_border: int = 0,
     controls: Optional[List[str]] = None,
     position: str = "top-right",
-    fit_bounds_options: Dict = None,
+    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

geoai/hf.py

@@ -2,7 +2,7 @@
 
 import csv
 import os
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -12,7 +12,9 @@ from tqdm import tqdm
 from transformers import AutoConfig, AutoModelForMaskedImageModeling, pipeline
 
 
-def get_model_config(model_id):
+def get_model_config(
+    model_id: str,
+) -> "transformers.configuration_utils.PretrainedConfig":
     """
     Get the model configuration for a Hugging Face model.
 
@@ -25,7 +27,7 @@ def get_model_config(model_id):
     return AutoConfig.from_pretrained(model_id)
 
 
-def get_model_input_channels(model_id):
+def get_model_input_channels(model_id: str) -> int:
     """
     Check the number of input channels supported by a Hugging Face model.
 
@@ -63,14 +65,14 @@ def get_model_input_channels(model_id):
 
 
 def image_segmentation(
-    tif_path,
-    output_path,
-    labels_to_extract=None,
-    dtype="uint8",
-    model_name=None,
-    segmenter_args=None,
-    **kwargs,
-):
+    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.
@@ -204,7 +206,7 @@ def mask_generation(
     band_indices: Optional[List[int]] = None,
     min_object_size: int = 0,
     generator_kwargs: Optional[Dict] = None,
-    **kwargs,
+    **kwargs: Any,
 ) -> Tuple[str, str]:
     """
     Process a GeoTIFF using SAM mask generation and save results as a GeoTIFF and CSV.

geoai/segment.py

@@ -2,23 +2,23 @@
 
 import os
 from dataclasses import dataclass
-from typing import Any, List, Dict, Optional, Union, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import cv2
+import geopandas as gpd
 import numpy as np
 import rasterio
 import torch
-import geopandas as gpd
 from PIL import Image
-from rasterio.windows import Window
 from rasterio.warp import transform_bounds
-from shapely.geometry import box, Polygon
+from rasterio.windows import Window
+from shapely.geometry import Polygon, box
 from tqdm import tqdm
 from transformers import (
-    CLIPSegForImageSegmentation,
-    CLIPSegProcessor,
     AutoModelForMaskGeneration,
     AutoProcessor,
+    CLIPSegForImageSegmentation,
+    CLIPSegProcessor,
     pipeline,
 )
 
@@ -90,13 +90,13 @@ class GroundedSAM:
 
     def __init__(
         self,
-        detector_id="IDEA-Research/grounding-dino-tiny",
-        segmenter_id="facebook/sam-vit-base",
-        device=None,
-        tile_size=1024,
-        overlap=128,
-        threshold=0.3,
-    ):
+        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.
 
@@ -125,7 +125,7 @@ class GroundedSAM:
 
         print(f"GroundedSAM initialized on {self.device}")
 
-    def _load_models(self):
+    def _load_models(self) -> None:
         """Load the Grounding DINO and SAM models."""
         # Load Grounding DINO
         self.object_detector = pipeline(
@@ -408,17 +408,17 @@ class GroundedSAM:
 
     def segment_image(
         self,
-        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,
-    ):
+        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.
 
@@ -742,11 +742,11 @@ class CLIPSegmentation:
 
     def __init__(
         self,
-        model_name="CIDAS/clipseg-rd64-refined",
-        device=None,
-        tile_size=512,
-        overlap=32,
-    ):
+        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.
 
@@ -774,8 +774,13 @@ class CLIPSegmentation:
         print(f"Model loaded on {self.device}")
 
     def segment_image(
-        self, input_path, output_path, text_prompt, threshold=0.5, smoothing_sigma=1.0
-    ):
+        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.
 
@@ -974,13 +979,13 @@ class CLIPSegmentation:
 
     def segment_image_batch(
         self,
-        input_paths,
-        output_dir,
-        text_prompt,
-        threshold=0.5,
-        smoothing_sigma=1.0,
-        suffix="_segmented",
-    ):
+        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.
 

geoai/segmentation.py

@@ -1,4 +1,5 @@
 import os
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import albumentations as A
 import matplotlib.pyplot as plt
@@ -24,10 +25,10 @@ class CustomDataset(Dataset):
         self,
         images_dir: str,
         masks_dir: str,
-        transform: A.Compose = None,
-        target_size: tuple = (256, 256),
+        transform: Optional[A.Compose] = None,
+        target_size: Tuple[int, int] = (256, 256),
         num_classes: int = 2,
-    ):
+    ) -> None:
         """
         Args:
             images_dir (str): Directory containing images.
@@ -48,7 +49,7 @@ class CustomDataset(Dataset):
         """Returns the total number of samples."""
         return len(self.images)
 
-    def __getitem__(self, idx: int) -> dict:
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
         """
         Args:
             idx (int): Index of the sample to fetch.
@@ -109,7 +110,7 @@ def prepare_datasets(
     transform: A.Compose,
     test_size: float = 0.2,
     random_state: int = 42,
-) -> tuple:
+) -> Tuple[Subset, Subset]:
     """
     Args:
         images_dir (str): Directory containing images.
@@ -233,7 +234,7 @@ def preprocess_image(image_path: str, target_size: tuple = (256, 256)) -> torch.
 def predict_image(
     model: SegformerForSemanticSegmentation,
     image_tensor: torch.Tensor,
-    original_size: tuple,
+    original_size: Tuple[int, int],
     device: torch.device,
 ) -> np.ndarray:
     """
@@ -262,7 +263,7 @@ def predict_image(
 def segment_image(
     image_path: str,
     model_path: str,
-    target_size: tuple = (256, 256),
+    target_size: Tuple[int, int] = (256, 256),
     device: torch.device = torch.device("cuda" if torch.cuda.is_available() else "cpu"),
 ) -> np.ndarray:
     """
@@ -288,8 +289,8 @@ def segment_image(
 def visualize_predictions(
     image_path: str,
     segmented_mask: np.ndarray,
-    target_size: tuple = (256, 256),
-    reference_image_path: str = None,
+    target_size: Tuple[int, int] = (256, 256),
+    reference_image_path: Optional[str] = None,
 ) -> None:
     """
     Visualizes the original image, segmented mask, and optionally the reference image.