From 9e084de37201d347471d12ad1b53fc75d002bc88 Mon Sep 17 00:00:00 2001 From: Corey White Date: Fri, 15 Nov 2024 22:06:28 -0500 Subject: [PATCH 01/10] i.sam2: SamGeo2 model --- src/imagery/i.sam2/Makefile | 7 + src/imagery/i.sam2/i.sam2.html | 16 ++ src/imagery/i.sam2/i.sam2.py | 227 ++++++++++++++++++++ src/imagery/i.sam2/requirements.txt | 1 + src/imagery/i.sam2/testsuite/test_i_sam2.py | 0 5 files changed, 251 insertions(+) create mode 100644 src/imagery/i.sam2/Makefile create mode 100644 src/imagery/i.sam2/i.sam2.html create mode 100644 src/imagery/i.sam2/i.sam2.py create mode 100644 src/imagery/i.sam2/requirements.txt create mode 100644 src/imagery/i.sam2/testsuite/test_i_sam2.py diff --git a/src/imagery/i.sam2/Makefile b/src/imagery/i.sam2/Makefile new file mode 100644 index 0000000000..3c16d5bb2b --- /dev/null +++ b/src/imagery/i.sam2/Makefile @@ -0,0 +1,7 @@ +MODULE_TOPDIR = ../.. + +PGM = i.sam2 + +include $(MODULE_TOPDIR)/include/Make/Script.make + +default: script diff --git a/src/imagery/i.sam2/i.sam2.html b/src/imagery/i.sam2/i.sam2.html new file mode 100644 index 0000000000..1fa3560437 --- /dev/null +++ b/src/imagery/i.sam2/i.sam2.html @@ -0,0 +1,16 @@ +

DESCRIPTION

+ +i.sam2 allows users to segment orthoimagery based on text prompts using SamGeo2. + +

EXAMPLES

+ +Segment orthoimagery using SamGeo2: + +
+    i.sam2 group=rgb_255 output=tree_mask text_prompt="trees"
+
+
+ + +

AUTHOR

+Corey T. White (NCSU GeoForAll Lab & OpenPlains Inc.) diff --git a/src/imagery/i.sam2/i.sam2.py b/src/imagery/i.sam2/i.sam2.py new file mode 100644 index 0000000000..6db2cb73e0 --- /dev/null +++ b/src/imagery/i.sam2/i.sam2.py @@ -0,0 +1,227 @@ +#!/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 for segmentation in GRASS GIS. +# % keyword: raster +# % keyword: segmentation +# % 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: Optional 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 + +# %flag +# % key: u +# % description: Update the default SAMGeo model +# %end + +import os +import sys +import grass.script as gs +import torch +import requests +import numpy as np +from PIL import Image +from PIL.Image import Image as PILImage +from grass.pygrass import raster as r +from grass.script import array as garray + + +def update_model(default_model_path): + gs.message("Updating the SAMGeo model...") + model_url = "https://example.com/samgeo/default_model.pth" # Replace with the actual model URL + response = requests.get(model_url) + if response.status_code == 200: + with open(default_model_path, "wb") as model_file: + model_file.write(response.content) + gs.message("Model updated successfully.") + else: + gs.fatal( + "Failed to update the model. Please check the URL or your internet connection." + ) + + +def main(): + from samgeo import SamGeo + from samgeo.text_sam import LangSAM + + gisenv = gs.gisenv() + 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")) + update_flag = flags["u"] + + # 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 + + # Update model if flag is set + if update_flag: + update_model(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").strip().split("\n") + input_image_np = list([garray.array(raster, dtype=np.uint8) for raster in rasters]) + + 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) + + image_size = rgb_array.shape + + np_image = Image.fromarray(rgb_array[:, :, :3]) + image_np = np.array(np_image) + + gs.message( + f"np_image type: {type(np_image)}, {isinstance(np_image, PILImage)}, image_np size: {image_np.size}" + ) + + # 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="vit_h", + model_type="sam2-hiera-large", + # checkpoint=model_path + ) + from torch.amp.autocast_mode import autocast + + with autocast(device_type=device): + masks, boxes, phrases, logits = sam.predict( + image=np_image, + # output=temp_output_path, + 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}") + # Check GPU memory usage + # allocated = torch.cuda.memory_allocated() + # reserved = torch.cuda.memory_reserved() + # free = torch.cuda.get_device_properties(0).total_memory - reserved + + # gs.message(f"Allocated Memory: {allocated / 1024**2:.2f} MB") + # gs.message(f"Reserved Memory: {reserved / 1024**2:.2f} MB") + # gs.message(f"Free Memory: {free / 1024**2:.2f} MB") + 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 input_np_array.shape[0] == 1: + gs.message("Writing single-band raster...") + mask_raster = garray.array() + for y in range(mask_raster.shape[0]): + for x in range(mask_raster.shape[1]): + mask_raster[y][x] = input_np_array[0][y][x] + mask_raster.write(mapname=output_raster) + else: + gs.message("Writing multi-band raster...") + # TODO: Merge the bands into a single raster + for idx, band in enumerate(input_np_array): + if band.shape != input_np_array[0].shape: + gs.fatal("All bands must have the same shape.") + mask_raster = garray.array() + for y in range(mask_raster.shape[0]): + for x in range(mask_raster.shape[1]): + mask_raster[y][x] = band[y][x] + + mask_raster.write(mapname=f"{output_raster}.{idx}") + + +if __name__ == "__main__": + options, flags = gs.parser() + sys.exit(main()) diff --git a/src/imagery/i.sam2/requirements.txt b/src/imagery/i.sam2/requirements.txt new file mode 100644 index 0000000000..ae46d19e2a --- /dev/null +++ b/src/imagery/i.sam2/requirements.txt @@ -0,0 +1 @@ +segment-geospatial diff --git a/src/imagery/i.sam2/testsuite/test_i_sam2.py b/src/imagery/i.sam2/testsuite/test_i_sam2.py new file mode 100644 index 0000000000..e69de29bb2 From 8f35e5d8b39d76d45e1e03d5d87fbdfb6047dbc3 Mon Sep 17 00:00:00 2001 From: Corey White Date: Fri, 15 Nov 2024 22:16:56 -0500 Subject: [PATCH 02/10] Updated docs --- src/imagery/i.sam2/i.sam2.html | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/imagery/i.sam2/i.sam2.html b/src/imagery/i.sam2/i.sam2.html index 1fa3560437..8bce8d56c6 100644 --- a/src/imagery/i.sam2/i.sam2.html +++ b/src/imagery/i.sam2/i.sam2.html @@ -1,6 +1,6 @@

DESCRIPTION

-i.sam2 allows users to segment orthoimagery based on text prompts using SamGeo2. +i.sam2 allows users to segment orthoimagery based on text prompts using SamGeo.

EXAMPLES

@@ -11,6 +11,12 @@

EXAMPLES

+

REFERENCES

+ +

AUTHOR

Corey T. White (NCSU GeoForAll Lab & OpenPlains Inc.) From 1706d70d0ac82847eb32f8285467672aee0fe71f Mon Sep 17 00:00:00 2001 From: Corey White Date: Sat, 16 Nov 2024 12:02:58 -0500 Subject: [PATCH 03/10] Update i.sam2.py keywords Co-authored-by: Markus Neteler --- src/imagery/i.sam2/i.sam2.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/imagery/i.sam2/i.sam2.py b/src/imagery/i.sam2/i.sam2.py index 6db2cb73e0..276f599634 100644 --- a/src/imagery/i.sam2/i.sam2.py +++ b/src/imagery/i.sam2/i.sam2.py @@ -14,8 +14,9 @@ # %module # % description: Integrates SAMGeo model for segmentation in GRASS GIS. -# % keyword: raster +# % keyword: imagery # % keyword: segmentation +# % keyword: object recognition # % keyword: deep learning # %end From 462be0ef372e07d7b8341e72d0c946f044b866c1 Mon Sep 17 00:00:00 2001 From: Corey White Date: Sat, 16 Nov 2024 12:07:54 -0500 Subject: [PATCH 04/10] Update i.sam2.py: updated description Co-authored-by: Markus Neteler --- src/imagery/i.sam2/i.sam2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/imagery/i.sam2/i.sam2.py b/src/imagery/i.sam2/i.sam2.py index 276f599634..1bd261f2a7 100644 --- a/src/imagery/i.sam2/i.sam2.py +++ b/src/imagery/i.sam2/i.sam2.py @@ -13,7 +13,7 @@ ############################################################################# # %module -# % description: Integrates SAMGeo model for segmentation in GRASS GIS. +# % description: Integrates SAMGeo model with text prompt for segmentation in GRASS GIS. # % keyword: imagery # % keyword: segmentation # % keyword: object recognition From 0fb59f2e862284e30bf5964b390b50e3496975bc Mon Sep 17 00:00:00 2001 From: Corey White Date: Sun, 17 Nov 2024 20:26:54 -0500 Subject: [PATCH 05/10] Code clean up, improved docs --- src/imagery/i.sam2/i.sam2.html | 2 + src/imagery/i.sam2/i.sam2.py | 87 +++++++++------------------------ src/imagery/i.sam2/trees.png | Bin 0 -> 475023 bytes 3 files changed, 25 insertions(+), 64 deletions(-) create mode 100644 src/imagery/i.sam2/trees.png diff --git a/src/imagery/i.sam2/i.sam2.html b/src/imagery/i.sam2/i.sam2.html index 8bce8d56c6..48574e9fd6 100644 --- a/src/imagery/i.sam2/i.sam2.html +++ b/src/imagery/i.sam2/i.sam2.html @@ -11,6 +11,8 @@

EXAMPLES

+ +

REFERENCES