bugfix: Check if prediction output is list or np.array (#49)

vanvalenlab · Aug 3, 2019 · 2ad050a · 2ad050a
1 parent a741332
commit 2ad050a
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/redis_consumer/consumers.py b/redis_consumer/consumers.py
@@ -602,7 +602,10 @@ def _consume(self, redis_hash):
             post_funcs = hvals.get('postprocess_function', '').split(',')
             # image[:-1] is targeted at a two semantic head panoptic model
             # TODO This may need to be modified and generalized in the future
-            image = self.postprocess(image[:-1], post_funcs, True)
+            if isinstance(image, list):
+                image = self.postprocess(image[:-1], post_funcs, True)
+            else:
+                image = self.postprocess(image, post_funcs, True)
 
             # Save the post-processed results to a file
             _ = timeit.default_timer()