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

geoai/__init__.py

@@ -2,34 +2,96 @@
 
 __author__ = """Qiusheng Wu"""
 __email__ = "giswqs@gmail.com"
-__version__ = "0.3.6"
+__version__ = "0.4.1"
 
 
 import os
 import sys
 
 
-def set_proj_lib_path():
-    """Set the PROJ_LIB environment variable based on the current conda environment."""
+def set_proj_lib_path(verbose=False):
+    """
+    Set the PROJ_LIB and GDAL_DATA environment variables based on the current conda environment.
+
+    This function attempts to locate and set the correct paths for PROJ_LIB and GDAL_DATA
+    by checking multiple possible locations within the conda environment structure.
+
+    Args:
+        verbose (bool): If True, print additional information during the process.
+
+    Returns:
+        bool: True if both paths were set successfully, False otherwise.
+    """
     try:
         # Get conda environment path
         conda_env_path = os.environ.get("CONDA_PREFIX") or sys.prefix
 
+        # Define possible paths for PROJ_LIB
+        possible_proj_paths = [
+            os.path.join(conda_env_path, "share", "proj"),
+            os.path.join(conda_env_path, "Library", "share", "proj"),
+            os.path.join(conda_env_path, "Library", "share"),
+        ]
+
+        # Define possible paths for GDAL_DATA
+        possible_gdal_paths = [
+            os.path.join(conda_env_path, "share", "gdal"),
+            os.path.join(conda_env_path, "Library", "share", "gdal"),
+            os.path.join(conda_env_path, "Library", "data", "gdal"),
+            os.path.join(conda_env_path, "Library", "share"),
+        ]
+
         # Set PROJ_LIB environment variable
-        proj_path = os.path.join(conda_env_path, "share", "proj")
-        gdal_path = os.path.join(conda_env_path, "share", "gdal")
-
-        # Check if the directory exists before setting
-        if os.path.exists(proj_path):
-            os.environ["PROJ_LIB"] = proj_path
-        if os.path.exists(gdal_path):
-            os.environ["GDAL_DATA"] = gdal_path
+        proj_set = False
+        for proj_path in possible_proj_paths:
+            if os.path.exists(proj_path) and os.path.isdir(proj_path):
+                # Verify it contains projection data
+                if os.path.exists(os.path.join(proj_path, "proj.db")):
+                    os.environ["PROJ_LIB"] = proj_path
+                    if verbose:
+                        print(f"PROJ_LIB set to: {proj_path}")
+                    proj_set = True
+                    break
+
+        # Set GDAL_DATA environment variable
+        gdal_set = False
+        for gdal_path in possible_gdal_paths:
+            if os.path.exists(gdal_path) and os.path.isdir(gdal_path):
+                # Verify it contains the header.dxf file or other critical GDAL files
+                if os.path.exists(
+                    os.path.join(gdal_path, "header.dxf")
+                ) or os.path.exists(os.path.join(gdal_path, "gcs.csv")):
+                    os.environ["GDAL_DATA"] = gdal_path
+                    if verbose:
+                        print(f"GDAL_DATA set to: {gdal_path}")
+                    gdal_set = True
+                    break
+
+        # If paths still not found, try a last-resort approach
+        if not proj_set or not gdal_set:
+            # Try a deep search in the conda environment
+            for root, dirs, files in os.walk(conda_env_path):
+                if not gdal_set and "header.dxf" in files:
+                    os.environ["GDAL_DATA"] = root
+                    if verbose:
+                        print(f"GDAL_DATA set to: {root} (deep search)")
+                    gdal_set = True
+
+                if not proj_set and "proj.db" in files:
+                    os.environ["PROJ_LIB"] = root
+                    if verbose:
+                        print(f"PROJ_LIB set to: {root} (deep search)")
+                    proj_set = True
+
+                if proj_set and gdal_set:
+                    break
+
     except Exception as e:
-        print(e)
+        print(f"Error setting projection library paths: {e}")
         return
 
 
-if "google.colab" not in sys.modules:
-    set_proj_lib_path()
+# if ("google.colab" not in sys.modules) and (sys.platform != "windows"):
+#     set_proj_lib_path()
 
 from .geoai import *

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

@@ -14,16 +14,15 @@ import torch
 from huggingface_hub import hf_hub_download
 from rasterio.windows import Window
 from shapely.geometry import Polygon, box
-from tqdm import tqdm
 from torchvision.models.detection import (
-    maskrcnn_resnet50_fpn,
     fasterrcnn_resnet50_fpn_v2,
+    maskrcnn_resnet50_fpn,
 )
+from tqdm import tqdm
 
 # Local Imports
 from .utils import get_raster_stats
 
-
 try:
     from torchgeo.datasets import NonGeoDataset
 except ImportError as e:
@@ -270,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.
 
@@ -278,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
@@ -297,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)):
@@ -342,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.
@@ -361,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,
@@ -1306,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)"""
@@ -2112,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,
     ):
         """
@@ -2123,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:
@@ -2147,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"):
@@ -2179,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
@@ -2197,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,
                 )
@@ -2218,7 +2225,7 @@ class ObjectDetector:
 
                 return gdf
             else:
-                print("No valid car polygons found")
+                print("No valid polygons found")
                 return None
 
 
@@ -2356,3 +2363,37 @@ class SolarPanelDetector(ObjectDetector):
         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,5 +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