Export test zstacks, too

extensions/chaeo/workflows.py

@@ -238,30 +238,31 @@ def transfer_ecotaxa_labels_to_patch_stacks(
         df_labels.loc[df_labels[col].isna(), col] = 0
     df_labels.to_csv(Path(where_output) / 'labels_key.csv', index=False)
 
-    # export patches as a single z-stack
-    zstack_keys = ['mask', 'label', 'raw']
-    zstacks = {k: np.zeros((*patch_size, 1, len(df_tr)), dtype='uint8') for k in zstack_keys}
-    stack_meta = []
-    for fi, pl in enumerate(df_tr.itertuples(name='PatchFile')):
-        fn = pl._asdict()['patch_filename']
-        ac = pl._asdict()['annotation_class']
-        aci = pl._asdict()['annotation_class_id']
+    # export patches as z-stacks
+    for (dfk, dfv) in {'train': df_tr, 'test': df_te}.items():
+        zstack_keys = ['mask', 'label', 'raw']
+        zstacks = {f'{dfk}_{zsk}': np.zeros((*patch_size, 1, len(df_tr)), dtype='uint8') for zsk in zstack_keys}
+        stack_meta = []
+        for fi, pl in enumerate(dfv.itertuples(name='PatchFile')):
+            fn = pl._asdict()['patch_filename']
+            ac = pl._asdict()['annotation_class']
+            aci = pl._asdict()['annotation_class_id']
 
-        stack_meta.append({'zi': fi, 'patch_filename': fn, 'annotation_class': ac, 'annotation_class_id': aci})
-        acc_bm = generate_file_accessor(Path(where_masks) / fn)
-        assert acc_bm.hw == patch_size, f'Unexpected patch size {patch_size}'
-        assert acc_bm.chroma == 1
-        assert acc_bm.nz == 1
-        mask = acc_bm.data[:, :, 0, 0]
-        zstacks['mask'][:, :, 0, fi] = mask
-        zstacks['label'][:, :, 0, fi] = (mask == 255) * aci
+            stack_meta.append({'zi': fi, 'patch_filename': fn, 'annotation_class': ac, 'annotation_class_id': aci})
+            acc_bm = generate_file_accessor(Path(where_masks) / fn)
+            assert acc_bm.hw == patch_size, f'Unexpected patch size {patch_size}'
+            assert acc_bm.chroma == 1
+            assert acc_bm.nz == 1
+            mask = acc_bm.data[:, :, 0, 0]
+            zstacks[dfk + '_mask'][:, :, 0, fi] = mask
+            zstacks[dfk + '_label'][:, :, 0, fi] = (mask == 255) * aci
 
-        acc_pa = generate_file_accessor(Path(where_patches) / fn)
-        zstacks['raw'][:, :, :, fi] = acc_pa.data[:, :, :, 0]
+            acc_pa = generate_file_accessor(Path(where_patches) / fn)
+            zstacks[dfk + '_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]))
+        for k in zstacks.keys():
+            write_accessor_data_to_file(Path(where_output) / f'zstack_{k}.tif', InMemoryDataAccessor(zstacks[k]))
 
-    pd.DataFrame(stack_meta).to_csv(Path(where_output) / 'training_stack.csv', index=False)
+        pd.DataFrame(stack_meta).to_csv(Path(where_output) / f'{dfk}_stack.csv', index=False)