More robust handling of GCPs geojson #731

rioxarray/rioxarray.py

@@ -1248,18 +1248,20 @@ def get_gcps(self) -> Optional[list[GroundControlPoint]]:
         except (KeyError, AttributeError):
             return None
 
-        gcps = [
-            GroundControlPoint(
-                x=gcp["geometry"]["coordinates"][0],
-                y=gcp["geometry"]["coordinates"][1],
-                z=gcp["geometry"]["coordinates"][2],
+        def parse_gcp(gcp) -> GroundControlPoint:
+            x, y, *z = gcp["geometry"]["coordinates"]
+            z = z[0] if len(z) > 0 else None
+            return GroundControlPoint(
+                x=x,
+                y=y,
+                z=z,
                 row=gcp["properties"]["row"],
                 col=gcp["properties"]["col"],
                 id=gcp["properties"]["id"],
                 info=gcp["properties"]["info"],
             )
-            for gcp in geojson_gcps["features"]
-        ]
+
+        gcps = [parse_gcp(gcp) for gcp in geojson_gcps["features"]]
         return gcps
 
 
@@ -1277,6 +1279,12 @@ def _convert_gcps_to_geojson(
     -------
     A FeatureCollection dict.
     """
+
+    def gcp_coordinates(gcp):
+        if gcp.z is None:
+            return [gcp.x, gcp.y]
+        return [gcp.x, gcp.y, gcp.z]
+
     features = [
         {
             "type": "Feature",
@@ -1286,7 +1294,7 @@ def _convert_gcps_to_geojson(
                 "row": gcp.row,
                 "col": gcp.col,
             },
-            "geometry": {"type": "Point", "coordinates": [gcp.x, gcp.y, gcp.z]},
+            "geometry": {"type": "Point", "coordinates": gcp_coordinates(gcp)},
         }
         for gcp in gcps
     ]

test/integration/test_integration_rioxarray.py

@@ -2895,19 +2895,20 @@ def test_interpolate_na_missing_nodata():
         test_da.to_dataset().rio.interpolate_na()
 
 
-def test_rio_write_gcps():
+@pytest.mark.parametrize("with_z", [True, False])
+def test_rio_write_gcps(with_z):
     """
     Test setting gcps in dataarray.
     """
-    gdal_gcps, gcp_crs = _create_gdal_gcps()
+    gdal_gcps, gcp_crs = _create_gdal_gcps(with_z)
 
     darr = xarray.DataArray(1)
     darr.rio.write_gcps(gdal_gcps, gcp_crs, inplace=True)
 
     _check_rio_gcps(darr, gdal_gcps, gcp_crs)
 
 
-def _create_gdal_gcps():
+def _create_gdal_gcps(with_z=True):
     src_gcps = [
         GroundControlPoint(
             row=0, col=0, x=0.0, y=0.0, z=12.0, id="1", info="the first gcp"
@@ -2922,6 +2923,9 @@ def _create_gdal_gcps():
             row=800, col=0, x=0.0, y=10.0, z=5.5, id="4", info="the fourth gcp"
         ),
     ]
+    if with_z is False:
+        for gcp in src_gcps:
+            gcp.z = None
     crs = CRS.from_epsg(4326)
     gdal_gcps = (src_gcps, crs)
     return gdal_gcps
@@ -2943,14 +2947,18 @@ def _check_rio_gcps(darr, src_gcps, crs):
         assert feature["properties"]["row"] == gcp.row
         assert feature["properties"]["col"] == gcp.col
         assert feature["geometry"]["type"] == "Point"
-        assert feature["geometry"]["coordinates"] == [gcp.x, gcp.y, gcp.z]
+        if gcp.z is not None:
+            assert feature["geometry"]["coordinates"] == [gcp.x, gcp.y, gcp.z]
+        else:
+            assert feature["geometry"]["coordinates"] == [gcp.x, gcp.y]
 
 
-def test_rio_get_gcps():
+@pytest.mark.parametrize("with_z", [True, False])
+def test_rio_get_gcps(with_z):
     """
     Test setting gcps in dataarray.
     """
-    gdal_gcps, gdal_crs = _create_gdal_gcps()
+    gdal_gcps, gdal_crs = _create_gdal_gcps(with_z)
 
     darr = xarray.DataArray(1)
     darr.rio.write_gcps(gdal_gcps, gdal_crs, inplace=True)