fix live inference and display, black and isort

Author: dikraMasrour

dlclive/__init__.py

@@ -5,9 +5,10 @@
 Licensed under GNU Lesser General Public License v3.0
 """
 
-from dlclive.version import __version__, VERSION
-from dlclive.dlclive import DLCLive
 from dlclive.display import Display
-from dlclive.processor import Processor
+from dlclive.dlclive import DLCLive
 from dlclive.predictor import HeatmapPredictor
+from dlclive.processor import Processor
+from dlclive.version import VERSION, __version__
+
 # from dlclive.benchmark import benchmark, benchmark_videos, download_benchmarking_data

dlclive/benchmark.py

@@ -5,33 +5,31 @@
 Licensed under GNU Lesser General Public License v3.0
 """
 
-
-import platform
 import os
-import time
-import sys
-import warnings
+import pickle
+import platform
 import subprocess
+import sys
+import time
 import typing
-import pickle
+import warnings
+
 import colorcet as cc
-from PIL import ImageColor
 import ruamel
+from PIL import ImageColor
 
 try:
     from pip._internal.operations import freeze
 except ImportError:
     from pip.operations import freeze
 
-from tqdm import tqdm
+import cv2
 import numpy as np
 import tensorflow as tf
-import cv2
+from tqdm import tqdm
 
-from dlclive import DLCLive
-from dlclive import VERSION
+from dlclive import VERSION, DLCLive
 from dlclive import __file__ as dlcfile
-
 from dlclive.utils import decode_fourcc
 
 
@@ -42,8 +40,9 @@ def download_benchmarking_data(
     """
     Downloads a DeepLabCut-Live benchmarking Data (videos & DLC models).
     """
-    import urllib.request
     import tarfile
+    import urllib.request
+
     from tqdm import tqdm
 
     def show_progress(count, block_size, total_size):
@@ -75,7 +74,7 @@ def tarfilenamecutting(tarf):
 
 
 def get_system_info() -> dict:
-    """ Return summary info for system running benchmark
+    """Return summary info for system running benchmark
     Returns
     -------
     dict
@@ -165,7 +164,7 @@ def benchmark(
     save_video=False,
     output=None,
 ) -> typing.Tuple[np.ndarray, tuple, bool, dict]:
-    """ Analyze DeepLabCut-live exported model on a video:
+    """Analyze DeepLabCut-live exported model on a video:
     Calculate inference time,
     display keypoints, or
     get poses/create a labeled video
@@ -193,7 +192,7 @@ def benchmark(
     n_frames : int, optional
         number of frames to run inference on, by default 1000
     print_rate : bool, optional
-        flat to print inference rate frame by frame, by default False
+        flag to print inference rate frame by frame, by default False
     display : bool, optional
         flag to display keypoints on images. Useful for checking the accuracy of exported models.
     pcutoff : float, optional
@@ -440,7 +439,7 @@ def benchmark(
 def save_inf_times(
     sys_info, inf_times, im_size, TFGPUinference, model=None, meta=None, output=None
 ):
-    """ Save inference time data collected using :function:`benchmark` with system information to a pickle file.
+    """Save inference time data collected using :function:`benchmark` with system information to a pickle file.
     This is primarily used through :function:`benchmark_videos`
 
 
@@ -666,8 +665,7 @@ def benchmark_videos(
 
 
 def main():
-    """Provides a command line interface :function:`benchmark_videos`
-    """
+    """Provides a command line interface :function:`benchmark_videos`"""
 
     import argparse
 

dlclive/benchmark_pytorch.py

@@ -13,7 +13,8 @@
 from PIL import ImageColor
 from pip._internal.operations import freeze
 
-from dlclive import VERSION, DLCLive
+from dlclive import DLCLive
+from dlclive.version import VERSION
 
 
 def get_system_info() -> dict:
@@ -482,8 +483,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-
-
-# Example how to run in command line:
-# python benchmark_pytorch.py /path/to/model /path/to/video DLC cuda -p FP32 -d -c 0.5 -dr 5 -r 0.5 -x 10 630 10 470 --save-poses --save-video --draw-keypoint-names --cmap bmy --save-dir
-# python benchmark_pytorch.py /Users/annastuckert/Documents/DLC_AI_Residency/DLC_AI2024/DeepLabCut-live/dlc-live-dummy/ventral-gait/resnet.onnx /Users/annastuckert/Documents/DLC_AI_Residency/DLC_AI2024/DeepLabCut-live/dlc-live-dummy/ventral-gait/1_20cms_0degUP_first_1s.avi DLC cuda -p FP32 -d -r 0.5 --save-poses --save-video --draw-keypoint-names --cmap bmy --save-dir /Users/annastuckert/Documents/DLC_AI_Residency/DLC_AI2024/DeepLabCut-live/dlc-live-dummy/ventral-gait/out

dlclive/check_install/check_install.py

@@ -5,19 +5,16 @@
 Licensed under GNU Lesser General Public License v3.0
 """
 
-
-import sys
+import argparse
 import shutil
-import warnings
-
-from dlclive import benchmark_videos
+import sys
 import urllib.request
-import argparse
+import warnings
 from pathlib import Path
-from dlclibrary.dlcmodelzoo.modelzoo_download import (
-    download_huggingface_model,
-)
 
+from dlclibrary.dlcmodelzoo.modelzoo_download import download_huggingface_model
+
+from dlclive import benchmark_videos
 
 MODEL_NAME = "superanimal_quadruped"
 SNAPSHOT_NAME = "snapshot-700000.pb"
@@ -27,28 +24,33 @@ def urllib_pbar(count, blockSize, totalSize):
     percent = int(count * blockSize * 100 / totalSize)
     outstr = f"{round(percent)}%"
     sys.stdout.write(outstr)
-    sys.stdout.write("\b"*len(outstr))
+    sys.stdout.write("\b" * len(outstr))
     sys.stdout.flush()
 
 
 def main():
     parser = argparse.ArgumentParser(
-        description="Test DLC-Live installation by downloading and evaluating a demo DLC project!")
-    parser.add_argument('--nodisplay', action='store_false', help="Run the test without displaying tracking")
+        description="Test DLC-Live installation by downloading and evaluating a demo DLC project!"
+    )
+    parser.add_argument(
+        "--nodisplay",
+        action="store_false",
+        help="Run the test without displaying tracking",
+    )
     args = parser.parse_args()
     display = args.nodisplay
 
     if not display:
-        print('Running without displaying video')
+        print("Running without displaying video")
 
     # make temporary directory in $HOME
     # TODO: why create this temp directory in $HOME?
     print("\nCreating temporary directory...\n")
-    tmp_dir = Path().home() / 'dlc-live-tmp'
-    tmp_dir.mkdir(mode=0o775,exist_ok=True)
+    tmp_dir = Path().home() / "dlc-live-tmp"
+    tmp_dir.mkdir(mode=0o775, exist_ok=True)
 
-    video_file = str(tmp_dir / 'dog_clip.avi')
-    model_dir = tmp_dir / 'DLC_Dog_resnet_50_iteration-0_shuffle-0'
+    video_file = str(tmp_dir / "dog_clip.avi")
+    model_dir = tmp_dir / "DLC_Dog_resnet_50_iteration-0_shuffle-0"
 
     # download dog test video from github:
     # TODO: Should check if the video's already there before downloading it (should have been cloned with the files)
@@ -58,25 +60,31 @@ def main():
 
     # download model from the DeepLabCut Model Zoo
     if Path(model_dir / SNAPSHOT_NAME).exists():
-        print('Model already downloaded, using cached version')
+        print("Model already downloaded, using cached version")
     else:
         print("Downloading full_dog model from the DeepLabCut Model Zoo...")
         download_huggingface_model(MODEL_NAME, model_dir)
 
     # assert these things exist so we can give informative error messages
     assert Path(video_file).exists(), f"Missing video file {video_file}"
-    assert Path(model_dir / SNAPSHOT_NAME).exists(), f"Missing model file {model_dir / SNAPSHOT_NAME}"
+    assert Path(
+        model_dir / SNAPSHOT_NAME
+    ).exists(), f"Missing model file {model_dir / SNAPSHOT_NAME}"
 
     # run benchmark videos
     print("\n Running inference...\n")
-    benchmark_videos(str(model_dir), video_file, display=display, resize=0.5, pcutoff=0.25)
+    benchmark_videos(
+        str(model_dir), video_file, display=display, resize=0.5, pcutoff=0.25
+    )
 
     # deleting temporary files
     print("\n Deleting temporary files...\n")
     try:
         shutil.rmtree(tmp_dir)
     except PermissionError:
-        warnings.warn(f'Could not delete temporary directory {str(tmp_dir)} due to a permissions error, but otherwise dlc-live seems to be working fine!')
+        warnings.warn(
+            f"Could not delete temporary directory {str(tmp_dir)} due to a permissions error, but otherwise dlc-live seems to be working fine!"
+        )
 
     print("\nDone!\n")
 

dlclive/display.py

@@ -9,9 +9,10 @@
 
 import colorcet as cc
 import numpy as np
-from dlclive import utils
 from PIL import Image, ImageDraw, ImageTk
 
+from dlclive import utils
+
 
 class Display(object):
     """

dlclive/dlclive.py

@@ -24,7 +24,8 @@
 from dlclive import utils
 from dlclive.display import Display
 from dlclive.exceptions import DLCLiveError, DLCLiveWarning
-from dlclive.pose import argmax_pose_predict, extract_cnn_output, multi_pose_predict
+from dlclive.pose import (argmax_pose_predict, extract_cnn_output,
+                          multi_pose_predict)
 from dlclive.predictor import HeatmapPredictor
 
 if typing.TYPE_CHECKING:
@@ -277,7 +278,11 @@ def load_model(self):
         elif self.model_type == "onnx":
             model_paths = glob.glob(os.path.normpath(self.path + "/*.onnx"))
             if self.precision == "FP16":
-                model_path = [model_paths[i] for i in range(len(model_paths)) if "fp16" in model_paths[i]][0]
+                model_path = [
+                    model_paths[i]
+                    for i in range(len(model_paths))
+                    if "fp16" in model_paths[i]
+                ][0]
                 print(model_path)
             else:
                 model_path = model_paths[0]
@@ -294,13 +299,16 @@ def load_model(self):
                 )
             # ! TODO implement if statements for choice of tensorrt engine options (precision, and caching)
             elif self.device == "tensorrt":
-                provider = [("TensorrtExecutionProvider", {
-                    "trt_engine_cache_enable": True,
-                    "trt_engine_cache_path": "./trt_engines"
-                })]
-                self.sess = ort.InferenceSession(
-                    model_path, opts, providers=provider
-                )
+                provider = [
+                    (
+                        "TensorrtExecutionProvider",
+                        {
+                            "trt_engine_cache_enable": True,
+                            "trt_engine_cache_path": "./trt_engines",
+                        },
+                    )
+                ]
+                self.sess = ort.InferenceSession(model_path, opts, providers=provider)
             self.predictor = HeatmapPredictor.build(self.cfg)
 
             if not os.path.isfile(model_path):
@@ -333,7 +341,7 @@ def init_inference(self, frame=None, **kwargs):
         # load model
         self.load_model()
 
-        inf_time = 0.
+        inf_time = 0.0
         # get pose of first frame (first inference is often very slow)
         if frame is not None:
             pose, inf_time = self.get_pose(frame, **kwargs)
@@ -356,7 +364,7 @@ def get_pose(self, frame=None, **kwargs):
         pose :class:`numpy.ndarray`
             the pose estimated by DeepLabCut for the input image
         """
-        inf_time = 0.
+        inf_time = 0.0
         if frame is None:
             raise DLCLiveError("No frame provided for live pose estimation")
 
@@ -381,8 +389,10 @@ def get_pose(self, frame=None, **kwargs):
             self.pose = self.pose["bodypart"]
 
         elif self.model_type == "onnx":
-            if self.precision == "FP32": frame = processed_frame.astype(np.float32)
-            elif self.precision == "FP16": frame = processed_frame.astype(np.float16)
+            if self.precision == "FP32":
+                frame = processed_frame.astype(np.float32)
+            elif self.precision == "FP16":
+                frame = processed_frame.astype(np.float16)
 
             frame = np.transpose(frame, (2, 0, 1))
             frame = np.expand_dims(frame, axis=0)

dlclive/exceptions.py

@@ -7,12 +7,12 @@
 
 
 class DLCLiveError(Exception):
-    """ Generic error type for incorrect use of the DLCLive class """
+    """Generic error type for incorrect use of the DLCLive class"""
 
     pass
 
 
 class DLCLiveWarning(Warning):
-    """ Generic warning for incorrect use of the DLCLive class """
+    """Generic warning for incorrect use of the DLCLive class"""
 
     pass

dlclive/graph.py

@@ -5,7 +5,6 @@
 Licensed under GNU Lesser General Public License v3.0
 """
 
-
 import tensorflow as tf
 
 vers = (tf.__version__).split(".")