i.sam2: SamGeo2 model

src/imagery/i.sam2/Makefile

@@ -0,0 +1,7 @@
+MODULE_TOPDIR = ../..
+
+PGM = i.sam2
+
+include $(MODULE_TOPDIR)/include/Make/Script.make
+
+default: script

src/imagery/i.sam2/i.sam2.html

@@ -0,0 +1,50 @@
+<h2>DESCRIPTION</h2>
+
+<em>i.sam2</em> allows users to segment orthoimagery based on text prompts using <a href="https://samgeo.gishub.org/">SamGeo</a>.
+
+<h2>REQUIREMENTS</h2>
+
+<ul>
+    <li><a href="https://pillow.readthedocs.io/en/stable/">Pillow>=10.2.0</a></li>
+    <li><a href="https://numpy.org/">numpy>=1.26.1</a></li>
+    <li><a href="https://pytorch.org/">torch>=2.5.1</a></li>
+    <li><a href="https://samgeo.gishub.org/">segment-geospatial>=0.12.3</a></li>
+</ul>
+
+<div class="code">
+    <pre>
+        pip install pillow numpy torch segment-geospatial
+    </pre>
+</div>
+
+<h2>EXAMPLES</h2>
+
+Segment orthoimagery using SamGeo2:
+
+<div class="code">
+    <pre>
+    i.sam2 group=rgb_255 output=tree_mask text_prompt="trees"
+    </pre>
+</div>
+
+<img src="./i_sam2_trees.jpg" height="600" alt="i.sam2 example" />
+
+<h2>NOTES</h2>
+The first time use will be longer as the model needs to be downloaded. Subsequent runs will be faster.
+Additionally, Cuda is required for GPU acceleration. If you do not have a GPU, you can use the CPU by setting the environment variable `CUDA_VISIBLE_DEVICES` to `-1`.
+
+<h2>REFERENCES</h2>
+<ul>
+<li>Wu, Q., & Osco, L. (2023). samgeo: A Python package for segmenting geospatial data with the Segment Anything Model (SAM). Journal of Open Source Software, 8(89), 5663. <a href="https://doi.org/10.21105/joss.05663">https://doi.org/10.21105/joss.05663</a></li>
+<li>Osco, L. P., Wu, Q., de Lemos, E. L., Gonçalves, W. N., Ramos, A. P. M., Li, J., & Junior, J. M. (2023). The Segment Anything Model (SAM) for remote sensing applications: From zero to one shot. International Journal of Applied Earth Observation and Geoinformation, 124, 103540. <a href="https://doi.org/10.1016/j.jag.2023.103540">https://doi.org/10.1016/j.jag.2023.103540</a></li>
+</ul>
+
+<h2>SEE ALSO</h2>
+<em>
+    <a href="i.segment.gsoc.html">i.segment.gsoc</a> for region growing and merging segmentation,
+    <a href="i.segment.hierarchical">i.segment.hierarchical</a> performs a hierarchical segmentation,
+    <a href="i.superpixels.slic">i.superpixels.slic</a> for superpixel segmentation.
+</em>
+
+<h2>AUTHOR</h2>
+Corey T. White (NCSU GeoForAll Lab & OpenPlains Inc.)

src/imagery/i.sam2/i.sam2.py

@@ -0,0 +1,190 @@
+#!/usr/bin/env python3
+
+############################################################################
+#
+# MODULE:       i.sam2
+# AUTHOR:       Corey T. White, OpenPlains Inc.
+# PURPOSE:      Uses the SAMGeo model for segmentation in GRASS GIS.
+# COPYRIGHT:    (C) 2023-2024 Corey White
+#               This program is free software under the GNU General
+#               Public License (>=v2). Read the file COPYING that
+#               comes with GRASS for details.
+#
+#############################################################################
+
+# %module
+# % description: Integrates SAMGeo model with text prompt for segmentation in GRASS GIS.
+# % keyword: imagery
+# % keyword: segmentation
+# % keyword: object recognition
+# % keyword: deep learning
+# %end
+
+# %option G_OPT_I_GROUP
+# % key: group
+# % description: Name of input imagery group
+# % required: yes
+# %end
+
+# %option
+# % key: output
+# % type: string
+# % description: Name of output segmented raster map
+# % gisprompt: new,cell,raster
+# % required: yes
+# %end
+
+# %option
+# % key: model_path
+# % type: string
+# % description: Path to the SAMGeo model file (optional if using default model)
+# % required: no
+# %end
+
+# %option
+# % key: text_prompt
+# % type: string
+# % description: Text prompt to guide segmentation
+# % required: yes
+# %end
+
+# %option
+# % key: text_threshold
+# % type: double
+# % answer: 0.24
+# % description: Text threshold for text segmentation
+# % required: no
+# % multiple: no
+# %end
+
+# %option
+# % key: box_threshold
+# % type: double
+# % answer: 0.24
+# % description: Box threshold for text segmentation
+# % required: no
+# % multiple: no
+# %end
+
+
+import os
+import sys
+import grass.script as gs
+import torch
+import numpy as np
+from PIL import Image
+from grass.script import array as garray
+
+
+def main():
+    from samgeo import SamGeo
+    from samgeo.text_sam import LangSAM
+
+    group = options["group"]
+
+    output_raster = options["output"]
+    model_path = options.get("model_path")
+    text_prompt = options.get("text_prompt")
+    text_threshold = float(options.get("text_threshold"))
+    box_threshold = float(options.get("box_threshold"))
+
+    # Set default model path if not provided
+    default_model_path = os.path.join(
+        gs.gisenv()["GISDBASE"], "samgeo_default_model.pth"
+    )
+    if not model_path:
+        model_path = default_model_path
+
+    # Set up paths to access the raster files
+    tmp_dir = gs.tempdir()
+    temp_input_path = os.path.join(tmp_dir, "input.tif")
+
+    temp_output_path = tmp_dir
+    guide_input_path = None
+
+    rasters = gs.read_command("i.group", group=group, flags="lg")
+    raster_list = [raster.strip() for raster in rasters.splitlines()]
+    input_image_np = list(
+        [garray.array(raster, dtype=np.uint8) for raster in raster_list]
+    )
+
+    rgb_array = np.stack(input_image_np, axis=-1)
+
+    if rgb_array.dtype != np.uint8:
+        rgb_array = (
+            (rgb_array - rgb_array.min()) / (rgb_array.max() - rgb_array.min()) * 255
+        )
+        rgb_array = rgb_array.astype(np.uint8)
+
+    np_image = Image.fromarray(rgb_array[:, :, :3])
+
+    # Get device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    if device == "cuda":
+        torch.cuda.empty_cache()
+
+    try:
+        if text_prompt:
+            gs.message("Running LangSAM segmentation...")
+            sam = LangSAM(
+                model_type="sam2-hiera-large",
+            )
+            from torch.amp.autocast_mode import autocast
+
+            with autocast(device_type=device):
+                masks, boxes, phrases, logits = sam.predict(
+                    image=np_image,
+                    text_prompt=text_prompt,
+                    box_threshold=box_threshold,
+                    text_threshold=text_threshold,
+                    return_results=True,
+                )
+        else:
+            gs.message("Running SAMGeo segmentation...")
+            sam = SamGeo(model_type="vit_h", model_path=model_path, device=device)
+            sam.generate(
+                input_image=temp_input_path,
+                output=temp_output_path,
+                guide=guide_input_path,
+            )
+    except Exception as e:
+        gs.message(torch.cuda.memory_summary())
+        gs.fatal(f"Error while running SAMGeo: {e}")
+        return 1
+
+    gs.message("Segmentation complete.")
+
+    write_raster(masks, output_raster)
+    return 0
+
+
+def write_raster(input_np_array, output_raster):
+    gs.message("Importing the segmented raster into GRASS GIS...")
+
+    if len(input_np_array) == 0:
+        gs.fatal("No masks found.")
+
+    # Initialize the merged raster with zeros
+    merged_raster = np.zeros_like(input_np_array[0], dtype=np.int32)
+
+    for idx, band in enumerate(input_np_array):
+        if band.shape != input_np_array[0].shape:
+            gs.fatal("All masks must have the same shape.")
+
+        unique_value = idx + 1  # Start unique values from 1
+
+        # Use NumPy's vectorized operations to assign unique values
+        mask = band != 0  # Create a mask where band is not zero
+        merged_raster[mask] = unique_value
+
+    # Convert the merged raster to a GRASS array
+    mask_raster = garray.array()
+    mask_raster[...] = merged_raster
+
+    # Write the merged raster to the output file
+    mask_raster.write(mapname=output_raster)
+
+
+if __name__ == "__main__":
+    options, flags = gs.parser()
+    sys.exit(main())

src/imagery/i.sam2/requirements.txt

@@ -0,0 +1,4 @@
+Pillow>=10.2.0
+numpy>=1.26.1
+torch>=2.5.1
+segment-geospatial>=0.12.3