New version.

Author: isaacrobinson2000

RELEASE_NOTES.md

@@ -1 +1,2 @@
-Reduces GPU memory usage when running inference on top of sleap.
+ - Adds support for DeepLabCut's csv format.
+ - Unifies multithreading configuration to be through a single argument (thread_count).

diplomat/__init__.py

@@ -2,7 +2,7 @@
 A tool providing multi-animal tracking capabilities on top of other Deep learning based tracking software.
 """
 
-__version__ = "0.3.7"
+__version__ = "0.3.8"
 # Can be used by functions to determine if diplomat was invoked through it's CLI interface.
 CLI_RUN = False
 

diplomat/predictors/fpe/frame_pass.py

@@ -30,7 +30,7 @@ def __init__(
         self,
         width: int,
         height: int,
-        multi_threading_allowed: bool,
+        thread_count: int,
         config: Dict[str, Any]
     ):
         # Set defaults to forward iteration...
@@ -41,7 +41,8 @@ def __init__(
 
         self.__width = width
         self.__height = height
-        self.__multi_threading_allowed = multi_threading_allowed
+        self.__multi_threading_allowed = thread_count > 1
+        self.__thread_count = thread_count
 
         self._config = Config(config, self.get_config_options())
         self._frame_data = None
@@ -58,6 +59,10 @@ def height(self) -> int:
     def multi_threading_allowed(self) -> bool:
         return self.__multi_threading_allowed
 
+    @property
+    def thread_count(self) -> int:
+        return self.__thread_count
+
     def _get_step_controls(self) -> Tuple[oint, oint, oint, oint]:
         return self._start, self._stop, self._step, self._prior_off
 

diplomat/predictors/fpe/frame_pass_loader.py

@@ -25,10 +25,10 @@ def __call__(
         self,
         width: int,
         height: int,
-        allow_multi_threading: bool = True
+        thread_count: int = 1
     ) -> FramePass:
         return self._clazz(
-            width, height, allow_multi_threading, self._config
+            width, height, thread_count, self._config
         )
 
     @classmethod

diplomat/predictors/fpe/frame_passes/cluster_frames.py

@@ -56,17 +56,15 @@ def run_pass(
         if(not in_place):
             raise ValueError("Clustering must be done in place!")
 
-        thread_count = os.cpu_count() if(self.config.thread_count is None) else self.config.thread_count
-
-        if(self.multi_threading_allowed and (thread_count > 0)):
+        if(self.multi_threading_allowed and (self.thread_count > 1)):
             from diplomat.predictors.sfpe.segmented_frame_pass_engine import PoolWithProgress
 
             self._frame_data = fb_data
             self._frame_data.allow_pickle = False
 
             iter_range = RangeSlicer(self._frame_data.frames)[self._start:self._stop:self._step]
 
-            with PoolWithProgress(prog_bar, process_count=thread_count, sub_ticks=1) as pool:
+            with PoolWithProgress(prog_bar, process_count=self.thread_count, sub_ticks=1) as pool:
                 pool.fast_map(
                     ClusterFrames._cluster_frames,
                     lambda i: self._get_frame(iter_range[i]),
@@ -298,12 +296,6 @@ def get_config_options(cls) -> ConfigSpec:
             "max_throwaway_count": (
                 10, float, "The maximum number of clusters to throw away before giving up on clustering a given frame."
             ),
-            "thread_count": (
-                None,
-                tc.Union(tc.Literal(None), tc.RangedInteger(0, np.inf)),
-                "The number of threads to use during processing. If None, uses os.cpu_count(). "
-                "If 0 disables multithreading."
-            ),
             "clustering_mode": (
                 ClusteringMethod.COMPLETE.name,
                 tc.Literal(*[n.name for n in ClusteringMethod]),

diplomat/predictors/fpe/frame_passes/fix_frame.py

@@ -542,7 +542,7 @@ def compute_scores(
         to_index = lambda i: slice(i * cls.SCORES_PER_CHUNK, (i + 1) * cls.SCORES_PER_CHUNK)
         max_dist = np.sqrt(fb_data.metadata.width ** 2 + fb_data.metadata.height ** 2)
 
-        if(thread_count > 0):
+        if(thread_count > 1):
             from ...sfpe.segmented_frame_pass_engine import PoolWithProgress
             with PoolWithProgress(prog_bar, process_count=thread_count, sub_ticks=1) as pool:
                 pool.fast_map(

diplomat/predictors/fpe/frame_passes/mit_viterbi.py

@@ -17,6 +17,9 @@ class NotAPool:
     T = TypeVar("T")
     E = TypeVar("E")
 
+    def __init__(self, *args, **kwargs):
+        pass
+
     def __enter__(self):
         return self
 
@@ -332,7 +335,7 @@ def _run_backtrace(
         backtrace_priors = [None for __ in range(fb_data.num_bodyparts)]
         backtrace_current = [None for __ in range(fb_data.num_bodyparts)]
 
-        with pool_cls() as pool:
+        with pool_cls(self.thread_count) as pool:
             exit_prob = fb_data.metadata.enter_trans_prob
             transition_function = ViterbiTransitionTable(self._gaussian_table, 1, exit_prob, 1 - exit_prob)
 
@@ -422,7 +425,7 @@ def _run_backtrace(
         return fb_data
 
     @staticmethod
-    def _get_pool():
+    def _get_pool(processes):
         # Check globals for a pool...
         """This function sets up a multiprocessing pool for parallel processing,
         improving the efficiency of the algorithm by allowing it to process 
@@ -434,11 +437,11 @@ def _get_pool():
         for ctx, args in [("forkserver", {}), ("spawn", {}), ("fork", {"maxtasksperchild": 1})]:
             try:
                 ctx = get_context(ctx)
-                return ctx.Pool(**args)
+                return ctx.Pool(processes, **args)
             except ValueError:
                 continue
 
-        return get_context().Pool()
+        return get_context().Pool(processes)
 
     def _run_forward(
         self,
@@ -475,7 +478,7 @@ def _run_forward(
         # We only use a pool if the body part group is high enough...
         pool_cls = self._get_pool if(self.multi_threading_allowed and (fb_data.num_bodyparts // meta.num_outputs) > 2) else NotAPool
 
-        with pool_cls() as pool:
+        with pool_cls(self.thread_count) as pool:
             exit_prob = fb_data.metadata.enter_trans_prob
             transition_func = ViterbiTransitionTable(self._gaussian_table, 1, exit_prob, 1 - exit_prob)
 

diplomat/predictors/sfpe/segmented_frame_pass_engine.py

@@ -271,7 +271,7 @@ class PoolWithProgress:
     def __init__(
         self,
         progress_bar: ProgressBar,
-        process_count: int = os.cpu_count(),
+        process_count: int,
         refresh_rate_seconds: float = DEF_MAX_REFRESH_RATE,
         sub_ticks: int = DEF_SUB_TICKS,
         max_worker_reuse: Optional[int] = None,
@@ -764,8 +764,10 @@ def get_maximums(
         return poses
 
     def _run_full_passes(self, progress_bar: Optional[ProgressBar]):
+        thread_count = self._get_thread_count()
+
         for (i, frame_pass_builder) in enumerate(self.FULL_PASSES):
-            frame_pass = frame_pass_builder(self._width, self._height, True)
+            frame_pass = frame_pass_builder(self._width, self._height, thread_count)
 
             if(progress_bar is not None):
                 progress_bar.message(
@@ -861,7 +863,7 @@ def _run_segment_pre_initialized(
         frame_pass_builders: List[FramePassBuilder],
         width: int,
         height: int,
-        allow_multi_threading: bool,
+        thread_count: int,
         fix_frame_idx: int, 
         fix_frame_score: float,
         progress_bar: Optional[ProgressBar] = None,
@@ -871,7 +873,7 @@ def _run_segment_pre_initialized(
             frame_pass_builders,
             width,
             height,
-            allow_multi_threading,
+            thread_count,
             fix_frame_idx, 
             fix_frame_score,
             progress_bar,
@@ -885,7 +887,7 @@ def _run_segment(
         frame_pass_builders: List[FramePassBuilder],
         width: int,
         height: int,
-        allow_multi_threading: bool,
+        thread_count: int,
         fix_frame_idx: int, 
         fix_frame_score,
         progress_bar: Optional[ProgressBar] = None,
@@ -925,7 +927,7 @@ def _run_segment(
             progress_bar.inc_rerun_counter()
 
         for (i, frame_pass_builder) in enumerate(frame_pass_builders):
-            frame_pass = frame_pass_builder(width, height, allow_multi_threading)
+            frame_pass = frame_pass_builder(width, height, thread_count)
 
             sub_frame = frame_pass.run_pass(
                 sub_frame,
@@ -950,7 +952,14 @@ def _get_segment(self, index: int):
         sub_frame.frames = self._frame_holder.frames[start:end]
         sub_frame.metadata = self._frame_holder.metadata
 
-        return (sub_frame, self.SEGMENTED_PASSES, self._width, self._height, False, fix_frame - start, segment_score)
+        return (
+            sub_frame,
+            self.SEGMENTED_PASSES,
+            self._width,
+            self._height,
+            0,
+            fix_frame - start,
+            segment_score)
 
     def _set_segment(self, index: int, frame_data: ForwardBackwardData):
         start, end, fix_frame = self._segments[index]
@@ -1047,31 +1056,20 @@ def _run_segmented_passes(
 
         self._frame_holder.allow_pickle = False
 
-        if(thread_count <= 0):
+        if(thread_count <= 1):
             pass_count = (len(self.SEGMENTED_PASSES) + 1) * total_segments
 
-            passes_can_use_pool = any(b.clazz.UTILIZE_GLOBAL_POOL for b in self.SEGMENTED_PASSES)
-            allow_multithread = self.settings.allow_pass_multithreading
-
             wrapper_bar = NestedProgressIndicator(
                 progress_bar,
                 total=pass_count,
                 ticks=int(self._frame_holder.num_frames / pass_count)
             )
             progress_bar.message("Running on Segments...")
 
-            if(passes_can_use_pool and allow_multithread):
-                with PoolWithProgress.get_optimal_ctx().Pool(processes=os.cpu_count()) as pool:
-                    FramePass.GLOBAL_POOL = AntiCloseObject(pool)
-                    for is_pre_init, segment_idx in self._iter_run_levels(segment_idxs, run_level_data):
-                        for idx in segment_idx:
-                            frm, segs, width, height, __, fix_frame_idx, fix_frame_score = self._get_segment(idx)
-                            self._run_segment(frm, segs, width, height, self.settings.allow_pass_multithreading, fix_frame_idx, fix_frame_score, wrapper_bar, is_pre_init)
-            else:
-                for is_pre_init, segment_idx in self._iter_run_levels(segment_idxs, run_level_data):
-                    for idx in segment_idx:
-                        frm, segs, width, height, __, fix_frame_idx, fix_frame_score = self._get_segment(idx)
-                        self._run_segment(frm, segs, width, height, self.settings.allow_pass_multithreading, fix_frame_idx, fix_frame_score, wrapper_bar, is_pre_init)
+            for is_pre_init, segment_idx in self._iter_run_levels(segment_idxs, run_level_data):
+                for idx in segment_idx:
+                    frm, segs, width, height, __, fix_frame_idx, fix_frame_score = self._get_segment(idx)
+                    self._run_segment(frm, segs, width, height, thread_count, fix_frame_idx, fix_frame_score, wrapper_bar, is_pre_init)
 
             FramePass.GLOBAL_POOL = None
         else:
@@ -1534,11 +1532,6 @@ def get_settings(cls) -> ConfigSpec:
                 "Defaults to None, which resolves to os.cpu_count() at runtime. "
                 "If set to 0, disables multithreading..."
             ),
-            "allow_pass_multithreading": (
-                True,
-                bool,
-                "Whether or not to allow frame passes to utilize multithreading. Defaults to True."
-            ),
             "segment_size": (
                 200,
                 type_casters.RangedInteger(10, np.inf),

diplomat/predictors/supervised_sfpe/supervised_segmented_frame_pass_engine.py

@@ -504,7 +504,6 @@ def _partial_rerun(
             poses[s_i:e_i, :] = poses[s_i:e_i, seg_ord]
         old_poses.get_all()[:] = poses.reshape(old_poses.get_frame_count(), old_poses.get_bodypart_count() * 3)
 
-        
         return (
             self.get_maximums(
                 self._frame_holder,

setup.py

@@ -19,6 +19,7 @@ def _get_readme():
 
 ONNX_TF_DEPS = [
     "h5py",
+    "tables",
     "tensorflow",
     "tf2onnx>=1.16.1",
     "onnx",