From 14d76dfa58cc336021be87af53eb1d0439aca257 Mon Sep 17 00:00:00 2001
From: Christopher Rhodes <christopher.rhodes@embl.de>
Date: Thu, 21 Dec 2023 15:13:25 +0100
Subject: [PATCH] Export DataFrame from workflow

---
 extensions/chaeo/params.py            |  1 +
 extensions/chaeo/tests/test_zstack.py |  3 ++-
 extensions/chaeo/zmask.py             | 13 ++++---------
 3 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/extensions/chaeo/params.py b/extensions/chaeo/params.py
index a1f6aa8c..e9fc2461 100644
--- a/extensions/chaeo/params.py
+++ b/extensions/chaeo/params.py
@@ -40,5 +40,6 @@ class RoiSetExportParams(BaseModel):
     patch_masks: bool = False
     annotated_zstacks: Union[AnnotatedZStackParams, None] = None
     object_classes: bool = False
+    dataframe: bool = False
 
 
diff --git a/extensions/chaeo/tests/test_zstack.py b/extensions/chaeo/tests/test_zstack.py
index 1779286b..45256258 100644
--- a/extensions/chaeo/tests/test_zstack.py
+++ b/extensions/chaeo/tests/test_zstack.py
@@ -232,7 +232,8 @@ class TestZStackDerivedDataProducts(unittest.TestCase):
             },
             'patch_masks': True,
             'annotated_zstacks': {},
-            'object_classes': True
+            'object_classes': True,
+            'dataframe': True,
         })
 
         infer_object_map_from_zstack(
diff --git a/extensions/chaeo/zmask.py b/extensions/chaeo/zmask.py
index 053f1d44..d1e9301e 100644
--- a/extensions/chaeo/zmask.py
+++ b/extensions/chaeo/zmask.py
@@ -81,17 +81,9 @@ class RoiSet(object):
 
         self.object_class_map = InMemoryDataAccessor(om)
 
-    # TODO: test
-    def get_object_mask_by_id(self, obj_id):
-        return self.object_id_labels == obj_id
-
     def get_object_mask_by_class(self, class_id):
         return self.object_id_labels == class_id
 
-    # TODO: implement
-    def get_object_patch_by_id(self, obj_id):
-        pass
-
     def run_exports(self, where, channel, prefix, params: RoiSetExportParams):
         if not self.count:
             return
@@ -129,7 +121,10 @@ class RoiSet(object):
                 write_accessor_data_to_file(subdir / (pr + '.tif'), annotated)
             if k == 'object_classes':
                 write_accessor_data_to_file(subdir / (pr + '.tif'), self.object_class_map)
-
+            if k == 'dataframe':
+                dfpa = subdir / (pr + '.csv')
+                dfpa.parent.mkdir(parents=True, exist_ok=True)
+                self.df.to_csv(dfpa, index=False)
 
 
 def build_zmask_from_object_mask(
-- 
GitLab