Fix rebase issues

Clean up and minor fixes in dlclive/benchmark.py and dlclive/check_install/check_install.py.

Changes include: reorder and consolidate imports; remove duplicated/extra blank lines and repeated statements (duplicate prints, duplicated init_inference call, duplicated function args); normalize long-line wrapping; switch some prints to f-strings; tighten exception handling for video download (catch OSError and URLError); add stacklevel to warnings.warn calls; and minor formatting improvements (single-line tuples, joined long expressions). These are non-functional refactors and small bug/safety fixes to reduce redundancy and improve diagnostics.

Author: C-Achard

dlclive/benchmark.py

@@ -13,10 +13,9 @@
 import sys
 import time
 import warnings
-from typing import TYPE_CHECKING
 from pathlib import Path
-import argparse
-import os
+from typing import TYPE_CHECKING
+
 import colorcet as cc
 import cv2
 import numpy as np
@@ -25,8 +24,7 @@
 from pip._internal.operations import freeze
 from tqdm import tqdm
 
-from dlclive import DLCLive
-from dlclive import VERSION
+from dlclive import VERSION, DLCLive
 from dlclive.engine import Engine
 from dlclive.utils import decode_fourcc
 
@@ -56,20 +54,16 @@ def download_benchmarking_data(
         print(f"{zip_path} already exists. Skipping download.")
     else:
 
-
         def show_progress(count, block_size, total_size):
             pbar.update(block_size)
 
         print(f"Downloading the benchmarking data from {url} ...")
         pbar = tqdm(unit="B", total=0, position=0, desc="Downloading")
 
-        filename, _ = urllib.request.urlretrieve(
-            url, filename=zip_path, reporthook=show_progress
-        )
+        filename, _ = urllib.request.urlretrieve(url, filename=zip_path, reporthook=show_progress)
         pbar.close()
 
     print(f"Extracting {zip_path} to {target_dir} ...")
-    with zipfile.ZipFile(zip_path, "r") as zip_ref:
     with zipfile.ZipFile(zip_path, "r") as zip_ref:
         zip_ref.extractall(target_dir)
 
@@ -192,7 +186,6 @@ def benchmark_videos(
 
         for i in range(len(resize)):
             print(f"\nRun {i + 1} / {len(resize)}\n")
-            print(f"\nRun {i + 1} / {len(resize)}\n")
 
             this_inf_times, this_im_size, meta = benchmark(
                 model_path=model_path,
@@ -296,7 +289,6 @@ def get_system_info() -> dict:
     }
 
 
-def save_inf_times(sys_info, inf_times, im_size, model=None, meta=None, output=None):
 def save_inf_times(sys_info, inf_times, im_size, model=None, meta=None, output=None):
     """Save inference time data collected using :function:`benchmark` with system information to a pickle file.
     This is primarily used through :function:`benchmark_videos`
@@ -366,7 +358,6 @@ def save_inf_times(sys_info, inf_times, im_size, model=None, meta=None, output=N
     return True
 
 
-
 def benchmark(
     model_path: str,
     model_type: str,
@@ -380,8 +371,6 @@ def benchmark(
     dynamic: tuple[bool, float, int] = (False, 0.5, 10),
     n_frames: int = 1000,
     print_rate: bool = False,
-    n_frames: int = 1000,
-    print_rate: bool = False,
     precision: str = "FP32",
     display: bool = True,
     pcutoff: float = 0.5,
@@ -526,9 +515,7 @@ def benchmark(
     frame_index = 0
 
     total_n_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
-    n_frames = int(
-        n_frames if (n_frames > 0) and n_frames < total_n_frames else total_n_frames
-    )
+    n_frames = int(n_frames if (n_frames > 0) and n_frames < total_n_frames else total_n_frames)
     iterator = range(n_frames) if print_rate or display else tqdm(range(n_frames))
     for _ in iterator:
         ret, frame = cap.read()
@@ -543,7 +530,6 @@ def benchmark(
         start_time = time.perf_counter()
         if frame_index == 0:
             pose = dlc_live.init_inference(frame)  # Loads model
-            pose = dlc_live.init_inference(frame)  # Loads model
         else:
             pose = dlc_live.get_pose(frame)
 
@@ -552,9 +538,7 @@ def benchmark(
         times.append(inf_time)
 
         if print_rate:
-            print(
-                "Inference rate = {:.3f} FPS".format(1 / inf_time), end="\r", flush=True
-            )
+            print(f"Inference rate = {1 / inf_time:.3f} FPS", end="\r", flush=True)
 
         if save_video:
             draw_pose_and_write(
@@ -567,15 +551,12 @@ def benchmark(
                 display_radius=display_radius,
                 draw_keypoint_names=draw_keypoint_names,
                 vwriter=vwriter,
-                vwriter=vwriter,
             )
 
         frame_index += 1
 
     if print_rate:
-        print(
-            "Mean inference rate: {:.3f} FPS".format(np.mean(1 / np.array(times)[1:]))
-        )
+        print(f"Mean inference rate: {np.mean(1 / np.array(times)[1:]):.3f} FPS")
 
     metadata = _get_metadata(video_path=video_path, cap=cap, dlc_live=dlc_live)
     metadata = _get_metadata(video_path=video_path, cap=cap, dlc_live=dlc_live)
@@ -593,9 +574,7 @@ def benchmark(
         else:
             individuals = []
         n_individuals = len(individuals) or 1
-        save_poses_to_files(
-            video_path, save_dir, n_individuals, bodyparts, poses, timestamp=timestamp
-        )
+        save_poses_to_files(video_path, save_dir, n_individuals, bodyparts, poses, timestamp=timestamp)
 
     return times, im_size, metadata
 
@@ -608,13 +587,6 @@ def setup_video_writer(
     cmap: str,
     fps: float,
     frame_size: tuple[int, int],
-    video_path: str,
-    save_dir: str,
-    timestamp: str,
-    num_keypoints: int,
-    cmap: str,
-    fps: float,
-    frame_size: tuple[int, int],
 ):
     # Set colors and convert to RGB
     cmap_colors = getattr(cc, cmap)
@@ -623,9 +595,7 @@ def setup_video_writer(
     # Define output video path
     video_path = Path(video_path)
     video_name = video_path.stem  # filename without extension
-    output_video_path = (
-        Path(save_dir) / f"{video_name}_DLCLIVE_LABELLED_{timestamp}.mp4"
-    )
+    output_video_path = Path(save_dir) / f"{video_name}_DLCLIVE_LABELLED_{timestamp}.mp4"
 
     # Get video writer setup
     fourcc = cv2.VideoWriter_fourcc(*"mp4v")
@@ -639,7 +609,6 @@ def setup_video_writer(
     return colors, vwriter
 
 
-
 def draw_pose_and_write(
     frame: np.ndarray,
     pose: np.ndarray,
@@ -656,9 +625,7 @@ def draw_pose_and_write(
 
     if resize is not None and resize != 1.0:
         # Resize the frame
-        frame = cv2.resize(
-            frame, None, fx=resize, fy=resize, interpolation=cv2.INTER_LINEAR
-        )
+        frame = cv2.resize(frame, None, fx=resize, fy=resize, interpolation=cv2.INTER_LINEAR)
 
         # Scale pose coordinates
         pose = pose.copy()
@@ -692,7 +659,6 @@ def draw_pose_and_write(
     vwriter.write(image=frame)
 
 
-def _get_metadata(video_path: str, cap: cv2.VideoCapture, dlc_live: DLCLive):
 def _get_metadata(video_path: str, cap: cv2.VideoCapture, dlc_live: DLCLive):
     try:
         fourcc = decode_fourcc(cap.get(cv2.CAP_PROP_FOURCC))
@@ -730,9 +696,7 @@ def _get_metadata(video_path: str, cap: cv2.VideoCapture, dlc_live: DLCLive):
     return meta
 
 
-def save_poses_to_files(
-    video_path, save_dir, n_individuals, bodyparts, poses, timestamp
-):
+def save_poses_to_files(video_path, save_dir, n_individuals, bodyparts, poses, timestamp):
     """
     Saves the detected keypoint poses from the video to CSV and HDF5 files.
 
@@ -778,7 +742,6 @@ def save_poses_to_files(
     pose_df.to_csv(csv_save_path, index=False)
 
 
-
 def _create_poses_np_array(n_individuals: int, bodyparts: list, poses: list):
     # Create numpy array with poses:
     max_frame = max(p["frame"] for p in poses)
@@ -791,9 +754,7 @@ def _create_poses_np_array(n_individuals: int, bodyparts: list, poses: list):
         if pose.ndim == 2:
             pose = pose[np.newaxis, :, :]
         padded_pose = np.full(pose_target_shape, np.nan)
-        slices = tuple(
-            slice(0, min(pose.shape[i], pose_target_shape[i])) for i in range(3)
-        )
+        slices = tuple(slice(0, min(pose.shape[i], pose_target_shape[i])) for i in range(3))
         padded_pose[slices] = pose[slices]
         poses_array[frame] = padded_pose
 

dlclive/check_install/check_install.py

@@ -8,16 +8,15 @@
 import argparse
 import shutil
 import urllib
-import warnings
 import urllib.error
+import warnings
 from pathlib import Path
 
 from dlclibrary.dlcmodelzoo.modelzoo_download import download_huggingface_model
 
-from dlclive.utils import download_file
 from dlclive.benchmark import benchmark_videos
 from dlclive.engine import Engine
-from dlclive.utils import get_available_backends
+from dlclive.utils import download_file, get_available_backends
 
 MODEL_NAME = "superanimal_quadruped"
 SNAPSHOT_NAME = "snapshot-700000.pb"
@@ -46,9 +45,7 @@ def run_pytorch_test(video_file: str, display: bool = False):
         model_name=TORCH_MODEL,
     )
     if not TORCH_CONFIG["checkpoint"].exists():
-        raise FileNotFoundError(
-            f"Failed to export {TORCH_CONFIG['super_animal']} model"
-        )
+        raise FileNotFoundError(f"Failed to export {TORCH_CONFIG['super_animal']} model")
     if TORCH_CONFIG["checkpoint"].stat().st_size == 0:
         raise ValueError(f"Exported {TORCH_CONFIG['super_animal']} model is empty")
     benchmark_videos(
@@ -71,14 +68,10 @@ def run_tensorflow_test(video_file: str, display: bool = False):
     if Path(model_dir / SNAPSHOT_NAME).exists():
         print("Model already downloaded, using cached version")
     else:
-        print(
-            "Downloading superanimal_quadruped model from the DeepLabCut Model Zoo..."
-        )
+        print("Downloading superanimal_quadruped model from the DeepLabCut Model Zoo...")
         download_huggingface_model(MODEL_NAME, str(model_dir))
 
-    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}"
 
     benchmark_videos(
         model_path=str(model_dir),
@@ -131,7 +124,7 @@ def main():
             url_link = "https://raw.githubusercontent.com/DeepLabCut/DeepLabCut-live/master/check_install/dog_clip.avi"
             try:
                 download_file(url_link, video_file)
-            except (urllib.error.URLError, IOError) as e:
+            except (OSError, urllib.error.URLError) as e:
                 raise RuntimeError(f"Failed to download video file: {e}") from e
         else:
             print(f"Video file already exists at {video_file}, skipping download.")
@@ -155,9 +148,7 @@ def main():
                     any_backend_succeeded = True
                     backend_results["tensorflow"] = ("SUCCESS", None)
                 else:
-                    warnings.warn(
-                        f"Unrecognized backend {backend}, skipping...", UserWarning
-                    )
+                    warnings.warn(f"Unrecognized backend {backend}, skipping...", UserWarning, stacklevel=2)
             except Exception as e:
                 backend_name = (
                     "pytorch"
@@ -172,6 +163,7 @@ def main():
                     f"Error while running test for backend {backend}: {e}. "
                     "Continuing to test other available backends.",
                     UserWarning,
+                    stacklevel=2,
                 )
 
         print("\n---\nBackend test summary:")
@@ -184,9 +176,7 @@ def main():
                 print(f"{name.capitalize()} error:\n{error}\n")
 
         if not any_backend_succeeded and backend_failures:
-            failure_messages = "; ".join(
-                f"{b}: {exc}" for b, exc in backend_failures.items()
-            )
+            failure_messages = "; ".join(f"{b}: {exc}" for b, exc in backend_failures.items())
             raise RuntimeError(f"All backend tests failed. Details: {failure_messages}")
 
     finally:
@@ -197,7 +187,7 @@ def main():
                 shutil.rmtree(tmp_dir)
         except PermissionError:
             warnings.warn(
-                f"Could not delete temporary directory {str(tmp_dir)} due to a permissions error."
+                f"Could not delete temporary directory {str(tmp_dir)} due to a permissions error.", stacklevel=2
             )
 
 
@@ -207,7 +197,8 @@ def main():
     print(f"Available backends: {[b.value for b in available_backends]}")
     if len(available_backends) == 0:
         raise NotImplementedError(
-            "Neither TensorFlow nor PyTorch is installed. Please install at least one of these frameworks to run the installation test."
+            "Neither TensorFlow nor PyTorch is installed. "
+            "Please install at least one of these frameworks to run the installation test."
         )
 
     main()