[FIX] delete unwanted classes (1, 1.)

NINAnor · Oct 16, 2024 · 0ee39bf · 0ee39bf
1 parent bddcfa9
commit 0ee39bf
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/dataset/rasterize_masks.py b/src/dataset/rasterize_masks.py
@@ -36,6 +36,8 @@ def rasterize_masks(
     gdf = gpd.read_file(geojson_path)
 
     # Factorize the text labels to get numeric labels
+    gdf = gdf[gdf["labelTekst"] != "1"]
+    gdf = gdf[gdf["labelTekst"] != "1."]
     gdf["label_encoded"], _ = pd.factorize(gdf[label_column])
     gdf["label_encoded"] = gdf["label_encoded"]
 
@@ -91,9 +93,7 @@ def rasterize_masks(
 
         # Save the mask as .tif to ensure saving of negative values
         mask_filename = f"mask_{image_id}.tif"
-        mask_img = Image.fromarray(
-            mask.astype(np.int16)
-        )  # Ensure it's int16 for -1 values
+        mask_img = Image.fromarray(mask.astype(np.int16))
         mask_img.save(os.path.join(output_mask_dir, mask_filename))
 
         print(f"Saved mask: {mask_filename}")