Merge pull request #10 from TravisWheelerLab/dipui_file

Dipui file
TravisWheelerLab · Apr 18, 2024 · 748e7c6 · 748e7c6
2 parents 25859e7 + cdfcea4
commit 748e7c6
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 4 deletions.
diff --git a/diplomat/__init__.py b/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.1.2"
+__version__ = "0.1.3"
 # Can be used by functions to determine if diplomat was invoked through it's CLI interface.
 CLI_RUN = False
 

diff --git a/diplomat/core_ops.py b/diplomat/core_ops.py
@@ -174,6 +174,7 @@ def track_with(
     predictor: Optional[str] = None,
     predictor_settings: Optional[Dict[str, Any]] = None,
     help_extra: Flag = False,
+    dipui_file: Optional[PathLike] = None,
     **extra_args
 ):
     """
@@ -238,6 +239,7 @@ def track_with(
             num_outputs=num_outputs,
             predictor=predictor,
             predictor_settings=predictor_settings,
+            dipui_file=dipui_file,
             **_get_casted_args(selected_frontend.analyze_frames, extra_args)
         )
 

diff --git a/diplomat/frontends/deeplabcut/predict_videos_dlc.py b/diplomat/frontends/deeplabcut/predict_videos_dlc.py
@@ -169,7 +169,8 @@ def analyze_videos(
                 save_as_csv,
                 destination_folder,
                 predictor_cls,
-                predictor_settings
+                predictor_settings,
+                dipui_file=None
             )
     else:
         print("No videos found!")

diff --git a/diplomat/predictors/fpe/frame_pass_engine.py b/diplomat/predictors/fpe/frame_pass_engine.py
@@ -477,7 +477,11 @@ def get_settings(cls) -> ConfigSpec:
                 SparseTrackingData.SparseModes.OFFSET_DOMINATION.name,
                 type_casters.Literal(*[mode.name for mode in SparseTrackingData.SparseModes]),
                 "The mode to utilize during sparsification."
-            )
+            ),
+            "dipui_file": (
+                None, type_casters.Union(type_casters.Literal(None), str), 
+                "A path specifying where to save the dipui file"
+            ),
         }
 
 

diff --git a/diplomat/predictors/sfpe/segmented_frame_pass_engine.py b/diplomat/predictors/sfpe/segmented_frame_pass_engine.py
@@ -2,6 +2,7 @@
 import os
 import shutil
 from enum import Enum
+from datetime import datetime
 from pathlib import Path
 from typing import List, Tuple, Optional, Any, Callable, Sequence, Iterable, BinaryIO
 import numpy as np
@@ -463,7 +464,16 @@ def get_frame_holder(self):
         Returns:
             DiskBackedForwardBackwardData: An object that holds and manages access to the frames data.
         """
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
         output_path = Path(self.video_metadata["output-file-path"]).resolve()
+        if self.settings.dipui_file is not None:
+            dipui_path = str(self.settings.dipui_file)
+            if os.path.exists(dipui_path):
+                output_path = dipui_path.replace(".dipui","") + f"_{timestamp}.dipui"
+            output_path = Path(output_path).resolve()
+        else:
+            output_path = output_path.parent / f"{output_path.stem}_{timestamp}{output_path.suffix}"
+
         video_path = Path(self.video_metadata["orig-video-path"]).resolve()
         disk_path = output_path.parent / (output_path.stem + ".dipui")
 
@@ -1584,7 +1594,8 @@ def get_settings(cls) -> ConfigSpec:
                 100,
                 type_casters.RangedInteger(1, np.inf),
                 "Size of lifo cache used to temporarily store frames loaded from disk if running in disk storage_mode."
-            )
+            ),
+            "dipui_file": (None, type_casters.Union(type_casters.Literal(None), str), "A path specifying where to save the dipui file"),
         }
 
     @classmethod