Fix len() method in SingleCellTrajectory class and add filter_empty p…

…arameter to write_json method

livecellx/core/single_cell.py

@@ -955,7 +955,8 @@ def __repr__(self) -> str:
         return f"SingleCellTrajectory(track_id={self.track_id}, #timeframe set={len(self)})"
 
     def __len__(self):
-        return self.get_timeframe_span_length()
+        # return self.get_timeframe_span_length()
+        return len(self.timeframe_to_single_cell)
 
     def __getitem__(self, timeframe: int) -> SingleCellStatic:
         if timeframe not in self.timeframe_set:
@@ -1402,7 +1403,10 @@ def load_from_json_dict_parallel(self, json_dict):
         self._post_load_trajectories()
         return self
 
-    def write_json(self, path, dataset_json_dir=None):
+
+    def write_json(self, path, dataset_json_dir=None, filter_empty=True):
+        if filter_empty:
+            self.remove_empty_sct(inplace=True)
         with open(path, "w+") as f:
             json.dump(self.to_json_dict(dataset_json_dir=dataset_json_dir), f, cls=LiveCellEncoder)
 
@@ -1473,7 +1477,7 @@ def _next_track_id(self):
             return 0
         return max(self.get_track_ids()) + 1
 
-    def remove_empty_sct(self):
+    def remove_empty_sct(self, inplace=True):
         remove_tids = []
         remove_scs = []
         for tid, sct in self:
@@ -1486,9 +1490,17 @@ def remove_empty_sct(self):
             if to_be_removed:
                 remove_tids.append(tid)
                 remove_scs.extend(_tmp_scs)
-        for tid in remove_tids:
-            self.pop_trajectory(tid)
-
+        if inplace:
+            for tid in remove_tids:
+                self.pop_trajectory(tid)
+            return self
+        else:
+            new_sctc = SingleCellTrajectoryCollection()
+            for tid, sct in self:
+                if tid not in remove_tids:
+                    new_sctc.add_trajectory(sct)
+            return new_sctc
+
 
 def create_sctc_from_scs(scs: List[SingleCellStatic]) -> SingleCellTrajectoryCollection:
     temp_sc_trajs = SingleCellTrajectoryCollection()
@@ -1887,3 +1899,4 @@ def largest_bbox(scs):
         if bbox[3] > largest_bbox[3]:
             largest_bbox[3] = bbox[3]
     return largest_bbox
+

tests/sctc_test.py

@@ -51,3 +51,20 @@ def test_filter_sctc_by_time_span(self):
         for _, sct in new_sctc:
             for timeframe in sct.timeframe_set:
                 self.assertTrue(0 <= timeframe <= 1, f"Timeframe {timeframe} is outside of the span (0, 1)")
+
+    def test_filter_empty_case_sc_empty_contour(self):
+        new_sctc = SingleCellTrajectoryCollection()
+        new_sct = SingleCellTrajectory()
+        empty_sc = SingleCellStatic(contour=[])
+        new_sct.add_single_cell(timeframe=1, sc=empty_sc)
+        new_sctc.add_trajectory(new_sct)
+        new_sctc.remove_empty_sct(inplace=True)
+        self.assertEqual(len(new_sctc), 0)
+
+
+    def test_filter_empty_case_empty_sct(self):
+        new_sctc = SingleCellTrajectoryCollection()
+        new_sct = SingleCellTrajectory()
+        new_sctc.add_trajectory(new_sct)
+        new_sctc.remove_empty_sct(inplace=True)
+        self.assertEqual(len(new_sctc), 0)