From 526b65b9aff539e340bac0aad148c762cf5eadba Mon Sep 17 00:00:00 2001
From: willgraf <7930703+willgraf@users.noreply.github.com>
Date: Wed, 18 Aug 2021 13:28:17 -0700
Subject: [PATCH] `is_valid_lineage` should check if the daughter ID is valid.
 (#70)

* Check for invalid daughters (KeyError)

* Add simple test for invalid daughter ID.

* Bump version to 0.4.5
---
 deepcell_tracking/utils.py      | 16 +++++++++++-----
 deepcell_tracking/utils_test.py |  5 +++++
 setup.py                        |  2 +-
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/deepcell_tracking/utils.py b/deepcell_tracking/utils.py
index d83df47..31ab53c 100644
--- a/deepcell_tracking/utils.py
+++ b/deepcell_tracking/utils.py
@@ -35,6 +35,7 @@
 import re
 import tarfile
 import tempfile
+import warnings
 
 import numpy as np
 
@@ -481,13 +482,18 @@ def is_valid_lineage(lineage):
     Returns:
         bool: Whether or not the lineage is valid.
     """
-    for cell_lineage in lineage.values():
+    for cell_label, cell_lineage in lineage.items():
         # Get last frame of parent
         last_parent_frame = cell_lineage['frames'][-1]
 
         for daughter in cell_lineage['daughters']:
-            # get first frame of daughter
-            first_daughter_frame = lineage[daughter]['frames'][0]
+            try:
+                # get first frame of daughter
+                first_daughter_frame = lineage[daughter]['frames'][0]
+            except KeyError:
+                warnings.warn('lineage {} has invalid daughters: {}'.format(
+                    cell_label, cell_lineage['daughters']))
+                return False
 
             # Check that daughter's start frame is one larger than parent end frame
             if first_daughter_frame - last_parent_frame != 1:
@@ -660,13 +666,13 @@ def __init__(self, path=None, tracked_data=None,
 
     def _correct_lineages(self):
         """Ensure sequential labels for all batches"""
-        new_lineages = {}
+        new_lineages = []
         for batch in range(self.y.shape[0]):
 
             y_relabel, new_lineage = relabel_sequential_lineage(
                 self.y[batch], self.lineages[batch])
 
-            new_lineages[batch] = new_lineage
+            new_lineages.append(new_lineage)
             self.y[batch] = y_relabel
 
         self.lineages = new_lineages
diff --git a/deepcell_tracking/utils_test.py b/deepcell_tracking/utils_test.py
index e789628..a7888d3 100644
--- a/deepcell_tracking/utils_test.py
+++ b/deepcell_tracking/utils_test.py
@@ -349,6 +349,11 @@ def test_is_valid_lineage(self):
         bad_lineage[2]['frames'] = [2]
         assert not utils.is_valid_lineage(bad_lineage)
 
+        # change one of cell 0's daughters to an invalid ID.
+        bad_lineage = copy.copy(lineage)
+        bad_lineage[0]['daughters'][0] = 3
+        assert not utils.is_valid_lineage(bad_lineage)
+
     def test_get_image_features(self):
         num_labels = 3
         y = get_annotated_image(num_labels=num_labels, sequential=True)
diff --git a/setup.py b/setup.py
index 4a26170..fa0960f 100644
--- a/setup.py
+++ b/setup.py
@@ -37,7 +37,7 @@
     readme = f.read()
 
 
-VERSION = '0.4.4'
+VERSION = '0.4.5'
 NAME = 'DeepCell_Tracking'
 DESCRIPTION = 'Tracking cells and lineage with deep learning.'
 LICENSE = 'LICENSE'