Add docstrings

Author: giswqs

docs/examples/satellite.ipynb

@@ -13,7 +13,9 @@
         "[![image](https://img.shields.io/badge/Open-Planetary%20Computer-black?style=flat&logo=microsoft)](https://pccompute.westeurope.cloudapp.azure.com/compute/hub/user-redirect/git-pull?repo=https://github.com/opengeos/segment-geospatial&urlpath=lab/tree/segment-geospatial/docs/examples/satellite.ipynb&branch=main)\n",
         "[![image](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/opengeos/segment-geospatial/blob/main/docs/examples/satellite.ipynb)\n",
         "\n",
-        "This notebook shows how to use segment satellite imagery using the Segment Anything Model (SAM) with a few lines of code. "
+        "This notebook shows how to use segment satellite imagery using the Segment Anything Model (SAM) with a few lines of code. \n",
+        "\n",
+        "Make sure you use GPU runtime for this notebook. For Google Colab, go to `Runtime` -> `Change runtime type` and select `GPU` as the hardware accelerator. "
       ]
     },
     {
@@ -118,18 +120,7 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "out_dir = os.path.join(os.path.expanduser('~'), 'Downloads')\n",
-        "if not os.path.exists(out_dir):\n",
-        "    os.makedirs(out_dir)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "image = os.path.join(out_dir, 'satellite.tif')\n",
+        "image = 'satellite.tif'\n",
         "# image = '/path/to/your/own/image.tif'"
       ]
     },
@@ -174,7 +165,16 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "checkpoint = os.path.join(out_dir, 'sam_vit_h_4b8939.pth')\n",
+        "out_dir = os.path.join(os.path.expanduser('~'), 'Downloads')\n",
+        "checkpoint = os.path.join(out_dir, 'sam_vit_h_4b8939.pth')"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
         "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
         "sam = SamGeo(checkpoint=checkpoint,\n",
         "               model_type='vit_h',\n",
@@ -207,7 +207,9 @@
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "## Polygonize the raster data"
+        "## Polygonize the raster data\n",
+        "\n",
+        "Save the segmentation results as a GeoPackage file."
       ]
     },
     {
@@ -216,10 +218,28 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "vector = os.path.join(out_dir, 'segment.gpkg')\n",
+        "vector = 'segment.gpkg'\n",
         "sam.tiff_to_gpkg(mask, vector, simplify_tolerance=None)"
       ]
     },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "You can also save the segmentation results as any vector data format supported by GeoPandas."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "shapefile = 'segment.shp'\n",
+        "sam.tiff_to_vector(mask, vector)"
+      ]
+    },
     {
       "attachments": {},
       "cell_type": "markdown",

mkdocs.yml

@@ -24,7 +24,7 @@ plugins:
     - mkdocs-jupyter:
           include_source: True
           execute: False
-          #   ignore_h1_titles: True
+          ignore_h1_titles: True
           #   execute_ignore: "*.ipynb"
 
 markdown_extensions:

samgeo/samgeo.py

@@ -32,6 +32,10 @@
 
 
 class SamGeo:
+    """The main class for segmenting geospatial data with the Segment Anything Model (SAM). See 
+        https://github.com/facebookresearch/segment-anything
+    """
+
     def __init__(
         self,
         checkpoint="sam_vit_h_4b8939.pth",
@@ -41,7 +45,17 @@ def __init__(
         mask_multiplier=255,
         sam_kwargs=None,
     ):
-        
+        """Initialize the class. 
+
+        Args:
+            checkpoint (str, optional): The path to the checkpoint. It can be one of the following:
+                sam_vit_h_4b8939.pth, sam_vit_l_0b3195.pth, sam_vit_b_01ec64.pth. Defaults to "sam_vit_h_4b8939.pth".
+            model_type (str, optional): The model type. It can be one of the following: vit_h, vit_l, vit_l. Defaults to 'vit_h'.
+            device (str, optional): The device to use. It can be one of the following: cpu, cuda. Defaults to 'cpu'.
+            erosion_kernel (tuple, optional): The erosion kernel. Defaults to (3, 3).
+            mask_multiplier (int, optional): The mask multiplier. Defaults to 255.
+            sam_kwargs (_type_, optional): The arguments for the SAM model. Defaults to None.
+        """        
         if not os.path.exists(checkpoint):
             print(f'Checkpoint {checkpoint} does not exist.')
             download_checkpoint(output=checkpoint)
@@ -89,6 +103,13 @@ def __call__(self, image):
         return resulting_mask_with_borders * self.mask_multiplier
 
     def generate(self, in_path, out_path, **kwargs):
+        """Segment the input image and save the result to the output path.
+
+        Args:
+            in_path (str): The path to the input image.
+            out_path (str): The path to the output image.
+        """
+
         return tiff_to_tiff(in_path, out_path, self, **kwargs)
 
     def image_to_image(self, image, **kwargs):
@@ -98,7 +119,43 @@ def download_tms_as_tiff(self, source, pt1, pt2, zoom, dist):
         image = draw_tile(source, pt1[0], pt1[1], pt2[0], pt2[1], zoom, dist)
         return image
 
-    def tiff_to_gpkg(self, tiff_path, gpkg_path, simplify_tolerance=None):
+    def tiff_to_gpkg(self, tiff_path, gpkg_path, simplify_tolerance=None, **kwargs):
+        """Convert a tiff file to a gpkg file.
+
+        Args:
+            tiff_path (str): The path to the tiff file.
+            gpkg_path (str): The path to the gpkg file.
+            simplify_tolerance (_type_, optional): The simplify tolerance. Defaults to None.
+        """
+    
+        with rasterio.open(tiff_path) as src:
+            band = src.read()
+
+            mask = band != 0
+            shapes = features.shapes(band, mask=mask, transform=src.transform)
+
+        fc = [
+            {"geometry": shapely.geometry.shape(shape), "properties": {"value": value}}
+            for shape, value in shapes
+        ]
+        if simplify_tolerance is not None:
+            for i in fc:
+                i["geometry"] = i["geometry"].simplify(tolerance=simplify_tolerance)
+
+        gdf = gpd.GeoDataFrame.from_features(fc)
+        gdf.set_crs(epsg=src.crs.to_epsg(), inplace=True)
+        gdf.to_file(gpkg_path, driver='GPKG', **kwargs)
+
+
+    def tiff_to_vector(self, tiff_path, output, simplify_tolerance=None, **kwargs):
+        """Convert a tiff file to a gpkg file.
+
+        Args:
+            tiff_path (str): The path to the tiff file.
+            output (str): The path to the vector file.
+            simplify_tolerance (_type_, optional): The simplify tolerance. Defaults to None.
+        """
+    
         with rasterio.open(tiff_path) as src:
             band = src.read()
 
@@ -115,4 +172,4 @@ def tiff_to_gpkg(self, tiff_path, gpkg_path, simplify_tolerance=None):
 
         gdf = gpd.GeoDataFrame.from_features(fc)
         gdf.set_crs(epsg=src.crs.to_epsg(), inplace=True)
-        gdf.to_file(gpkg_path, driver='GPKG')
+        gdf.to_file(output, **kwargs)