Add data_column_name argument to Raster.polygonize (#580)

geoutils/raster/raster.py

@@ -4132,15 +4132,19 @@ def from_pointcloud_regular(
         return raster_arr
 
     def polygonize(
-        self, target_values: Number | tuple[Number, Number] | list[Number] | NDArrayNum | Literal["all"] = "all"
+        self,
+        target_values: Number | tuple[Number, Number] | list[Number] | NDArrayNum | Literal["all"] = "all",
+        data_column_name: str = "id",
     ) -> Vector:
         """
         Polygonize the raster into a vector.
 
         :param target_values: Value or range of values of the raster from which to
-          create geometries (defaults to 'all', for which all unique pixel values of the raster are used).
+          create geometries (defaults to "all", for which all unique pixel values of the raster are used).
+        :param data_column_name: Data column name to be associated with target values in the output vector
+            (defaults to "id").
 
-        :returns: Vector containing the polygonized geometries.
+        :returns: Vector containing the polygonized geometries associated to target values.
         """
 
         # Mask a unique value set by a number
@@ -4191,7 +4195,7 @@ def polygonize(
         )
 
         gdf = gpd.GeoDataFrame.from_features(list(results))
-        gdf.insert(0, "New_ID", range(0, 0 + len(gdf)))
+        gdf.insert(0, data_column_name, range(0, 0 + len(gdf)))
         gdf = gdf.set_geometry(col="geometry")
         gdf = gdf.set_crs(self.crs)
 
@@ -4562,7 +4566,9 @@ def crop(
             return super().crop(crop_geom=crop_geom, mode=mode, inplace=inplace)
 
     def polygonize(
-        self, target_values: Number | tuple[Number, Number] | list[Number] | NDArrayNum | Literal["all"] = 1
+        self,
+        target_values: Number | tuple[Number, Number] | list[Number] | NDArrayNum | Literal["all"] = 1,
+        data_column_name: str = "id",
     ) -> Vector:
         # If target values is passed but does not correspond to 0 or 1, raise a warning
         if not isinstance(target_values, (int, np.integer, float, np.floating)) or target_values not in [0, 1]:

tests/test_raster/test_raster.py

@@ -3116,6 +3116,12 @@ def test_polygonize(self, example: str) -> None:
         assert isinstance(polygonized, gu.Vector)
         assert polygonized.crs == img.crs
 
+        # Check default name of data column, and that defining a custom name works the same
+        assert "id" in polygonized.ds.columns
+        polygonized2 = img.polygonize(target_values=value, data_column_name="myname")
+        assert "myname" in polygonized2.ds.columns
+        assert np.array_equal(polygonized2.ds["myname"].values, polygonized.ds["id"].values)
+
         # -- Test 2: data types --
 
         # Check that polygonize works as expected for any input dtype (e.g. float64 being not supported by GeoPandas)