Merge fedeb3a into c4883ae

vanvalenlab · Aug 29, 2019 · 27096c3 · 27096c3
2 parents c4883ae + fedeb3a
commit 27096c3
Show file tree

Hide file tree

Showing 4 changed files with 353 additions and 77 deletions.
diff --git a/redis_consumer/consumers.py b/redis_consumer/consumers.py
@@ -1040,7 +1040,7 @@ def _load_data(self, hvalues, subdir, fname):
 
         # remove the last dimensions added by `get_image`
 
-        tiff_stack = np.squeeze(raw, -1)
+        tiff_stack = np.squeeze(raw, -1)  # TODO: required? check the ndim?
         if len(tiff_stack.shape) != 3:
             raise ValueError("This tiff file has shape {}, which is not 3 "
                              "dimensions. Tracking can only be done on images "
@@ -1115,7 +1115,7 @@ def _load_data(self, hvalues, subdir, fname):
                     self.logger.debug('Hash %s has status %s',
                                       segment_hash, status)
 
-                    if status == 'failed':
+                    if status == self.failed_status:
                         reason = self.redis.hget(segment_hash, 'reason')
                         raise RuntimeError(
                             'Tracking failed during segmentation on frame {}.'
@@ -1186,7 +1186,8 @@ def _consume(self, redis_hash):
             save_name = os.path.join(
                 tempdir, hvalues.get('original_name', fname)) + '.trk'
 
-            tracker.dump(save_name, file_format='.trk')
+            # Post-process and save the output file
+            tracker.postprocess(save_name)
             output_file_name, output_url = self.storage.upload(save_name)
 
             self.update_key(redis_hash, {

diff --git a/redis_consumer/consumers_test.py b/redis_consumer/consumers_test.py
@@ -174,7 +174,10 @@ class DummyTracker(object):
     def _track_cells(self):
         return None
 
-    def dump(*_, **__):
+    def dump(self, *_, **__):
+        return None
+
+    def postprocess(self, *_, **__):
         return None