Migrate DeepCell Tracking to TF2

deepcell/image_generators/tracking.py

@@ -894,8 +894,9 @@ def _get_batches_of_transformed_samples(self, index_array):
 
             batch_y[i, type_cell] = 1
 
-        # prepare final batch list
-        batch_list = []
+        # create dictionary to house generator outputs
+        # Start with batch inputs to model
+        batch_inputs = {}
         for feature_i, feature in enumerate(self.features):
             batch_feature_1, batch_feature_2 = batch_features[feature_i]
             # Remove singleton dimensions (if min_track_length is 1)
@@ -904,10 +905,13 @@ def _get_batches_of_transformed_samples(self, index_array):
                 batch_feature_1 = np.squeeze(batch_feature_1, axis=axis)
                 batch_feature_2 = np.squeeze(batch_feature_2, axis=axis)
 
-            batch_list.append(batch_feature_1)
-            batch_list.append(batch_feature_2)
+            batch_inputs['{}_input1'.format(feature)] = batch_feature_1
+            batch_inputs['{}_input2'.format(feature)] = batch_feature_2
 
-        return batch_list, batch_y
+        # Dict to house training output (model target)
+        batch_outputs = {'classification': batch_y}
+
+        return batch_inputs, batch_outputs
 
     def next(self):
         """For python 2.x. Returns the next batch.

deepcell/model_zoo/featurenet.py

@@ -693,8 +693,8 @@ def compute_feature_extractor(feature, shape):
         re_shape = compute_reshape(feature)
         feature_extractor = compute_feature_extractor(feature, in_shape)
 
-        layer_1 = Input(shape=in_shape)
-        layer_2 = Input(shape=in_shape)
+        layer_1 = Input(shape=in_shape, name='{}_input1'.format(feature))
+        layer_2 = Input(shape=in_shape, name='{}_input2'.format(feature))
 
         inputs.extend([layer_1, layer_2])
 
@@ -727,7 +727,7 @@ def compute_feature_extractor(feature, shape):
     dense2 = Dense(128)(relu1)
     bn2 = BatchNormalization(axis=channel_axis)(dense2)
     relu2 = Activation('relu')(bn2)
-    dense3 = Dense(3, activation='softmax', dtype=K.floatx())(relu2)
+    dense3 = Dense(3, activation='softmax', name='classification', dtype=K.floatx())(relu2)
 
     # Instantiate model
     final_layer = dense3

deepcell/training.py

@@ -486,12 +486,10 @@ def loss_function(y_true, y_pred):
         horizontal_flip=0,
         vertical_flip=0)
 
+    # same_probability values have varied from 0.5 to 5.0
     total_train_pairs = tracking_utils.count_pairs(train_dict['y'], same_probability=5.0)
     total_test_pairs = tracking_utils.count_pairs(val_dict['y'], same_probability=5.0)
 
-    # total_train_pairs = tracking_utils.count_pairs(train_dict['y'], same_probability=0.5)
-    # total_test_pairs = tracking_utils.count_pairs(val_dict['y'], same_probability=0.5)
-
     train_data = datagen.flow(
         train_dict,
         seed=seed,
@@ -515,6 +513,64 @@ def loss_function(y_true, y_pred):
     print('batch size:', batch_size)
     print('validation_steps: ', total_test_pairs // batch_size)
 
+    # Make dicts to map the two generator outputs to the Dataset and model
+    # input here is model input and output is model output
+    features = sorted(features)
+
+    input_type_dict = {}
+    input_shape_dict = {}
+    for feature in features:
+
+        feature_name1 = '{}_input1'.format(feature)
+        feature_name2 = '{}_input2'.format(feature)
+
+        input_type_dict[feature_name1] = tf.float32
+        input_type_dict[feature_name2] = tf.float32
+
+        if feature == 'appearance':
+            app1 = tuple([None, train_data.min_track_length,
+                          train_data.crop_dim, train_data.crop_dim, 1])
+            app2 = tuple([None, 1, train_data.crop_dim, train_data.crop_dim, 1])
+
+            input_shape_dict[feature_name1] = app1
+            input_shape_dict[feature_name2] = app2
+
+        elif feature == 'distance':
+            dist1 = tuple([None, train_data.min_track_length, 2])
+            dist2 = tuple([None, 1, 2])
+
+            input_shape_dict[feature_name1] = dist1
+            input_shape_dict[feature_name2] = dist2
+
+        elif feature == 'neighborhood':
+            neighborhood_size = 2 * train_data.neighborhood_scale_size + 1
+            neigh1 = tuple([None, train_data.min_track_length,
+                            neighborhood_size, neighborhood_size, 1])
+            neigh2 = tuple([None, 1, neighborhood_size, neighborhood_size, 1])
+
+            input_shape_dict[feature_name1] = neigh1
+            input_shape_dict[feature_name2] = neigh2
+
+        elif feature == 'regionprop':
+            rprop1 = tuple([None, train_data.min_track_length, 3])
+            rprop2 = tuple([None, 1, 3])
+
+            input_shape_dict[feature_name1] = rprop1
+            input_shape_dict[feature_name2] = rprop2
+
+    output_type_dict = {'classification': tf.int32}
+    # Ouput_shape has to be None because we dont know how many cells
+    output_shape_dict = {'classification': (None, 3)}
+
+    train_dataset = Dataset.from_generator(
+        lambda: train_data,
+        (input_type_dict, output_type_dict),
+        output_shapes=(input_shape_dict, output_shape_dict))
+    val_dataset = Dataset.from_generator(
+        lambda: val_data,
+        (input_type_dict, output_type_dict),
+        output_shapes=(input_shape_dict, output_shape_dict))
+
     train_callbacks = get_callbacks(
         model_path, lr_sched=lr_sched,
         tensorboard_log_dir=log_dir,
@@ -523,10 +579,10 @@ def loss_function(y_true, y_pred):
 
     # fit the model on the batches generated by datagen.flow()
     loss_history = model.fit(
-        train_data,
+        train_dataset,
         steps_per_epoch=total_train_pairs // batch_size,
         epochs=n_epoch,
-        validation_data=val_data,
+        validation_data=val_dataset,
         validation_steps=total_test_pairs // batch_size,
         callbacks=train_callbacks)