Merge branch 'rdm-point-extract' of https://github.com/WorldCereal/wo…

…rldcereal-classification into rdm-point-extract
WorldCereal · Nov 26, 2024 · 6b281d3 · 6b281d3
2 parents 7fcd5fc + be9b46b
commit 6b281d3
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/worldcereal/rdm_api/rdm_interaction.py b/src/worldcereal/rdm_api/rdm_interaction.py
@@ -37,6 +37,7 @@ class RdmInteraction:
         "quality_score_ct",
         "extract",
         "h3_l3_cell",
+        "collection_id",
     ]
 
     # RDM API Endpoint
@@ -456,6 +457,8 @@ def download_samples(
         # Convert the WKB geometry to a Shapely geometry
         df["geometry"] = df["wkb_geometry"].apply(lambda x: wkb.loads(bytes(x)))
         df.drop(columns=["wkb_geometry"], inplace=True)
+        # Make sure df contains only requested columns
+        df = df[columns]
 
         # Convert the DataFrame to a GeoDataFrame
         gdf = gpd.GeoDataFrame(df, geometry="geometry", crs="EPSG:4326")