geoai-py 0.8.0__py2.py3-none-any.whl → 0.8.2__py2.py3-none-any.whl

geoai/__init__.py

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

geoai/train.py

@@ -511,7 +511,7 @@ def visualize_predictions(model, dataset, device, num_samples=5, output_dir=None
 
     # Create output directory if needed
     if output_dir:
-        os.makedirs(output_dir, exist_ok=True)
+        os.makedirs(os.path.abspath(output_dir), exist_ok=True)
 
     # Select random samples
     indices = random.sample(range(len(dataset)), min(num_samples, len(dataset)))
@@ -648,7 +648,7 @@ def train_MaskRCNN_model(
     torch.backends.cudnn.benchmark = False
 
     # Create output directory
-    os.makedirs(output_dir, exist_ok=True)
+    os.makedirs(os.path.abspath(output_dir), exist_ok=True)
 
     # Get device
     if device is None:
@@ -1216,7 +1216,7 @@ def object_detection_batch(
         )
 
     if not os.path.exists(output_dir):
-        os.makedirs(output_dir, exist_ok=True)
+        os.makedirs(os.path.abspath(output_dir), exist_ok=True)
 
     if not os.path.exists(model_path):
         try:
@@ -1905,7 +1905,7 @@ def train_segmentation_model(
     torch.backends.cudnn.benchmark = False
 
     # Create output directory
-    os.makedirs(output_dir, exist_ok=True)
+    os.makedirs(os.path.abspath(output_dir), exist_ok=True)
 
     # Get device
     if device is None:
@@ -2511,7 +2511,7 @@ def semantic_inference_on_geotiff(
             print(f"Inference completed in {inference_time:.2f} seconds")
 
         # Save output
-        out_dir = os.path.dirname(output_path)
+        out_dir = os.path.abspath(os.path.dirname(output_path))
         os.makedirs(out_dir, exist_ok=True)
         with rasterio.open(output_path, "w", **out_meta) as dst:
             dst.write(mask, 1)
@@ -2778,7 +2778,7 @@ def semantic_inference_on_image(
             # Change extension to PNG if binary output to preserve exact values
             output_path_png = os.path.splitext(output_path)[0] + ".png"
             output_img = Image.fromarray(mask, mode="L")
-            out_dir = os.path.dirname(output_path)
+            out_dir = os.path.abspath(os.path.dirname(output_path))
             os.makedirs(out_dir, exist_ok=True)
             output_img.save(output_path_png)
             if not quiet:
@@ -2887,7 +2887,7 @@ def semantic_segmentation(
         )
     else:
         # Create output directory if it doesn't exist
-        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        os.makedirs(os.path.abspath(os.path.dirname(output_path)), exist_ok=True)
 
         semantic_inference_on_image(
             model=model,
@@ -2962,7 +2962,7 @@ def semantic_segmentation_batch(
         raise FileNotFoundError(f"Input directory does not exist: {input_dir}")
 
     # Create output directory if it doesn't exist
-    os.makedirs(output_dir, exist_ok=True)
+    os.makedirs(os.path.abspath(output_dir), exist_ok=True)
 
     # Get all supported image files
     image_extensions = (".tif", ".tiff", ".png", ".jpg", ".jpeg")

geoai/utils.py

@@ -5,6 +5,7 @@ import glob
 import json
 import math
 import os
+import subprocess
 import warnings
 import xml.etree.ElementTree as ET
 from collections import OrderedDict
@@ -7279,106 +7280,213 @@ def plot_prediction_comparison(
     prediction_colormap: str = "gray",
     ground_truth_colormap: str = "gray",
     original_colormap: Optional[str] = None,
+    indexes: Optional[List[int]] = None,
+    divider: Optional[float] = None,
 ):
-    """
-    Plot original image, prediction image, and optionally ground truth image side by side.
+    """Plot original image, prediction, and optional ground truth side by side.
+
+    Supports input as file paths, NumPy arrays, or PIL Images. For multi-band
+    images, selected channels can be specified via `indexes`. If the image data
+    is not normalized (e.g., Sentinel-2 [0, 10000]), the `divider` can be used
+    to scale values for visualization.
 
     Args:
-        original_image: Original input image (file path, numpy array, or PIL Image)
-        prediction_image: Prediction/segmentation mask (file path, numpy array, or PIL Image)
-        ground_truth_image: Optional ground truth mask (file path, numpy array, or PIL Image)
-        titles: Optional list of titles for each subplot
-        figsize: Figure size tuple (width, height)
-        save_path: Optional path to save the plot
-        show_plot: Whether to display the plot
-        prediction_colormap: Colormap for prediction image
-        ground_truth_colormap: Colormap for ground truth image
-        original_colormap: Colormap for original image (None for RGB)
+        original_image (Union[str, np.ndarray, Image.Image]):
+            Original input image as a file path, NumPy array, or PIL Image.
+        prediction_image (Union[str, np.ndarray, Image.Image]):
+            Predicted segmentation mask image.
+        ground_truth_image (Optional[Union[str, np.ndarray, Image.Image]], optional):
+            Ground truth mask image. Defaults to None.
+        titles (Optional[List[str]], optional):
+            List of titles for the subplots. If not provided, default titles are used.
+        figsize (Tuple[int, int], optional):
+            Size of the entire figure in inches. Defaults to (15, 5).
+        save_path (Optional[str], optional):
+            If specified, saves the figure to this path. Defaults to None.
+        show_plot (bool, optional):
+            Whether to display the figure using plt.show(). Defaults to True.
+        prediction_colormap (str, optional):
+            Colormap to use for the prediction mask. Defaults to "gray".
+        ground_truth_colormap (str, optional):
+            Colormap to use for the ground truth mask. Defaults to "gray".
+        original_colormap (Optional[str], optional):
+            Colormap to use for the original image if it's grayscale. Defaults to None.
+        indexes (Optional[List[int]], optional):
+            List of band/channel indexes (0-based for NumPy, 1-based for rasterio) to extract from the original image.
+            Useful for multi-band imagery like Sentinel-2. Defaults to None.
+        divider (Optional[float], optional):
+            Value to divide the original image by for normalization (e.g., 10000 for reflectance). Defaults to None.
 
     Returns:
-        matplotlib.figure.Figure: The figure object
+        matplotlib.figure.Figure:
+            The generated matplotlib figure object.
     """
 
-    def _load_image(img_input):
+    def _load_image(img_input, indexes=None):
         """Helper function to load image from various input types."""
         if isinstance(img_input, str):
-            # File path
             if img_input.lower().endswith((".tif", ".tiff")):
-                # Handle GeoTIFF files
                 with rasterio.open(img_input) as src:
-                    img = src.read()
-                    if img.shape[0] == 1:
-                        # Single band
-                        img = img[0]
+                    if indexes:
+                        img = src.read(indexes)  # 1-based
+                        img = (
+                            np.transpose(img, (1, 2, 0)) if len(indexes) > 1 else img[0]
+                        )
                     else:
-                        # Multi-band, transpose to (H, W, C)
-                        img = np.transpose(img, (1, 2, 0))
+                        img = src.read()
+                        if img.shape[0] == 1:
+                            img = img[0]
+                        else:
+                            img = np.transpose(img, (1, 2, 0))
             else:
-                # Regular image file
                 img = np.array(Image.open(img_input))
         elif isinstance(img_input, Image.Image):
-            # PIL Image
             img = np.array(img_input)
         elif isinstance(img_input, np.ndarray):
-            # NumPy array
             img = img_input
+            if indexes is not None and img.ndim == 3:
+                img = img[:, :, indexes]
         else:
             raise ValueError(f"Unsupported image type: {type(img_input)}")
-
         return img
 
     # Load images
-    original = _load_image(original_image)
+    original = _load_image(original_image, indexes=indexes)
     prediction = _load_image(prediction_image)
     ground_truth = (
         _load_image(ground_truth_image) if ground_truth_image is not None else None
     )
 
-    # Determine number of subplots
-    num_plots = 3 if ground_truth is not None else 2
+    # Apply divider normalization if requested
+    if divider is not None and isinstance(original, np.ndarray) and original.ndim == 3:
+        original = np.clip(original.astype(np.float32) / divider, 0, 1)
 
-    # Create figure and subplots
+    # Determine layout
+    num_plots = 3 if ground_truth is not None else 2
     fig, axes = plt.subplots(1, num_plots, figsize=figsize)
     if num_plots == 2:
         axes = [axes[0], axes[1]]
 
-    # Default titles
     if titles is None:
         titles = ["Original Image", "Prediction"]
         if ground_truth is not None:
             titles.append("Ground Truth")
 
-    # Plot original image
-    if len(original.shape) == 3 and original.shape[2] in [3, 4]:
-        # RGB or RGBA image
+    # Plot original
+    if original.ndim == 3 and original.shape[2] in [3, 4]:
         axes[0].imshow(original)
     else:
-        # Grayscale or single channel
         axes[0].imshow(original, cmap=original_colormap)
     axes[0].set_title(titles[0])
     axes[0].axis("off")
 
-    # Plot prediction image
+    # Prediction
     axes[1].imshow(prediction, cmap=prediction_colormap)
     axes[1].set_title(titles[1])
     axes[1].axis("off")
 
-    # Plot ground truth if provided
+    # Ground truth
     if ground_truth is not None:
         axes[2].imshow(ground_truth, cmap=ground_truth_colormap)
         axes[2].set_title(titles[2])
         axes[2].axis("off")
 
-    # Adjust layout
     plt.tight_layout()
 
-    # Save if requested
     if save_path:
         plt.savefig(save_path, dpi=300, bbox_inches="tight")
         print(f"Plot saved to: {save_path}")
 
-    # Show plot
     if show_plot:
         plt.show()
 
     return fig
+
+
+def get_raster_resolution(image_path: str) -> Tuple[float, float]:
+    """Get pixel resolution from the raster using rasterio.
+
+    Args:
+        image_path: The path to the raster image.
+
+    Returns:
+        A tuple of (x resolution, y resolution).
+    """
+    with rasterio.open(image_path) as src:
+        res = src.res
+    return res
+
+
+def stack_bands(
+    input_files: List[str],
+    output_file: str,
+    resolution: Optional[float] = None,
+    dtype: Optional[str] = None,  # e.g., "UInt16", "Float32"
+    temp_vrt: str = "stack.vrt",
+    overwrite: bool = False,
+    compress: str = "DEFLATE",
+    output_format: str = "COG",
+    extra_gdal_translate_args: Optional[List[str]] = None,
+) -> str:
+    """
+    Stack bands from multiple images into a single multi-band GeoTIFF.
+
+    Args:
+        input_files (List[str]): List of input image paths.
+        output_file (str): Path to the output stacked image.
+        resolution (float, optional): Output resolution. If None, inferred from first image.
+        dtype (str, optional): Output data type (e.g., "UInt16", "Float32").
+        temp_vrt (str): Temporary VRT filename.
+        overwrite (bool): Whether to overwrite the output file.
+        compress (str): Compression method.
+        output_format (str): GDAL output format (default is "COG").
+        extra_gdal_translate_args (List[str], optional): Extra arguments for gdal_translate.
+
+    Returns:
+        str: Path to the output file.
+    """
+    if not input_files:
+        raise ValueError("No input files provided.")
+
+    if os.path.exists(output_file) and not overwrite:
+        print(f"Output file already exists: {output_file}")
+        return output_file
+
+    # Infer resolution if not provided
+    if resolution is None:
+        resolution_x, resolution_y = get_raster_resolution(input_files[0])
+    else:
+        resolution_x = resolution_y = resolution
+
+    # Step 1: Build VRT
+    vrt_cmd = ["gdalbuildvrt", "-separate", temp_vrt] + input_files
+    subprocess.run(vrt_cmd, check=True)
+
+    # Step 2: Translate VRT to output GeoTIFF
+    translate_cmd = [
+        "gdal_translate",
+        "-tr",
+        str(resolution_x),
+        str(resolution_y),
+        temp_vrt,
+        output_file,
+        "-of",
+        output_format,
+        "-co",
+        f"COMPRESS={compress}",
+    ]
+
+    if dtype:
+        translate_cmd.insert(1, "-ot")
+        translate_cmd.insert(2, dtype)
+
+    if extra_gdal_translate_args:
+        translate_cmd += extra_gdal_translate_args
+
+    subprocess.run(translate_cmd, check=True)
+
+    # Step 3: Clean up VRT
+    if os.path.exists(temp_vrt):
+        os.remove(temp_vrt)
+
+    return output_file

{geoai_py-0.8.0.dist-info → geoai_py-0.8.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: geoai-py
-Version: 0.8.0
+Version: 0.8.2
 Summary: A Python package for using Artificial Intelligence (AI) with geospatial data
 Author-email: Qiusheng Wu <giswqs@gmail.com>
 License: MIT License

{geoai_py-0.8.0.dist-info → geoai_py-0.8.2.dist-info}/RECORD

@@ -1,4 +1,4 @@
-geoai/__init__.py,sha256=05t9D_IforM2DDpgmwh1BPzGOsmQqfRi1BC_-Ola9_w,3765
+geoai/__init__.py,sha256=wKdVPgcUCURuMI3TNzqh1MPw96Bs5XL7mNUST75bOjw,3765
 geoai/classify.py,sha256=O8fah3DOBDMZW7V_qfDYsUnjB-9Wo5fjA-0e4wvUeAE,35054
 geoai/download.py,sha256=EQpcrcqMsYhDpd7bpjf4hGS5xL2oO-jsjngLgjjP3cE,46599
 geoai/extract.py,sha256=vyHH1k5zaXiy1SLdCLXxbWiNLp8XKdu_MXZoREMtAOQ,119102
@@ -7,11 +7,11 @@ geoai/hf.py,sha256=mLKGxEAS5eHkxZLwuLpYc1o7e3-7QIXdBv-QUY-RkFk,17072
 geoai/sam.py,sha256=O6S-kGiFn7YEcFbfWFItZZQOhnsm6-GlunxQLY0daEs,34345
 geoai/segment.py,sha256=pThAyq8kjgVDhMwHMiWkZ2qL5ykzA5lRg7tyMmSEBxk,43434
 geoai/segmentation.py,sha256=AtPzCvguHAEeuyXafa4bzMFATvltEYcah1B8ZMfkM_s,11373
-geoai/train.py,sha256=bwcCjDwiWF3aGo2gRqhd-6Xlhg3xZV_SeHx0kskWsPM,112709
-geoai/utils.py,sha256=Ucards_Kb3H77bFn-pVpnt0f1oez-Vu3eFaaZhRV5D4,281943
-geoai_py-0.8.0.dist-info/licenses/LICENSE,sha256=vN2L5U7cZ6ZkOHFmc8WiGlsogWsZc5dllMeNxnKVOZg,1070
-geoai_py-0.8.0.dist-info/METADATA,sha256=iYkcRGlikK5Tdy1gBmiCBsd6t7tQhPi77e39btQMKgk,6661
-geoai_py-0.8.0.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
-geoai_py-0.8.0.dist-info/entry_points.txt,sha256=uGp3Az3HURIsRHP9v-ys0hIbUuBBNUfXv6VbYHIXeg4,41
-geoai_py-0.8.0.dist-info/top_level.txt,sha256=1YkCUWu-ii-0qIex7kbwAvfei-gos9ycyDyUCJPNWHY,6
-geoai_py-0.8.0.dist-info/RECORD,,
+geoai/train.py,sha256=Mrsb0yMVnprQHld3zDvA-puc-r8hGm1XgG0j2GGIn7E,112845
+geoai/utils.py,sha256=_JVEhFUzOdDS_Rmco8c4BJVnrs2VvUVHKf4ubtim5fg,286109
+geoai_py-0.8.2.dist-info/licenses/LICENSE,sha256=vN2L5U7cZ6ZkOHFmc8WiGlsogWsZc5dllMeNxnKVOZg,1070
+geoai_py-0.8.2.dist-info/METADATA,sha256=6LIaLEjT5AtbV5OOJwMT7NQCMI6pHcKXsfQbakWl6Ns,6661
+geoai_py-0.8.2.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
+geoai_py-0.8.2.dist-info/entry_points.txt,sha256=uGp3Az3HURIsRHP9v-ys0hIbUuBBNUfXv6VbYHIXeg4,41
+geoai_py-0.8.2.dist-info/top_level.txt,sha256=1YkCUWu-ii-0qIex7kbwAvfei-gos9ycyDyUCJPNWHY,6
+geoai_py-0.8.2.dist-info/RECORD,,