segment-geospatial 1.0.2__tar.gz → 1.1.0__tar.gz

{segment_geospatial-1.0.2 → segment_geospatial-1.1.0}/.github/workflows/ubuntu.yml

@@ -48,6 +48,8 @@ jobs:
                   pip install --user -r requirements.txt
                   pip install --user -r requirements_dev.txt
                   pip install --user .
-            - name: PKG-TEST
-              run: |
-                  python -m unittest discover tests/
+            - name: Test import
+              run: python -c "import samgeo; print('samgeo import successful')"
+            # - name: PKG-TEST
+            #   run: |
+            #       python -m unittest discover tests/

{segment_geospatial-1.0.2 → segment_geospatial-1.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: segment-geospatial
-Version: 1.0.2
+Version: 1.1.0
 Summary: Meta AI' Segment Anything Model (SAM) for Geospatial Data.
 Author-email: Qiusheng Wu <giswqs@gmail.com>
 License: MIT license

{segment_geospatial-1.0.2 → segment_geospatial-1.1.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "segment-geospatial"
-version = "1.0.2"
+version = "1.1.0"
 dynamic = [
     "dependencies",
 ]
@@ -111,7 +111,7 @@ dependencies = {file = ["requirements.txt"]}
 universal = true
 
 [tool.bumpversion]
-current_version = "1.0.2"
+current_version = "1.1.0"
 commit = true
 tag = true
 

{segment_geospatial-1.0.2 → segment_geospatial-1.1.0}/samgeo/__init__.py

@@ -2,7 +2,7 @@
 
 __author__ = """Qiusheng Wu"""
 __email__ = "giswqs@gmail.com"
-__version__ = "1.0.2"
+__version__ = "1.1.0"
 
 
 from .samgeo import *  # noqa: F403

{segment_geospatial-1.0.2 → segment_geospatial-1.1.0}/samgeo/common.py

@@ -3737,27 +3737,50 @@ def make_temp_dir(**kwargs) -> str:
     return temp_dir
 
 
-def geotiff_to_jpg(geotiff_path: str, output_path: str) -> None:
+def geotiff_to_jpg(
+    geotiff_path: str,
+    output_path: str,
+    bands: Optional[List[int]] = None,
+) -> None:
     """Convert a GeoTIFF file to a JPG file.
 
     Args:
         geotiff_path (str): The path to the input GeoTIFF file.
         output_path (str): The path to the output JPG file.
+        bands (List[int], optional): List of band indices (1-based) to use for RGB.
+            For example, [4, 3, 2] for NIR-R-G false color composite.
+            If None, uses the first 3 bands. Defaults to None.
     """
 
     from PIL import Image
 
     # Open the GeoTIFF file
     with rasterio.open(geotiff_path) as src:
-        # Read the first band (for grayscale) or all bands
-        array = src.read()
-
-        # If the array has more than 3 bands, reduce it to the first 3 (RGB)
-        if array.shape[0] >= 3:
-            array = array[:3, :, :]  # Select the first 3 bands (R, G, B)
-        elif array.shape[0] == 1:
-            # For single-band images, repeat the band to create a grayscale RGB
-            array = np.repeat(array, 3, axis=0)
+        if bands is not None:
+            # Validate band indices (1-based)
+            if len(bands) != 3:
+                raise ValueError("bands must contain exactly 3 band indices for RGB.")
+            for band in bands:
+                if band < 1 or band > src.count:
+                    raise ValueError(
+                        f"Band index {band} is out of range. "
+                        f"Image has {src.count} bands (1-indexed)."
+                    )
+            # Read specified bands (rasterio uses 1-based indexing)
+            array = np.stack([src.read(b) for b in bands], axis=0)
+        else:
+            # Read all bands
+            array = src.read()
+
+            # If the array has more than 3 bands, reduce it to the first 3 (RGB)
+            if array.shape[0] >= 3:
+                array = array[:3, :, :]  # Select the first 3 bands (R, G, B)
+            elif array.shape[0] == 1:
+                # For single-band images, repeat the band to create a grayscale RGB
+                array = np.repeat(array, 3, axis=0)
+            elif array.shape[0] == 2:
+                # For two-band images, repeat the first band to create a 3-band image
+                array = np.concatenate([array, array[0:1, :, :]], axis=0)
 
         # Transpose the array from (bands, height, width) to (height, width, bands)
         array = np.transpose(array, (1, 2, 0))
@@ -3774,12 +3797,19 @@ def geotiff_to_jpg(geotiff_path: str, output_path: str) -> None:
         image.save(output_path)
 
 
-def geotiff_to_jpg_batch(input_folder: str, output_folder: str = None) -> str:
+def geotiff_to_jpg_batch(
+    input_folder: str,
+    output_folder: str = None,
+    bands: Optional[List[int]] = None,
+) -> str:
     """Convert all GeoTIFF files in a folder to JPG files.
 
     Args:
         input_folder (str): The path to the folder containing GeoTIFF files.
         output_folder (str): The path to the folder to save the output JPG files.
+        bands (List[int], optional): List of band indices (1-based) to use for RGB.
+            For example, [4, 3, 2] for NIR-R-G false color composite.
+            If None, uses the first 3 bands. Defaults to None.
 
     Returns:
         str: The path to the output folder containing the JPG files.
@@ -3800,7 +3830,7 @@ def geotiff_to_jpg_batch(input_folder: str, output_folder: str = None) -> str:
         geotiff_path = os.path.join(input_folder, filename)
         jpg_filename = os.path.splitext(filename)[0] + ".jpg"
         output_path = os.path.join(output_folder, jpg_filename)
-        geotiff_to_jpg(geotiff_path, output_path)
+        geotiff_to_jpg(geotiff_path, output_path, bands=bands)
 
     return output_folder
 

{segment_geospatial-1.0.2 → segment_geospatial-1.1.0}/samgeo/samgeo3.py

@@ -244,6 +244,7 @@ class SamGeo3:
         self,
         image: Union[str, np.ndarray],
         state=None,
+        bands: Optional[List[int]] = None,
     ) -> None:
         """Set the input image as a numpy array.
 
@@ -251,6 +252,10 @@ class SamGeo3:
             image (Union[str, np.ndarray, Image]): The input image as a path,
                 a numpy array, or an Image.
             state (optional): An optional state object to pass to the processor's set_image method (Meta backend only).
+            bands (List[int], optional): List of band indices (1-based) to use for RGB
+                when the input is a GeoTIFF with more than 3 bands. For example,
+                [4, 3, 2] for NIR-R-G false color composite. If None, uses the
+                first 3 bands for multi-band images. Defaults to None.
         """
         if isinstance(image, str):
             if image.startswith("http"):
@@ -260,8 +265,50 @@ class SamGeo3:
                 raise ValueError(f"Input path {image} does not exist.")
 
             self.source = image
-            image = cv2.imread(image)
-            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+            # Check if image is a GeoTIFF and handle band selection
+            if image.lower().endswith((".tif", ".tiff")):
+                import rasterio
+
+                with rasterio.open(image) as src:
+                    if bands is not None:
+                        # Validate band indices (1-based)
+                        if len(bands) != 3:
+                            raise ValueError(
+                                "bands must contain exactly 3 band indices for RGB."
+                            )
+                        for band in bands:
+                            if band < 1 or band > src.count:
+                                raise ValueError(
+                                    f"Band index {band} is out of range. "
+                                    f"Image has {src.count} bands (1-indexed)."
+                                )
+                        # Read specified bands (rasterio uses 1-based indexing)
+                        array = np.stack([src.read(b) for b in bands], axis=0)
+                    else:
+                        # Read all bands
+                        array = src.read()
+                        # If more than 3 bands, use first 3
+                        if array.shape[0] >= 3:
+                            array = array[:3, :, :]
+                        elif array.shape[0] == 1:
+                            array = np.repeat(array, 3, axis=0)
+                        elif array.shape[0] == 2:
+                            # Repeat the first band to make 3 bands: [band1, band2, band1]
+                            array = np.concatenate([array, array[0:1, :, :]], axis=0)
+                    # Transpose from (bands, height, width) to (height, width, bands)
+                    array = np.transpose(array, (1, 2, 0))
+
+                    # Normalize to 8-bit (0-255) range
+                    array = array.astype(np.float32)
+                    array -= array.min()
+                    if array.max() > 0:
+                        array /= array.max()
+                    array *= 255
+                    image = array.astype(np.uint8)
+            else:
+                image = cv2.imread(image)
+                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
             self.image = image
         elif isinstance(image, np.ndarray):
             self.image = image
@@ -3467,7 +3514,7 @@ class SamGeo3Video:
             if not os.path.exists(bpe_path):
                 bpe_dir = os.path.dirname(bpe_path)
                 os.makedirs(bpe_dir, exist_ok=True)
-                url = "https://github.com/facebookresearch/sam3/raw/refs/heads/main/assets/bpe_simple_vocab_16e6.txt.gz"
+                url = "https://huggingface.co/datasets/giswqs/geospatial/resolve/main/bpe_simple_vocab_16e6.txt.gz"
                 bpe_path = common.download_file(url, bpe_path, quiet=True)
 
         print(f"Using GPUs: {gpus_to_use}")
@@ -3492,6 +3539,7 @@ class SamGeo3Video:
         output_dir: Optional[str] = None,
         frame_rate: Optional[int] = None,
         prefix: str = "",
+        bands: Optional[List[int]] = None,
     ) -> None:
         """Load a video or time series images for segmentation.
 
@@ -3507,12 +3555,17 @@ class SamGeo3Video:
             frame_rate (int, optional): Frame rate for extracting frames from video.
                 Only used when video_path is an MP4 file. Defaults to None.
             prefix (str): Prefix for extracted frame filenames. Defaults to "".
+            bands (List[int], optional): List of band indices (1-based) to use for RGB
+                when the input is a GeoTIFF directory with multi-band images. For example,
+                [4, 3, 2] for NIR-R-G false color composite. If None, uses the
+                first 3 bands for multi-band images. Defaults to None.
 
         Example:
             >>> sam = SamGeo3Video()
             >>> sam.set_video("video.mp4")  # Load MP4 video
             >>> sam.set_video("frames/")  # Load from JPEG frames directory
             >>> sam.set_video("landsat_ts/")  # Load GeoTIFF time series
+            >>> sam.set_video("landsat_ts/", bands=[4, 3, 2])  # NIR-R-G composite
         """
         if isinstance(video_path, str):
             if video_path.startswith("http"):
@@ -3541,7 +3594,7 @@ class SamGeo3Video:
                     self._tif_dir = video_path
                     self._tif_names = files
                     # Convert GeoTIFFs to JPEGs for SAM3
-                    video_path = common.geotiff_to_jpg_batch(video_path)
+                    video_path = common.geotiff_to_jpg_batch(video_path, bands=bands)
                     print(f"Converted GeoTIFFs to JPEGs: {video_path}")
 
             if not os.path.exists(video_path):

{segment_geospatial-1.0.2 → segment_geospatial-1.1.0}/segment_geospatial.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: segment-geospatial
-Version: 1.0.2
+Version: 1.1.0
 Summary: Meta AI' Segment Anything Model (SAM) for Geospatial Data.
 Author-email: Qiusheng Wu <giswqs@gmail.com>
 License: MIT license