Export binary masks, too

extensions/chaeo/workflows.py

@@ -239,8 +239,8 @@ def transfer_ecotaxa_labels_to_patch_stacks(
     df_labels.to_csv(Path(where_output) / 'labels_key.csv', index=False)
 
     # export patches as a single z-stack
-    mask_zstack = np.zeros((*patch_size, 1, len(df_tr)), dtype='uint8')
-    patch_zstack = np.zeros((*patch_size, 1, len(df_tr)), dtype='uint8')
+    zstack_keys = ['mask', 'label', 'raw']
+    zstacks = {k: np.zeros((*patch_size, 1, len(df_tr)), dtype='uint8') for k in zstack_keys}
     for fi, pl in enumerate(df_tr.itertuples(name='PatchFile')):
         fn = pl._asdict()['patch_filename']
         ac = pl._asdict()['annotation_class_id']
@@ -248,11 +248,14 @@ def transfer_ecotaxa_labels_to_patch_stacks(
         assert acc_bm.hw == patch_size, f'Unexpected patch size {patch_size}'
         assert acc_bm.chroma == 1
         assert acc_bm.nz == 1
-        mask_zstack[:, :, 0, fi] = (acc_bm.data[:, :, 0, 0] == 255) * ac
+        mask = acc_bm.data[:, :, 0, 0]
+        zstacks['mask'][:, :, 0, fi] = mask
+        zstacks['label'][:, :, 0, fi] = (mask == 255) * ac
 
         acc_pa = generate_file_accessor(Path(where_patches) / fn)
-        patch_zstack[:, :, :, fi] = acc_pa.data[:, :, :, 0]
-    write_accessor_data_to_file(Path(where_output) / 'zstack_object_label.tif', InMemoryDataAccessor(mask_zstack))
-    write_accessor_data_to_file(Path(where_output) / 'zstack_patches.tif', InMemoryDataAccessor(patch_zstack))
+        zstacks['raw'][:, :, :, fi] = acc_pa.data[:, :, :, 0]
+
+    for k in zstack_keys:
+        write_accessor_data_to_file(Path(where_output) / f'zstack_{k}.tif', InMemoryDataAccessor(zstacks[k]))