Fix save masks bug (#332)

* Fix save masks bug * Fix save masks bug for SAM 1
opengeos · Oct 5, 2024 · 837f5a5 · 837f5a5
1 parent dab8b84
commit 837f5a5
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/samgeo/samgeo.py b/samgeo/samgeo.py
@@ -265,8 +265,8 @@ def save_masks(
 
         # Generate a mask of objects with unique values
         if unique:
-            # Sort the masks by area in ascending order
-            sorted_masks = sorted(masks, key=(lambda x: x["area"]), reverse=False)
+            # Sort the masks by area in descending order
+            sorted_masks = sorted(masks, key=(lambda x: x["area"]), reverse=True)
 
             # Create an output image with the same size as the input image
             objects = np.zeros(
@@ -276,9 +276,10 @@ def save_masks(
                 )
             )
             # Assign a unique value to each object
+            count = len(sorted_masks)
             for index, ann in enumerate(sorted_masks):
                 m = ann["segmentation"]
-                objects[m] = index + 1
+                objects[m] = count - index
 
         # Generate a binary mask
         else:

diff --git a/samgeo/samgeo2.py b/samgeo/samgeo2.py
@@ -293,8 +293,8 @@ def save_masks(
 
         # Generate a mask of objects with unique values
         if unique:
-            # Sort the masks by area in ascending order
-            sorted_masks = sorted(masks, key=(lambda x: x["area"]), reverse=False)
+            # Sort the masks by area in descending order
+            sorted_masks = sorted(masks, key=(lambda x: x["area"]), reverse=True)
 
             # Create an output image with the same size as the input image
             objects = np.zeros(
@@ -304,9 +304,10 @@ def save_masks(
                 )
             )
             # Assign a unique value to each object
+            count = len(sorted_masks)
             for index, ann in enumerate(sorted_masks):
                 m = ann["segmentation"]
-                objects[m] = index + 1
+                objects[m] = count - index
 
         # Generate a binary mask
         else: