Use defaultdict

ome · Jun 12, 2020 · 2075756 · 2075756 · mtbc · Jun 12, 2020
1 parent 7f2d8f1
commit 2075756
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/src/omero_rois/library.py b/src/omero_rois/library.py
@@ -18,6 +18,7 @@
 # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 
 import numpy as np
+from collections import defaultdict
 from omero.gateway import ColorHolder
 from omero.model import MaskI
 from omero.rtypes import (
@@ -154,7 +155,7 @@ def masks_from_3d_label_image(
            ({} if no labels found)
 
     """
-    masks = {}
+    masks = defaultdict(list)
     for i, plane in enumerate(planes):
         if z_stack:
             plane_masks = masks_from_label_image(plane, rgba, i, c, t,
@@ -163,8 +164,6 @@ def masks_from_3d_label_image(
             plane_masks = masks_from_label_image(plane, rgba, z, c, i,
                                                  text, False)
         for label, mask in enumerate(plane_masks):
-            if label not in masks:
-                masks[label] = []
             if mask.getBytes().any():
                 masks[label].append(mask)
     return masks