Refactor TrackingConsumer to CalibanConsumer. (#180)

* Refactor TrackingConsumer to CalibanConsumer.

* Update docs.

* Use CALIBAN_MODEL instead of TRACKING_MODEL.

* Remove DRIFT_CORRECT_ENABLED option.

* Remove scikit-image from dependencies.

* Update toolbox to latest patch.

docs/source/conf.py

@@ -200,8 +200,7 @@
 
 
 # -- Extension configuration -------------------------------------------------
-autodoc_mock_imports = ["skimage",
-                        "cv2",
+autodoc_mock_imports = ["cv2",
                         "keras_retinanet",
                         "keras_preprocessing",
                         "deepcell_tracking",

docs/source/redis_consumer.consumers.rst

@@ -10,19 +10,28 @@ redis_consumer.consumers.base_consumer module
     :show-inheritance:
     :private-members:
 
-redis_consumer.consumers.image_consumer module
-----------------------------------------------
+redis_consumer.consumers.segmentation_consumer module
+-----------------------------------------------------
 
-.. automodule:: redis_consumer.consumers.image_consumer
+.. automodule:: redis_consumer.consumers.segmentation_consumer
     :members:
     :undoc-members:
     :show-inheritance:
     :private-members:
 
-redis_consumer.consumers.tracking_consumer module
+redis_consumer.consumers.mesmer_consumer module
+-----------------------------------------------
+
+.. automodule:: redis_consumer.consumers.mesmer_consumer
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :private-members:
+
+redis_consumer.consumers.caliban_consumer module
 -------------------------------------------------
 
-.. automodule:: redis_consumer.consumers.tracking_consumer
+.. automodule:: redis_consumer.consumers.caliban_consumer
     :members:
     :undoc-members:
     :show-inheritance:

redis_consumer/consumers/__init__.py

@@ -34,7 +34,7 @@
 
 # Custom Workflow consumers
 from redis_consumer.consumers.segmentation_consumer import SegmentationConsumer
-from redis_consumer.consumers.tracking_consumer import TrackingConsumer
+from redis_consumer.consumers.caliban_consumer import CalibanConsumer
 from redis_consumer.consumers.mesmer_consumer import MesmerConsumer
 # TODO: Import future custom Consumer classes.
 
@@ -43,15 +43,17 @@
     'image': SegmentationConsumer,  # deprecated, use "segmentation" instead.
     'segmentation': SegmentationConsumer,
     'zip': ZipFileConsumer,
-    'tracking': TrackingConsumer,
+    'tracking': CalibanConsumer,  # deprecated, use "caliban" instead.
     'multiplex': MesmerConsumer,  # deprecated, use "mesmer" instead.
     'mesmer': MesmerConsumer,
+    'caliban': CalibanConsumer,
     # TODO: Add future custom Consumer classes here.
 }
 
 
-# Backwards compatibility for MultiplexConsumer
+# Backwards compatibility aliases
 MultiplexConsumer = MesmerConsumer
+TrackingConsumer = CalibanConsumer
 
 
 del absolute_import

redis_consumer/consumers/tracking_consumer.py → redis_consumer/consumers/caliban_consumer.py

@@ -23,7 +23,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-"""TrackingConsumer class for consuming cell tracking jobs."""
+"""CalibanConsumer class for consuming cell tracking jobs."""
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
@@ -38,16 +38,14 @@
 import numpy as np
 import tifffile
 
-from deepcell_toolbox.processing import correct_drift
-
 from deepcell.applications import CellTracking
 
 from redis_consumer.consumers import TensorFlowServingConsumer
 from redis_consumer import utils
 from redis_consumer import settings
 
 
-class TrackingConsumer(TensorFlowServingConsumer):
+class CalibanConsumer(TensorFlowServingConsumer):
     """Consumes some unspecified file format, tracks the images,
        and uploads the results
     """
@@ -225,19 +223,12 @@ def _consume(self, redis_hash):
         self.logger.debug('X shape: %s', data['X'].shape)
         self.logger.debug('y shape: %s', data['y'].shape)
 
-        # Correct for drift if enabled
-        if settings.DRIFT_CORRECT_ENABLED:
-            t = timeit.default_timer()
-            data['X'], data['y'] = correct_drift(data['X'], data['y'])
-            self.logger.debug('Drift correction complete in %s seconds.',
-                              timeit.default_timer() - t)
-
         # Prep Neighborhood_Encoder
         neighborhood_encoder = self.get_model_wrapper(settings.NEIGHBORHOOD_ENCODER,
                                                       batch_size=64)
 
         # Send data to the model
-        app = self.get_grpc_app(settings.TRACKING_MODEL, CellTracking,
+        app = self.get_grpc_app(settings.CALIBAN_MODEL, CellTracking,
                                 neighborhood_encoder=neighborhood_encoder,
                                 birth=settings.BIRTH,
                                 death=settings.DEATH,

redis_consumer/consumers/tracking_consumer_test.py → redis_consumer/consumers/caliban_consumer_test.py

@@ -45,12 +45,12 @@
 from redis_consumer.testing_utils import _get_image
 
 
-class TestTrackingConsumer(object):
+class TestCalibanConsumer(object):
     # pylint: disable=R0201,W0621
     def test_is_valid_hash(self, mocker, redis_client):
         queue = 'track'
         storage = DummyStorage()
-        consumer = consumers.TrackingConsumer(redis_client, storage, queue)
+        consumer = consumers.CalibanConsumer(redis_client, storage, queue)
 
         mocker.patch.object(redis_client, 'hget', lambda x, y: x.split(':')[-1])
 
@@ -68,7 +68,7 @@ def test_is_valid_hash(self, mocker, redis_client):
     def test__load_data(self, tmpdir, mocker, redis_client):
         queue = 'track'
         storage = DummyStorage()
-        consumer = consumers.TrackingConsumer(redis_client, storage, queue)
+        consumer = consumers.CalibanConsumer(redis_client, storage, queue)
         tmpdir = str(tmpdir)
         exp = random.randint(0, 99)
 
@@ -155,9 +155,8 @@ def test__consume(self, mocker, redis_client):
             model=mock_model,
         )
 
-        consumer = consumers.TrackingConsumer(redis_client, storage, queue)
+        consumer = consumers.CalibanConsumer(redis_client, storage, queue)
 
-        mocker.patch.object(settings, 'DRIFT_CORRECT_ENABLED', True)
         mocker.patch.object(consumer, 'get_grpc_app',
                             lambda *x, **y: mock_app)
         # mock get_model_wrapper for neighborhood encoder

redis_consumer/settings.py

@@ -92,8 +92,8 @@
 
 # Tracking settings
 TRACKING_MODEL = config('TRACKING_MODEL', default='TrackingModelInf:4', cast=str)
+CALIBAN_MODEL = config('CALIBAN_MODEL', default=TRACKING_MODEL, cast=str)
 NEIGHBORHOOD_ENCODER = config('NEIGHBORHOOD_ENCODER', default='TrackingModelNE:2', cast=str)
-DRIFT_CORRECT_ENABLED = config('DRIFT_CORRECT_ENABLED', default=False, cast=bool)
 
 # tracking.cell_tracker settings TODO: can we extract from model_metadata?
 MAX_DISTANCE = config('MAX_DISTANCE', default=50, cast=int)

requirements.txt

@@ -1,10 +1,9 @@
 # deepcell packages
 deepcell-cpu~=0.11.0
-deepcell-toolbox~=0.10.2
+deepcell-toolbox~=0.10.3
 deepcell-tracking~=0.5.2
 tensorflow-cpu~=2.5.2
 tifffile>=2020.9.3
-scikit-image>=0.14.5,<0.17.0
 numpy>=1.16.6
 
 # tensorflow-serving-apis and gRPC dependencies