diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index c0b0d9a..9bdf762 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -4,9 +4,9 @@ name: Python package on: push: - branches: [ master ] + branches: [ main ] pull_request: - branches: [ master ] + branches: [ main ] jobs: diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 77e0b23..9b4ffdf 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -2,9 +2,8 @@ name: Run Tests on: push: - branches: [master] + branches: [main] pull_request: - branches: [master] jobs: build: @@ -31,6 +30,14 @@ jobs: uses: actions/setup-python@v4 with: python-version: ${{ matrix.python-version }} + + - name: Set up Python + uses: conda-incubator/setup-miniconda@v3 + with: + channels: conda-forge,defaults + channel-priority: strict + python-version: ${{ matrix.python-version }} + - name: Install ffmpeg run: | if [ "$RUNNER_OS" == "Linux" ]; then @@ -42,8 +49,27 @@ jobs: choco install ffmpeg fi shell: bash + - name: Install and test + shell: bash -el {0} # Important: activates the conda environment run: | - python -m pip install --upgrade pip wheel poetry - python -m poetry install --extras "tf" - python -m poetry run dlc-live-test --nodisplay + conda install pytables==3.8.0 "numpy<2" + + - name: Install dependencies via Conda + shell: bash -el {0} + run: conda install -y "numpy>=1.26,<2.0" + + - name: Install Poetry + run: pip install --upgrade pip wheel poetry + + - name: Regenerate Poetry lock + run: poetry lock --no-cache + + - name: Install project dependencies + run: poetry install --with dev --extras "tf" --extras "pytorch" + + - name: Run DLC Live Tests + run: poetry run dlc-live-test --nodisplay + + - name: Run Functional Benchmark Test + run: poetry run pytest tests/test_benchmark_script.py diff --git a/.gitignore b/.gitignore index 7fa3a49..2ce1454 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,8 @@ benchmarking/results* **DS_Store* *vscode* +**/__MACOSX/ + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/README.md b/README.md index 5ae7593..3ddbedb 100644 --- a/README.md +++ b/README.md @@ -73,10 +73,13 @@ same version (i.e., export a PyTorch model, then install PyTorch, export with TF then use TF1.13 with DlC-Live; export with TF2.3, then use TF2.3 with DLC-live). - available on pypi as: `pip install deeplabcut-live` + -Note, you can then test your installation by running: +Note, you can then test your installation by installing poetry (`pip install poetry`), then running: -`dlc-live-test` +```python +poetry run dlc-live-test +``` If installed properly, this script will i) create a temporary folder ii) download the full_dog model from the [DeepLabCut Model Zoo]( diff --git a/benchmarking/run_dlclive_benchmark.py b/benchmarking/run_dlclive_benchmark.py index 1ea14e4..859843b 100644 --- a/benchmarking/run_dlclive_benchmark.py +++ b/benchmarking/run_dlclive_benchmark.py @@ -12,6 +12,7 @@ import glob from dlclive import benchmark_videos, download_benchmarking_data +from dlclive.engine import Engine datafolder = os.path.join( pathlib.Path(__file__).parent.absolute(), "Data-DLC-live-benchmark" @@ -36,8 +37,22 @@ if not os.path.isdir(out_dir): os.mkdir(out_dir) -for m in dog_models: - benchmark_videos(m, dog_video, output=out_dir, n_frames=n_frames, pixels=pixels) - -for m in mouse_models: - benchmark_videos(m, mouse_video, output=out_dir, n_frames=n_frames, pixels=pixels) +for model_path in dog_models: + benchmark_videos( + model_path=model_path, + model_type="base" if Engine.from_model_path(model_path) == Engine.TENSORFLOW else "pytorch", + video_path=dog_video, + output=out_dir, + n_frames=n_frames, + pixels=pixels + ) + +for model_path in mouse_models: + benchmark_videos( + model_path=model_path, + model_type="base" if Engine.from_model_path(model_path) == Engine.TENSORFLOW else "pytorch", + video_path=mouse_video, + output=out_dir, + n_frames=n_frames, + pixels=pixels + ) diff --git a/dlclive/__init__.py b/dlclive/__init__.py index 71a89d9..b1f5774 100644 --- a/dlclive/__init__.py +++ b/dlclive/__init__.py @@ -9,3 +9,4 @@ from dlclive.dlclive import DLCLive from dlclive.processor.processor import Processor from dlclive.version import VERSION, __version__ +from dlclive.benchmark import benchmark_videos, download_benchmarking_data \ No newline at end of file diff --git a/dlclive/benchmark.py b/dlclive/benchmark.py index 2f0f2af..1a41ca9 100644 --- a/dlclive/benchmark.py +++ b/dlclive/benchmark.py @@ -16,28 +16,198 @@ import colorcet as cc import cv2 import numpy as np -import torch +import pickle from PIL import ImageColor from pip._internal.operations import freeze +import torch +from tqdm import tqdm -try: - import pandas as pd +from dlclive import DLCLive +from dlclive import VERSION +from dlclive import __file__ as dlcfile +from dlclive.engine import Engine +from dlclive.utils import decode_fourcc - has_pandas = True -except ModuleNotFoundError as err: - has_pandas = False -try: - from tqdm import tqdm +def download_benchmarking_data( + target_dir=".", + url="https://huggingface.co/datasets/mwmathis/DLCspeed_benchmarking/resolve/main/Data-DLC-live-benchmark.zip", +): + """ + Downloads and extracts DeepLabCut-Live benchmarking data (videos & DLC models). + """ + import os + import urllib.request + import zipfile - has_tqdm = True -except ModuleNotFoundError as err: - has_tqdm = False + # Avoid nested folder issue + if os.path.basename(os.path.normpath(target_dir)) == "Data-DLC-live-benchmark": + target_dir = os.path.dirname(os.path.normpath(target_dir)) + os.makedirs(target_dir, exist_ok=True) # Ensure target directory exists + zip_path = os.path.join(target_dir, "Data-DLC-live-benchmark.zip") + + if os.path.exists(zip_path): + 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) + pbar.close() + + print(f"Extracting {zip_path} to {target_dir} ...") + with zipfile.ZipFile(zip_path, 'r') as zip_ref: + zip_ref.extractall(target_dir) -from dlclive import DLCLive -from dlclive.utils import decode_fourcc -from dlclive.version import VERSION + +def benchmark_videos( + model_path, + model_type, + video_path, + output=None, + n_frames=1000, + tf_config=None, + resize=None, + pixels=None, + cropping=None, + dynamic=(False, 0.5, 10), + print_rate=False, + display=False, + pcutoff=0.5, + display_radius=3, + cmap="bmy", + save_poses=False, + save_video=False, +): + """Analyze videos using DeepLabCut-live exported models. + Analyze multiple videos and/or multiple options for the size of the video + by specifying a resizing factor or the number of pixels to use in the image (keeping aspect ratio constant). + Options to record inference times (to examine inference speed), + display keypoints to visually check the accuracy, + or save poses to an hdf5 file as in :function:`deeplabcut.benchmark_videos` and + create a labeled video as in :function:`deeplabcut.create_labeled_video`. + + Parameters + ---------- + model_path : str + path to exported DeepLabCut model + model_type: string, optional + Which model to use. For the PyTorch engine, options are [`pytorch`]. For the + TensorFlow engine, options are [`base`, `tensorrt`, `lite`]. + video_path : str or list + path to video file or list of paths to video files + output : str + path to directory to save results + tf_config : :class:`tensorflow.ConfigProto` + tensorflow session configuration + resize : int, optional + resize factor. Can only use one of resize or pixels. If both are provided, will use pixels. by default None + pixels : int, optional + downsize image to this number of pixels, maintaining aspect ratio. Can only use one of resize or pixels. If both are provided, will use pixels. by default None + cropping : list of int + cropping parameters in pixel number: [x1, x2, y1, y2] + dynamic: triple containing (state, detectiontreshold, margin) + If the state is true, then dynamic cropping will be performed. That means that if an object is detected (i.e. any body part > detectiontreshold), + then object boundaries are computed according to the smallest/largest x position and smallest/largest y position of all body parts. This window is + expanded by the margin and from then on only the posture within this crop is analyzed (until the object is lost, i.e. `, by default "bmy" + save_poses : bool, optional + flag to save poses to an hdf5 file. If True, operates similar to :function:`DeepLabCut.benchmark_videos`, by default False + save_video : bool, optional + flag to save a labeled video. If True, operates similar to :function:`DeepLabCut.create_labeled_video`, by default False + + Example + ------- + Return a vector of inference times for 10000 frames on one video or two videos: + dlclive.benchmark_videos('/my/exported/model', 'my_video.avi', n_frames=10000) + dlclive.benchmark_videos('/my/exported/model', ['my_video1.avi', 'my_video2.avi'], n_frames=10000) + + Return a vector of inference times, testing full size and resizing images to half the width and height for inference, for two videos + dlclive.benchmark_videos('/my/exported/model', ['my_video1.avi', 'my_video2.avi'], n_frames=10000, resize=[1.0, 0.5]) + + Display keypoints to check the accuracy of an exported model + dlclive.benchmark_videos('/my/exported/model', 'my_video.avi', display=True) + + Analyze a video (save poses to hdf5) and create a labeled video, similar to :function:`DeepLabCut.benchmark_videos` and :function:`create_labeled_video` + dlclive.benchmark_videos('/my/exported/model', 'my_video.avi', save_poses=True, save_video=True) + """ + # convert video_paths to list + video_path = video_path if type(video_path) is list else [video_path] + + # fix resize + if pixels: + pixels = pixels if type(pixels) is list else [pixels] + resize = [None for p in pixels] + elif resize: + resize = resize if type(resize) is list else [resize] + pixels = [None for r in resize] + else: + resize = [None] + pixels = [None] + + # loop over videos + for video in video_path: + # initialize full inference times + inf_times = [] + im_size_out = [] + + for i in range(len(resize)): + print(f"\nRun {i+1} / {len(resize)}\n") + + this_inf_times, this_im_size, meta = benchmark( + model_path=model_path, + model_type=model_type, + video_path=video, + tf_config=tf_config, + resize=resize[i], + pixels=pixels[i], + cropping=cropping, + dynamic=dynamic, + n_frames=n_frames, + print_rate=print_rate, + display=display, + pcutoff=pcutoff, + display_radius=display_radius, + cmap=cmap, + save_poses=save_poses, + save_video=save_video, + save_dir=output, + ) + + inf_times.append(this_inf_times) + im_size_out.append(this_im_size) + + inf_times = np.array(inf_times) + im_size_out = np.array(im_size_out) + + # save results + if output is not None: + sys_info = get_system_info() + save_inf_times( + sys_info, + inf_times, + im_size_out, + model=os.path.basename(model_path), + meta=meta, + output=output, + ) def get_system_info() -> dict: @@ -106,344 +276,132 @@ def get_system_info() -> dict: } -def benchmark( - path: str | Path, - video_path: str | Path, - single_animal: bool = True, - resize: float | None = None, - pixels: int | None = None, - cropping: list[int] = None, - dynamic: tuple[bool, float, int] = (False, 0.5, 10), - n_frames: int = 1000, - print_rate: bool = False, - display: bool = False, - pcutoff: float = 0.0, - max_detections: int = 10, - display_radius: int = 3, - cmap: str = "bmy", - save_poses: bool = False, - save_video: bool = False, - output: str | Path | None = None, -) -> tuple[np.ndarray, tuple, dict]: - """Analyze DeepLabCut-live exported model on a video: - - Calculate inference time, display keypoints, or get poses/create a labeled video. +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` + Parameters ---------- - path : str - path to exported DeepLabCut model - video_path : str - path to video file - single_animal: bool - to make code behave like DLCLive for tensorflow models - resize : int, optional - Resize factor. Can only use one of resize or pixels. If both are provided, will - use pixels. by default None - pixels : int, optional - Downsize image to this number of pixels, maintaining aspect ratio. Can only use - one of resize or pixels. If both are provided, will use pixels. by default None - cropping : list of int - cropping parameters in pixel number: [x1, x2, y1, y2] - dynamic: triple containing (state, detectiontreshold, margin) - If the state is true, then dynamic cropping will be performed. That means that - if an object is detected (i.e. any body part > detectiontreshold), then object - boundaries are computed according to the smallest/largest x position and - smallest/largest y position of all body parts. This window is expanded by the - margin and from then on only the posture within this crop is analyzed (until the - object is lost, i.e. < detectiontreshold). The current position is utilized for - updating the crop window for the next frame (this is why the margin is important - and should be set large enough given the movement of the animal) - n_frames : int, optional - number of frames to run inference on, by default 1000 - print_rate : bool, optional - 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 - likelihood threshold to display keypoints - max_detections: int - for top-down models, the maximum number of individuals to detect in a frame - display_radius : int, optional - size (radius in pixels) of keypoint to display - cmap : str, optional - a string indicating the :package:`colorcet` colormap, `options here - `, by default "bmy" - save_poses : bool, optional - flag to save poses to an hdf5 file. If True, operates similar to - :function:`DeepLabCut.benchmark_videos`, by default False - save_video : bool, optional - flag to save a labeled video. If True, operates similar to - :function:`DeepLabCut.create_labeled_video`, by default False + sys_info : tuple + system information generated by :func:`get_system_info` + inf_times : :class:`numpy.ndarray` + array of inference times generated by :func:`benchmark` + im_size : tuple or :class:`numpy.ndarray` + image size (width, height) for each benchmark run. If an array, each row corresponds to a row in inf_times + model: str, optional + name of model + meta : dict, optional + metadata returned by :func:`benchmark` output : str, optional - path to directory to save pose and/or video file. If not specified, will use - the directory of video_path, by default None + path to directory to save data. If None, uses pwd, by default None Returns ------- - :class:`numpy.ndarray` - vector of inference times - tuple - (image width, image height) - dict - metadata for video - - Example - ------- - Return a vector of inference times for 10000 frames: - dlclive.benchmark('/my/exported/model', 'my_video.avi', n_frames=10000) - - Return a vector of inference times, resizing images to half the width and height for inference - dlclive.benchmark('/my/exported/model', 'my_video.avi', n_frames=10000, resize=0.5) - - Display keypoints to check the accuracy of an exported model - dlclive.benchmark('/my/exported/model', 'my_video.avi', display=True) - - Analyze a video (save poses to hdf5) and create a labeled video, similar to :function:`DeepLabCut.benchmark_videos` and :function:`create_labeled_video` - dlclive.benchmark('/my/exported/model', 'my_video.avi', save_poses=True, save_video=True) + bool + flag indicating successful save """ - path = Path(path) - video_path = Path(video_path) - if not video_path.exists(): - raise ValueError(f"Could not find video: {video_path}: check that it exists!") - if output is None: - output = video_path.parent - else: - output = Path(output) - output.mkdir(exist_ok=True, parents=True) + output = output if output is not None else os.getcwd() + model_type = None + if model is not None: + if "resnet" in model: + model_type = "resnet" + elif "mobilenet" in model: + model_type = "mobilenet" + else: + model_type = None - # load video - cap = cv2.VideoCapture(str(video_path)) - ret, frame = cap.read() - n_frames = ( - n_frames - if (n_frames > 0) and (n_frames < cap.get(cv2.CAP_PROP_FRAME_COUNT) - 1) - else (cap.get(cv2.CAP_PROP_FRAME_COUNT) - 1) - ) - n_frames = int(n_frames) - im_size = ( - int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), - int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), + fn_ind = 0 + base_name = ( + f"benchmark_{sys_info['host_name']}_{sys_info['device_type']}_{fn_ind}.pickle" ) - - # get resize factor - if pixels is not None: - resize = np.sqrt(pixels / (im_size[0] * im_size[1])) - - if resize is not None: - im_size = (int(im_size[0] * resize), int(im_size[1] * resize)) - - # create video writer - if save_video: - colors = None - out_vid_file = output / f"{video_path.stem}_DLCLIVE_LABELED.avi" - fourcc = cv2.VideoWriter_fourcc(*"DIVX") - fps = cap.get(cv2.CAP_PROP_FPS) - print(out_vid_file) - print(fourcc) - print(fps) - print(im_size) - vid_writer = cv2.VideoWriter(str(out_vid_file), fourcc, fps, im_size) - - # initialize DLCLive and perform inference - inf_times = np.zeros(n_frames) - poses = [] - - live = DLCLive( - model_path=path, - single_animal=single_animal, - resize=resize, - cropping=cropping, - dynamic=dynamic, - display=display, - max_detections=max_detections, - pcutoff=pcutoff, - display_radius=display_radius, - display_cmap=cmap, + out_file = os.path.normpath(f"{output}/{base_name}") + while os.path.isfile(out_file): + fn_ind += 1 + base_name = f"benchmark_{sys_info['host_name']}_{sys_info['device_type']}_{fn_ind}.pickle" + out_file = os.path.normpath(f"{output}/{base_name}") + + # summary stats (mean inference time & standard error of mean) + stats = zip( + np.mean(inf_times, 1), + np.std(inf_times, 1) * 1.0 / np.sqrt(np.shape(inf_times)[1]), ) - poses.append(live.init_inference(frame)) - - iterator = range(n_frames) - if print_rate or display: - iterator = tqdm(iterator) - - for i in iterator: - ret, frame = cap.read() - if not ret: - warnings.warn( - f"Did not complete {n_frames:d} frames. There probably were not enough " - f"frames in the video {video_path}." - ) - break - - start_pose = time.time() - poses.append(live.get_pose(frame)) - inf_times[i] = time.time() - start_pose - if save_video: - this_pose = poses[-1] - - if single_animal: - # expand individual dimension - this_pose = this_pose[None] - - num_idv, num_bpt = this_pose.shape[:2] - num_colors = num_bpt - - if colors is None: - all_colors = getattr(cc, cmap) - colors = [ - ImageColor.getcolor(c, "RGB")[::-1] - for c in all_colors[:: int(len(all_colors) / num_colors)] - ] - - for j in range(num_idv): - for k in range(num_bpt): - color_idx = k - if this_pose[j, k, 2] > pcutoff: - x = int(this_pose[j, k, 0]) - y = int(this_pose[j, k, 1]) - frame = cv2.circle( - frame, - (x, y), - display_radius, - colors[color_idx], - thickness=-1, - ) - - if resize is not None: - frame = cv2.resize(frame, im_size) - vid_writer.write(frame) - - if print_rate: - print(f"pose rate = {int(1 / inf_times[i]):d}") - - if print_rate: - print(f"mean pose rate = {int(np.mean(1 / inf_times)):d}") - - # gather video and test parameterization - # dont want to fail here so gracefully failing on exception -- - # eg. some packages of cv2 don't have CAP_PROP_CODEC_PIXEL_FORMAT - try: - fourcc = decode_fourcc(cap.get(cv2.CAP_PROP_FOURCC)) - except: - fourcc = "" - - try: - fps = round(cap.get(cv2.CAP_PROP_FPS)) - except Exception: - fps = None - - try: - pix_fmt = decode_fourcc(cap.get(cv2.CAP_PROP_CODEC_PIXEL_FORMAT)) - except Exception: - pix_fmt = "" - - try: - frame_count = round(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - except Exception: - frame_count = None - - try: - orig_im_size = ( - round(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), - round(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), - ) - except Exception: - orig_im_size = None - - meta = { - "video_path": video_path, - "video_codec": fourcc, - "video_pixel_format": pix_fmt, - "video_fps": fps, - "video_total_frames": frame_count, - "original_frame_size": orig_im_size, - "dlclive_params": live.parameterization, + data = { + "model": model, + "model_type": model_type, + "im_size": im_size, + "inference_times": inf_times, + "stats": stats, } - # close video - cap.release() - if save_video: - vid_writer.release() - - if save_poses: - bodyparts = live.cfg["metadata"]["bodyparts"] - max_idv = np.max([p.shape[0] for p in poses]) - - poses_array = -np.ones((len(poses), max_idv, len(bodyparts), 3)) - for i, p in enumerate(poses): - num_det = len(p) - poses_array[i, :num_det] = p - poses = poses_array - - num_frames, num_idv, num_bpts = poses.shape[:3] - individuals = [f"individual-{i}" for i in range(num_idv)] - - if has_pandas: - poses = poses.reshape((num_frames, num_idv * num_bpts * 3)) - col_index = pd.MultiIndex.from_product( - [individuals, bodyparts, ["x", "y", "likelihood"]], - names=["individual", "bodyparts", "coords"], - ) - pose_df = pd.DataFrame(poses, columns=col_index) - - out_dlc_file = output / (video_path.stem + "_DLCLIVE_POSES.h5") - try: - pose_df.to_hdf(out_dlc_file, key="df_with_missing", mode="w") - except ImportError as err: - print( - "Cannot export predictions to H5 file. Install ``pytables`` extra " - f"to export to HDF: {err}" - ) - out_csv = Path(out_dlc_file).with_suffix(".csv") - pose_df.to_csv(out_csv) - - else: - warnings.warn( - "Could not find installation of pandas; saving poses as a numpy array " - "with the dimensions (n_frames, n_keypoints, [x, y, likelihood])." - ) - np.save(str(output / (video_path.stem + "_DLCLIVE_POSES.npy")), poses) + data.update(sys_info) + if meta: + data.update(meta) - return inf_times, im_size, meta + os.makedirs(os.path.normpath(output), exist_ok=True) + pickle.dump(data, open(out_file, "wb")) + return True -def benchmark_videos( - video_path: str, +def benchmark( model_path: str, model_type: str, - device: str, + video_path: str, + tf_config=None, + device: str | None = None, + resize: float | None = None, + pixels: int | None = None, + single_animal: bool = True, + cropping=None, + dynamic=(False, 0.5, 10), + n_frames=1000, + print_rate=False, precision: str = "FP32", display=True, pcutoff=0.5, - display_radius=5, - resize=None, - cropping=None, # Adding cropping to the function parameters - dynamic=(False, 0.5, 10), - save_poses=False, - save_dir="model_predictions", - draw_keypoint_names=False, + display_radius=3, cmap="bmy", - get_sys_info=True, + save_dir=None, + save_poses=False, save_video=False, + draw_keypoint_names=False, ): """ - Analyzes a video to track keypoints using a DeepLabCut model, and optionally saves - the keypoint data and the labeled video. + Analyzes a video to track keypoints using a DeepLabCut model, and optionally saves the keypoint data and the labeled video. Parameters ---------- - video_path : str - Path to the video file to be analyzed. model_path : str Path to the DeepLabCut model. model_type : str - Type of the model (e.g., 'onnx'). + Which model to use. For the PyTorch engine, options are [`pytorch`]. For the + TensorFlow engine, options are [`base`, `tensorrt`, `lite`]. + video_path : str + Path to the video file to be analyzed. + TensorFlow engine, options are [`base`, `tensorrt`, `lite`]. + tf_config : :class:`tensorflow.ConfigProto` + Tensorflow only. Tensorflow session configuration device : str - Device to run the model on ('cpu' or 'cuda'). + Pytorch only. Device to run the model on ('cpu' or 'cuda'). + resize : float or None, optional + Resize dimensions for video frames. e.g. if resize = 0.5, the video will be processed in half the original size. If None, no resizing is applied. + pixels : int, optional + downsize image to this number of pixels, maintaining aspect ratio. + Can only use one of resize or pixels. If both are provided, will use pixels. + single_animal: bool, optional, default=True + Whether the video contains only one animal (True) or multiple animals (False). + cropping : list of int or None, optional + Cropping parameters [x1, x2, y1, y2] in pixels. If None, no cropping is applied. + dynamic : tuple, optional, default=(False, 0.5, 10) (True/false), p cutoff, margin) + Parameters for dynamic cropping. If the state is true, then dynamic cropping will be performed. That means that if an object is detected (i.e. any body part > detectiontreshold), then object boundaries are computed according to the smallest/largest x position and smallest/largest y position of all body parts. This window is expanded by the margin and from then on only the posture within this crop is analyzed (until the object is lost, i.e. detectiontreshold), then object boundaries are computed according to the - smallest/largest x position and smallest/largest y position of all body parts. - This window is expanded by the margin and from then on only the posture within - this crop is analyzed (until the object is lost, i.e. 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() if not ret: + warnings.warn( + ( + "Did not complete {:d} frames. " + "There probably were not enough frames in the video {}." + ).format(n_frames, video_path) + ) break - # if frame_index == 0: - # pose = dlc_live.init_inference(frame) # load DLC model - try: - # pose = dlc_live.get_pose(frame) - if frame_index == 0: - # TODO trying to fix issues with dynamic cropping jumping back and forth - # between dyanmic cropped and original image - # dlc_live.dynamic = (False, dynamic[1], dynamic[2]) - pose, inf_time = dlc_live.init_inference(frame) # load DLC model - else: - # dlc_live.dynamic = dynamic - pose, inf_time = dlc_live.get_pose(frame) - except Exception as e: - print(f"Error analyzing frame {frame_index}: {e}") - continue + start_time = time.perf_counter() + if frame_index == 0: + pose = dlc_live.init_inference(frame) # Loads model + else: + pose = dlc_live.get_pose(frame) + + inf_time = time.perf_counter() - start_time poses.append({"frame": frame_index, "pose": pose}) times.append(inf_time) + if print_rate: + print("Inference rate = {:.3f} FPS".format(1 / inf_time), end="\r", flush=True) + if save_video: - # Visualize keypoints - this_pose = pose["poses"][0][0] - for j in range(this_pose.shape[0]): - if this_pose[j, 2] > pcutoff: - x, y = map(int, this_pose[j, :2]) - cv2.circle( - frame, - center=(x, y), - radius=display_radius, - color=colors[j], - thickness=-1, - ) + draw_pose_and_write( + frame=frame, + pose=pose, + resize=resize, + colors=colors, + bodyparts=bodyparts, + pcutoff=pcutoff, + display_radius=display_radius, + draw_keypoint_names=draw_keypoint_names, + vwriter=vwriter + ) - if draw_keypoint_names: - cv2.putText( - frame, - text=bodyparts[j], - org=(x + 10, y), - fontFace=cv2.FONT_HERSHEY_SIMPLEX, - fontScale=0.5, - color=colors[j], - thickness=1, - lineType=cv2.LINE_AA, - ) - - vwriter.write(image=frame) frame_index += 1 + if print_rate: + print("Mean inference rate: {:.3f} FPS".format(np.mean(1 / np.array(times)[1:]))) + + metadata = _get_metadata( + video_path=video_path, + cap=cap, + dlc_live=dlc_live + ) + cap.release() + + dlc_live.close() + if save_video: vwriter.release() - if get_sys_info: - print(get_system_info()) - if save_poses: - save_poses_to_files(video_path, save_dir, bodyparts, poses, timestamp=timestamp) + if engine == Engine.PYTORCH: + individuals = dlc_live.read_config()["metadata"].get("individuals", []) + else: + individuals = [] + n_individuals = len(individuals) or 1 + save_poses_to_files(video_path, save_dir, n_individuals, bodyparts, poses, timestamp=timestamp) + + return times, im_size, metadata + - return poses, times +def setup_video_writer( + 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) + colors = [ + ImageColor.getrgb(color) + for color in cmap_colors[:: int(len(cmap_colors) / num_keypoints)] + ] + + # 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" + + # Get video writer setup + fourcc = cv2.VideoWriter_fourcc(*"mp4v") + vwriter = cv2.VideoWriter( + filename=output_video_path, + fourcc=fourcc, + fps=fps, + frameSize=frame_size, + ) + + return colors, vwriter + +def draw_pose_and_write( + frame: np.ndarray, + pose: np.ndarray, + resize: float, + colors: list[tuple[int, int, int]], + bodyparts: list[str], + pcutoff: float, + display_radius: int, + draw_keypoint_names: bool, + vwriter: cv2.VideoWriter, +): + if len(pose.shape) == 2: + pose = pose[None] + + 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) + + # Scale pose coordinates + pose = pose.copy() + pose[..., :2] *= resize + + # Visualize keypoints + for i in range(pose.shape[0]): + for j in range(pose.shape[1]): + if pose[i, j, 2] > pcutoff: + x, y = map(int, pose[i, j, :2]) + cv2.circle( + frame, + center=(x, y), + radius=display_radius, + color=colors[j], + thickness=-1, + ) + + if draw_keypoint_names: + cv2.putText( + frame, + text=bodyparts[j], + org=(x + 10, y), + fontFace=cv2.FONT_HERSHEY_SIMPLEX, + fontScale=0.5, + color=colors[j], + thickness=1, + lineType=cv2.LINE_AA, + ) -def save_poses_to_files(video_path, save_dir, bodyparts, poses, timestamp): + vwriter.write(image=frame) + + +def _get_metadata( + video_path: str, + cap: cv2.VideoCapture, + dlc_live: DLCLive +): + try: + fourcc = decode_fourcc(cap.get(cv2.CAP_PROP_FOURCC)) + except: + fourcc = "" + try: + fps = round(cap.get(cv2.CAP_PROP_FPS)) + except: + fps = None + try: + pix_fmt = decode_fourcc(cap.get(cv2.CAP_PROP_CODEC_PIXEL_FORMAT)) + except: + pix_fmt = "" + try: + frame_count = round(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + except: + frame_count = None + try: + orig_im_size = ( + round(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), + round(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), + ) + except: + orig_im_size = None + + meta = { + "video_path": video_path, + "video_codec": fourcc, + "video_pixel_format": pix_fmt, + "video_fps": fps, + "video_total_frames": frame_count, + "original_frame_size": orig_im_size, + "dlclive_params": dlc_live.parameterization, + } + return meta + + +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. @@ -623,6 +698,8 @@ def save_poses_to_files(video_path, save_dir, bodyparts, poses, timestamp): Path to the analyzed video file. save_dir : str Directory where the pose data files will be saved. + n_individuals: int + Number of individuals bodyparts : list of str List of body part names corresponding to the keypoints. poses : list of dict @@ -632,27 +709,48 @@ def save_poses_to_files(video_path, save_dir, bodyparts, poses, timestamp): ------- None """ + import pandas as pd + + base_filename = Path(video_path).stem + save_dir = Path(save_dir) + h5_save_path = save_dir / f"{base_filename}_poses_{timestamp}.h5" + csv_save_path = save_dir / f"{base_filename}_poses_{timestamp}.csv" + + poses_array = _create_poses_np_array(n_individuals, bodyparts, poses) + flattened_poses = poses_array.reshape(poses_array.shape[0], -1) + + if n_individuals == 1: + pdindex = pd.MultiIndex.from_product( + [bodyparts, ["x", "y", "likelihood"]], names=["bodyparts", "coords"] + ) + else: + individuals = [f"individual_{i}" for i in range(n_individuals)] + pdindex = pd.MultiIndex.from_product( + [individuals, bodyparts, ["x", "y", "likelihood"]], names=["individuals", "bodyparts", "coords"] + ) + + pose_df = pd.DataFrame(flattened_poses, columns=pdindex) + + pose_df.to_hdf(h5_save_path, key="df_with_missing", mode="w") + 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) + pose_target_shape = (n_individuals, len(bodyparts), 3) + poses_array = np.full((max_frame + 1, *pose_target_shape), np.nan) + + for item in poses: + frame = item["frame"] + pose = item["pose"] + 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)) + padded_pose[slices] = pose[slices] + poses_array[frame] = padded_pose - base_filename = os.path.splitext(os.path.basename(video_path))[0] - csv_save_path = os.path.join(save_dir, f"{base_filename}_poses_{timestamp}.csv") - h5_save_path = os.path.join(save_dir, f"{base_filename}_poses_{timestamp}.h5") - - # Save to CSV - with open(csv_save_path, mode="w", newline="") as file: - writer = csv.writer(file) - header = ["frame"] + [ - f"{bp}_{axis}" for bp in bodyparts for axis in ["x", "y", "confidence"] - ] - writer.writerow(header) - for entry in poses: - frame_num = entry["frame"] - pose = entry["pose"]["poses"][0][0] - row = [frame_num] + [ - item.item() if isinstance(item, torch.Tensor) else item - for kp in pose - for item in kp - ] - writer.writerow(row) + return poses_array import argparse diff --git a/dlclive/benchmark_pytorch.py b/dlclive/benchmark_pytorch.py deleted file mode 100644 index caf855a..0000000 --- a/dlclive/benchmark_pytorch.py +++ /dev/null @@ -1,516 +0,0 @@ -import csv -import platform -import subprocess -import sys -import time - -import colorcet as cc -import cv2 -import h5py -import numpy as np -from pathlib import Path -from PIL import ImageColor -from pip._internal.operations import freeze -import torch -from tqdm import tqdm - -# torch import needs to switch order with "from pip._internal.operations import freeze" because of crash -# see https://github.com/pytorch/pytorch/issues/140914 - -from dlclive import DLCLive -from dlclive.version import VERSION - - -def get_system_info() -> dict: - """ - Returns a summary of system information relevant to running benchmarking. - - Returns - ------- - dict - A dictionary containing the following system information: - - host_name (str): Name of the machine. - - op_sys (str): Operating system. - - python (str): Path to the Python executable, indicating the conda/virtual environment in use. - - device_type (str): Type of device used ('GPU' or 'CPU'). - - device (list): List containing the name of the GPU or CPU brand. - - freeze (list): List of installed Python packages with their versions. - - python_version (str): Version of Python in use. - - git_hash (str or None): If installed from git repository, hash of HEAD commit. - - dlclive_version (str): Version of the DLCLive package. - """ - - # Get OS and host name - op_sys = platform.platform() - host_name = platform.node().replace(" ", "") - - # Get Python executable path - if platform.system() == "Windows": - host_python = sys.executable.split(os.path.sep)[-2] - else: - host_python = sys.executable.split(os.path.sep)[-3] - - # Try to get git hash if possible - git_hash = None - dlc_basedir = os.path.dirname(os.path.dirname(__file__)) - try: - git_hash = ( - subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=dlc_basedir) - .decode("utf-8") - .strip() - ) - except subprocess.CalledProcessError: - # Not installed from git repo, e.g., pypi - pass - - # Get device info (GPU or CPU) - if torch.cuda.is_available(): - dev_type = "GPU" - dev = [torch.cuda.get_device_name(torch.cuda.current_device())] - else: - from cpuinfo import get_cpu_info - - dev_type = "CPU" - dev = [get_cpu_info()["brand_raw"]] - - return { - "host_name": host_name, - "op_sys": op_sys, - "python": host_python, - "device_type": dev_type, - "device": dev, - "freeze": list(freeze.freeze()), - "python_version": sys.version, - "git_hash": git_hash, - "dlclive_version": VERSION, - } - - -def benchmark( - video_path: str, - model_path: str, - model_type: str, - device: str, - single_animal: bool, - save_dir=None, - n_frames=1000, - precision: str = "FP32", - display=True, - pcutoff=0.5, - display_radius=3, - resize=None, - cropping=None, # Adding cropping to the function parameters - dynamic=(False, 0.5, 10), - save_poses=False, - draw_keypoint_names=False, - cmap="bmy", - get_sys_info=True, - save_video=False, -): - """ - Analyzes a video to track keypoints using a DeepLabCut model, and optionally saves the keypoint data and the labeled video. - - Parameters - ---------- - video_path : str - Path to the video file to be analyzed. - model_path : str - Path to the DeepLabCut model. - model_type : str - Type of the model (e.g., 'onnx'). - device : str - Device to run the model on ('cpu' or 'cuda'). - single_animal: bool - Whether the video contains only one animal (True) or multiple animals (False). - save_dir : str, optional - Directory to save output data and labeled video. - If not specified, will use the directory of video_path, by default None - n_frames : int, optional - Number of frames to run inference on, by default 1000 - precision : str, optional, default='FP32' - Precision type for the model ('FP32' or 'FP16'). - display : bool, optional, default=True - Whether to display frame with labelled key points. - pcutoff : float, optional, default=0.5 - Probability cutoff below which keypoints are not visualized. - display_radius : int, optional, default=5 - Radius of circles drawn for keypoints on video frames. - resize : tuple of int (width, height) or None, optional - Resize dimensions for video frames. e.g. if resize = 0.5, the video will be processed in half the original size. If None, no resizing is applied. - cropping : list of int or None, optional - Cropping parameters [x1, x2, y1, y2] in pixels. If None, no cropping is applied. - dynamic : tuple, optional, default=(False, 0.5, 10) (True/false), p cutoff, margin) - Parameters for dynamic cropping. If the state is true, then dynamic cropping will be performed. That means that if an object is detected (i.e. any body part > detectiontreshold), then object boundaries are computed according to the smallest/largest x position and smallest/largest y position of all body parts. This window is expanded by the margin and from then on only the posture within this crop is analyzed (until the object is lost, i.e. 0) and n_frames < total_n_frames - else total_n_frames - ) - iterator = range(n_frames) if display else tqdm(range(n_frames)) - for i in iterator: - ret, frame = cap.read() - if not ret: - break - - start_time = time.perf_counter() - if frame_index == 0: - pose = dlc_live.init_inference(frame) # Loads model - else: - pose = dlc_live.get_pose(frame) - - inf_time = time.perf_counter() - start_time - poses.append({"frame": frame_index, "pose": pose}) - times.append(inf_time) - - if save_video: - draw_pose_and_write( - frame=frame, - pose=pose, - colors=colors, - bodyparts=bodyparts, - pcutoff=pcutoff, - display_radius=display_radius, - draw_keypoint_names=draw_keypoint_names, - vwriter=vwriter - ) - - frame_index += 1 - - cap.release() - if save_video: - vwriter.release() - - if get_sys_info: - print(get_system_info()) - - if save_poses: - individuals = dlc_live.read_config()["metadata"].get("individuals", []) - n_individuals = len(individuals) or 1 - save_poses_to_files(video_path, save_dir, n_individuals, bodyparts, poses, timestamp=timestamp) - - return poses, times - -def setup_video_writer( - 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) - colors = [ - ImageColor.getrgb(color) - for color in cmap_colors[:: int(len(cmap_colors) / num_keypoints)] - ] - - # 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" - - # Get video writer setup - fourcc = cv2.VideoWriter_fourcc(*"mp4v") - vwriter = cv2.VideoWriter( - filename=output_video_path, - fourcc=fourcc, - fps=fps, - frameSize=frame_size, - ) - - return colors, vwriter - -def draw_pose_and_write( - frame: np.ndarray, - pose: np.ndarray, - colors: list[tuple[int, int, int]], - bodyparts: list[str], - pcutoff: float, - display_radius: int, - draw_keypoint_names: bool, - vwriter: cv2.VideoWriter, -): - if len(pose.shape) == 2: - pose = pose[None] - - # Visualize keypoints - for i in range(pose.shape[0]): - for j in range(pose.shape[1]): - if pose[i, j, 2] > pcutoff: - x, y = map(int, pose[i, j, :2]) - cv2.circle( - frame, - center=(x, y), - radius=display_radius, - color=colors[j], - thickness=-1, - ) - - if draw_keypoint_names: - cv2.putText( - frame, - text=bodyparts[j], - org=(x + 10, y), - fontFace=cv2.FONT_HERSHEY_SIMPLEX, - fontScale=0.5, - color=colors[j], - thickness=1, - lineType=cv2.LINE_AA, - ) - - - vwriter.write(image=frame) - -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. - - Parameters - ---------- - video_path : str - Path to the analyzed video file. - save_dir : str - Directory where the pose data files will be saved. - n_individuals: int - Number of individuals - bodyparts : list of str - List of body part names corresponding to the keypoints. - poses : list of dict - List of dictionaries containing frame numbers and corresponding pose data. - - Returns - ------- - None - """ - import pandas as pd - - base_filename = Path(video_path).stem - save_dir = Path(save_dir) - h5_save_path = save_dir / f"{base_filename}_poses_{timestamp}.h5" - csv_save_path = save_dir / f"{base_filename}_poses_{timestamp}.csv" - - poses_array = _create_poses_np_array(n_individuals, bodyparts, poses) - flattened_poses = poses_array.reshape(poses_array.shape[0], -1) - - if n_individuals == 1: - pdindex = pd.MultiIndex.from_product( - [bodyparts, ["x", "y", "likelihood"]], names=["bodyparts", "coords"] - ) - else: - individuals = [f"individual_{i}" for i in range(n_individuals)] - pdindex = pd.MultiIndex.from_product( - [individuals, bodyparts, ["x", "y", "likelihood"]], names=["individuals", "bodyparts", "coords"] - ) - - pose_df = pd.DataFrame(flattened_poses, columns=pdindex) - - pose_df.to_hdf(h5_save_path, key="df_with_missing", mode="w") - 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) - pose_target_shape = (n_individuals, len(bodyparts), 3) - poses_array = np.full((max_frame + 1, *pose_target_shape), np.nan) - - for item in poses: - frame = item["frame"] - pose = item["pose"] - 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)) - padded_pose[slices] = pose[slices] - poses_array[frame] = padded_pose - - return poses_array - - -import argparse -import os - - -def main(): - """Provides a command line interface to analyze_video function.""" - - parser = argparse.ArgumentParser( - description="Analyze a video using a DeepLabCut model and visualize keypoints." - ) - parser.add_argument("model_path", type=str, help="Path to the model.") - parser.add_argument("video_path", type=str, help="Path to the video file.") - parser.add_argument("model_type", type=str, help="Type of the model (e.g., 'DLC').") - parser.add_argument( - "device", type=str, help="Device to run the model on (e.g., 'cuda' or 'cpu')." - ) - parser.add_argument( - "-p", - "--precision", - type=str, - default="FP32", - help="Model precision (e.g., 'FP32', 'FP16').", - ) - parser.add_argument( - "-d", "--display", action="store_true", help="Display keypoints on the video." - ) - parser.add_argument( - "-c", - "--pcutoff", - type=float, - default=0.5, - help="Probability cutoff for keypoints visualization.", - ) - parser.add_argument( - "-dr", - "--display-radius", - type=int, - default=5, - help="Radius of keypoint circles in the display.", - ) - parser.add_argument( - "-r", - "--resize", - type=int, - default=None, - help="Resize video frames to [width, height].", - ) - parser.add_argument( - "-x", - "--cropping", - type=int, - nargs=4, - default=None, - help="Cropping parameters [x1, x2, y1, y2].", - ) - parser.add_argument( - "-y", - "--dynamic", - type=float, - nargs=3, - default=[False, 0.5, 10], - help="Dynamic cropping [flag, pcutoff, margin].", - ) - parser.add_argument( - "--save-poses", action="store_true", help="Save the keypoint poses to files." - ) - parser.add_argument( - "--save-video", - action="store_true", - help="Save the output video with keypoints.", - ) - parser.add_argument( - "--save-dir", - type=str, - default="model_predictions", - help="Directory to save output files.", - ) - parser.add_argument( - "--draw-keypoint-names", - action="store_true", - help="Draw keypoint names on the video.", - ) - parser.add_argument( - "--cmap", type=str, default="bmy", help="Colormap for keypoints visualization." - ) - parser.add_argument( - "--no-sys-info", - action="store_false", - help="Do not print system info.", - dest="get_sys_info", - ) - - args = parser.parse_args() - - # Call the analyze_video function with the parsed arguments - benchmark( - video_path=args.video_path, - model_path=args.model_path, - model_type=args.model_type, - device=args.device, - precision=args.precision, - display=args.display, - pcutoff=args.pcutoff, - display_radius=args.display_radius, - resize=tuple(args.resize) if args.resize else None, - cropping=args.cropping, - dynamic=tuple(args.dynamic), - save_poses=args.save_poses, - save_dir=args.save_dir, - draw_keypoint_names=args.draw_keypoint_names, - cmap=args.cmap, - get_sys_info=args.get_sys_info, - save_video=args.save_video, - ) - - -if __name__ == "__main__": - main() diff --git a/dlclive/benchmark_tf.py b/dlclive/benchmark_tf.py deleted file mode 100644 index ef4dfc9..0000000 --- a/dlclive/benchmark_tf.py +++ /dev/null @@ -1,717 +0,0 @@ -""" -DeepLabCut Toolbox (deeplabcut.org) -© A. & M. Mathis Labs - -Licensed under GNU Lesser General Public License v3.0 -""" - -import os -import pickle -import platform -import subprocess -import sys -import time -import typing -import warnings - -import colorcet as cc -import ruamel -from PIL import ImageColor - -try: - from pip._internal.operations import freeze -except ImportError: - from pip.operations import freeze - -import cv2 -import numpy as np -import tensorflow as tf -from dlclive import VERSION, DLCLive -from dlclive import __file__ as dlcfile -from dlclive.utils import decode_fourcc -from tqdm import tqdm - - -def download_benchmarking_data( - target_dir=".", - url="http://deeplabcut.rowland.harvard.edu/datasets/dlclivebenchmark.tar.gz", -): - """ - Downloads a DeepLabCut-Live benchmarking Data (videos & DLC models). - """ - import tarfile - import urllib.request - - from tqdm import tqdm - - def show_progress(count, block_size, total_size): - pbar.update(block_size) - - def tarfilenamecutting(tarf): - """' auxfun to extract folder path - ie. /xyz-trainsetxyshufflez/ - """ - for memberid, member in enumerate(tarf.getmembers()): - if memberid == 0: - parent = str(member.path) - l = len(parent) + 1 - if member.path.startswith(parent): - member.path = member.path[l:] - yield member - - response = urllib.request.urlopen(url) - print( - "Downloading the benchmarking data from the DeepLabCut server @Harvard -> Go Crimson!!! {}....".format( - url - ) - ) - total_size = int(response.getheader("Content-Length")) - pbar = tqdm(unit="B", total=total_size, position=0) - filename, _ = urllib.request.urlretrieve(url, reporthook=show_progress) - with tarfile.open(filename, mode="r:gz") as tar: - tar.extractall(target_dir, members=tarfilenamecutting(tar)) - - -def get_system_info() -> dict: - """Return summary info for system running benchmark - Returns - ------- - dict - Dictionary containing the following system information: - * ``host_name`` (str): name of machine - * ``op_sys`` (str): operating system - * ``python`` (str): path to python (which conda/virtual environment) - * ``device`` (tuple): (device type (``'GPU'`` or ``'CPU'```), device information) - * ``freeze`` (list): list of installed packages and versions - * ``python_version`` (str): python version - * ``git_hash`` (str, None): If installed from git repository, hash of HEAD commit - * ``dlclive_version`` (str): dlclive version from :data:`dlclive.VERSION` - """ - - # get os - - op_sys = platform.platform() - host_name = platform.node().replace(" ", "") - - # A string giving the absolute path of the executable binary for the Python interpreter, on systems where this makes sense. - if platform.system() == "Windows": - host_python = sys.executable.split(os.path.sep)[-2] - else: - host_python = sys.executable.split(os.path.sep)[-3] - - # try to get git hash if possible - dlc_basedir = os.path.dirname(os.path.dirname(dlcfile)) - git_hash = None - try: - git_hash = subprocess.check_output( - ["git", "rev-parse", "HEAD"], cwd=dlc_basedir - ) - git_hash = git_hash.decode("utf-8").rstrip("\n") - except subprocess.CalledProcessError: - # not installed from git repo, eg. pypi - # fine, pass quietly - pass - - # get device info (GPU or CPU) - dev = None - if tf.test.is_gpu_available(): - gpu_name = tf.test.gpu_device_name() - from tensorflow.python.client import device_lib - - dev_desc = [ - d.physical_device_desc - for d in device_lib.list_local_devices() - if d.name == gpu_name - ] - dev = [d.split(",")[1].split(":")[1].strip() for d in dev_desc] - dev_type = "GPU" - else: - from cpuinfo import get_cpu_info - - dev = [get_cpu_info()["brand"]] - dev_type = "CPU" - - return { - "host_name": host_name, - "op_sys": op_sys, - "python": host_python, - "device_type": dev_type, - "device": dev, - # pip freeze to get versions of all packages - "freeze": list(freeze.freeze()), - "python_version": sys.version, - "git_hash": git_hash, - "dlclive_version": VERSION, - } - - -def benchmark( - model_path, - video_path, - tf_config=None, - resize=None, - pixels=None, - cropping=None, - dynamic=(False, 0.5, 10), - n_frames=1000, - print_rate=False, - display=False, - pcutoff=0.0, - display_radius=3, - cmap="bmy", - save_poses=False, - save_video=False, - output=None, -) -> typing.Tuple[np.ndarray, tuple, bool, dict]: - """Analyze DeepLabCut-live exported model on a video: - Calculate inference time, - display keypoints, or - get poses/create a labeled video - - Parameters - ---------- - model_path : str - path to exported DeepLabCut model - video_path : str - path to video file - tf_config : :class:`tensorflow.ConfigProto` - tensorflow session configuration - resize : int, optional - resize factor. Can only use one of resize or pixels. If both are provided, will use pixels. by default None - pixels : int, optional - downsize image to this number of pixels, maintaining aspect ratio. Can only use one of resize or pixels. If both are provided, will use pixels. by default None - cropping : list of int - cropping parameters in pixel number: [x1, x2, y1, y2] - dynamic: triple containing (state, detectiontreshold, margin) - If the state is true, then dynamic cropping will be performed. That means that if an object is detected (i.e. any body part > detectiontreshold), - then object boundaries are computed according to the smallest/largest x position and smallest/largest y position of all body parts. This window is - expanded by the margin and from then on only the posture within this crop is analyzed (until the object is lost, i.e. `, by default "bmy" - save_poses : bool, optional - flag to save poses to an hdf5 file. If True, operates similar to :function:`DeepLabCut.benchmark_videos`, by default False - save_video : bool, optional - flag to save a labeled video. If True, operates similar to :function:`DeepLabCut.create_labeled_video`, by default False - output : str, optional - path to directory to save pose and/or video file. If not specified, will use the directory of video_path, by default None - - Returns - ------- - :class:`numpy.ndarray` - vector of inference times - tuple - (image width, image height) - bool - tensorflow inference flag - dict - metadata for video - - Example - ------- - Return a vector of inference times for 10000 frames: - dlclive.benchmark('/my/exported/model', 'my_video.avi', n_frames=10000) - - Return a vector of inference times, resizing images to half the width and height for inference - dlclive.benchmark('/my/exported/model', 'my_video.avi', n_frames=10000, resize=0.5) - - Display keypoints to check the accuracy of an exported model - dlclive.benchmark('/my/exported/model', 'my_video.avi', display=True) - - Analyze a video (save poses to hdf5) and create a labeled video, similar to :function:`DeepLabCut.benchmark_videos` and :function:`create_labeled_video` - dlclive.benchmark('/my/exported/model', 'my_video.avi', save_poses=True, save_video=True) - """ - - ### load video - - cap = cv2.VideoCapture(video_path) - ret, frame = cap.read() - n_frames = ( - n_frames - if (n_frames > 0) and (n_frames < cap.get(cv2.CAP_PROP_FRAME_COUNT) - 1) - else (cap.get(cv2.CAP_PROP_FRAME_COUNT) - 1) - ) - n_frames = int(n_frames) - im_size = (cap.get(cv2.CAP_PROP_FRAME_WIDTH), cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) - - ### get resize factor - - if pixels is not None: - resize = np.sqrt(pixels / (im_size[0] * im_size[1])) - if resize is not None: - im_size = (int(im_size[0] * resize), int(im_size[1] * resize)) - - ### create video writer - - if save_video: - colors = None - out_dir = ( - output - if output is not None - else os.path.dirname(os.path.realpath(video_path)) - ) - out_vid_base = os.path.basename(video_path) - out_vid_file = os.path.normpath( - f"{out_dir}/{os.path.splitext(out_vid_base)[0]}_DLCLIVE_LABELED.avi" - ) - fourcc = cv2.VideoWriter_fourcc(*"DIVX") - fps = cap.get(cv2.CAP_PROP_FPS) - vwriter = cv2.VideoWriter(out_vid_file, fourcc, fps, im_size) - - ### check for pandas installation if using save_poses flag - - if save_poses: - try: - import pandas as pd - - use_pandas = True - except: - use_pandas = False - warnings.warn( - "Could not find installation of pandas; saving poses as a numpy array with the dimensions (n_frames, n_keypoints, [x, y, likelihood])." - ) - - ### initialize DLCLive and perform inference - - inf_times = np.zeros(n_frames) - poses = [] - - live = DLCLive( - model_path=model_path, - model_type="base", - tf_config=tf_config, - resize=resize, - cropping=cropping, - dynamic=dynamic, - display=display, - pcutoff=pcutoff, - display_radius=display_radius, - display_cmap=cmap, - ) - - poses.append(live.init_inference(frame)) - TFGPUinference = True if len(live.runner.outputs) == 1 else False - - iterator = range(n_frames) if (print_rate) or (display) else tqdm(range(n_frames)) - for i in iterator: - ret, frame = cap.read() - - if not ret: - warnings.warn( - "Did not complete {:d} frames. There probably were not enough frames in the video {}.".format( - n_frames, video_path - ) - ) - break - - start_pose = time.time() - poses.append(live.get_pose(frame)) - inf_times[i] = time.time() - start_pose - - if save_video: - if colors is None: - all_colors = getattr(cc, cmap) - colors = [ - ImageColor.getcolor(c, "RGB")[::-1] - for c in all_colors[:: int(len(all_colors) / poses[-1].shape[0])] - ] - - this_pose = poses[-1] - for j in range(this_pose.shape[0]): - if this_pose[j, 2] > pcutoff: - x = int(this_pose[j, 0]) - y = int(this_pose[j, 1]) - frame = cv2.circle( - frame, (x, y), display_radius, colors[j], thickness=-1 - ) - - if resize is not None: - frame = cv2.resize(frame, im_size) - vwriter.write(frame) - - if print_rate: - print("pose rate = {:d}".format(int(1 / inf_times[i]))) - - if print_rate: - print("mean pose rate = {:d}".format(int(np.mean(1 / inf_times)))) - - ### gather video and test parameterization - - # dont want to fail here so gracefully failing on exception -- - # eg. some packages of cv2 don't have CAP_PROP_CODEC_PIXEL_FORMAT - try: - fourcc = decode_fourcc(cap.get(cv2.CAP_PROP_FOURCC)) - except: - fourcc = "" - - try: - fps = round(cap.get(cv2.CAP_PROP_FPS)) - except: - fps = None - - try: - pix_fmt = decode_fourcc(cap.get(cv2.CAP_PROP_CODEC_PIXEL_FORMAT)) - except: - pix_fmt = "" - - try: - frame_count = round(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - except: - frame_count = None - - try: - orig_im_size = ( - round(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), - round(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), - ) - except: - orig_im_size = None - - meta = { - "video_path": video_path, - "video_codec": fourcc, - "video_pixel_format": pix_fmt, - "video_fps": fps, - "video_total_frames": frame_count, - "original_frame_size": orig_im_size, - "dlclive_params": live.parameterization, - } - - ### close video and tensorflow session - - cap.release() - live.close() - - if save_video: - vwriter.release() - - if save_poses: - cfg_path = os.path.normpath(f"{model_path}/pose_cfg.yaml") - ruamel_file = ruamel.yaml.YAML() - dlc_cfg = ruamel_file.load(open(cfg_path, "r")) - bodyparts = dlc_cfg["all_joints_names"] - poses = np.array(poses) - - if use_pandas: - poses = poses.reshape((poses.shape[0], poses.shape[1] * poses.shape[2])) - pdindex = pd.MultiIndex.from_product( - [bodyparts, ["x", "y", "likelihood"]], names=["bodyparts", "coords"] - ) - pose_df = pd.DataFrame(poses, columns=pdindex) - - out_dir = ( - output - if output is not None - else os.path.dirname(os.path.realpath(video_path)) - ) - out_vid_base = os.path.basename(video_path) - out_dlc_file = os.path.normpath( - f"{out_dir}/{os.path.splitext(out_vid_base)[0]}_DLCLIVE_POSES.h5" - ) - pose_df.to_hdf(out_dlc_file, key="df_with_missing", mode="w") - - else: - out_vid_base = os.path.basename(video_path) - out_dlc_file = os.path.normpath( - f"{out_dir}/{os.path.splitext(out_vid_base)[0]}_DLCLIVE_POSES.npy" - ) - np.save(out_dlc_file, poses) - - return inf_times, im_size, TFGPUinference, meta - - -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. - This is primarily used through :function:`benchmark_videos` - - - Parameters - ---------- - sys_info : tuple - system information generated by :func:`get_system_info` - inf_times : :class:`numpy.ndarray` - array of inference times generated by :func:`benchmark` - im_size : tuple or :class:`numpy.ndarray` - image size (width, height) for each benchmark run. If an array, each row corresponds to a row in inf_times - TFGPUinference: bool - flag if using tensorflow inference or numpy inference DLC model - model: str, optional - name of model - meta : dict, optional - metadata returned by :func:`benchmark` - output : str, optional - path to directory to save data. If None, uses pwd, by default None - - Returns - ------- - bool - flag indicating successful save - """ - - output = output if output is not None else os.getcwd() - model_type = None - if model is not None: - if "resnet" in model: - model_type = "resnet" - elif "mobilenet" in model: - model_type = "mobilenet" - else: - model_type = None - - fn_ind = 0 - base_name = ( - f"benchmark_{sys_info['host_name']}_{sys_info['device_type']}_{fn_ind}.pickle" - ) - out_file = os.path.normpath(f"{output}/{base_name}") - while os.path.isfile(out_file): - fn_ind += 1 - base_name = f"benchmark_{sys_info['host_name']}_{sys_info['device_type']}_{fn_ind}.pickle" - out_file = os.path.normpath(f"{output}/{base_name}") - - # summary stats (mean inference time & standard error of mean) - stats = zip( - np.mean(inf_times, 1), - np.std(inf_times, 1) * 1.0 / np.sqrt(np.shape(inf_times)[1]), - ) - - # for stat in stats: - # print("Stats:", stat) - - data = { - "model": model, - "model_type": model_type, - "TFGPUinference": TFGPUinference, - "im_size": im_size, - "inference_times": inf_times, - "stats": stats, - } - - data.update(sys_info) - if meta: - data.update(meta) - - os.makedirs(os.path.normpath(output), exist_ok=True) - pickle.dump(data, open(out_file, "wb")) - - return True - - -def benchmark_videos( - model_path, - video_path, - output=None, - n_frames=1000, - tf_config=None, - resize=None, - pixels=None, - cropping=None, - dynamic=(False, 0.5, 10), - print_rate=False, - display=False, - pcutoff=0.5, - display_radius=3, - cmap="bmy", - save_poses=False, - save_video=False, -): - """Analyze videos using DeepLabCut-live exported models. - Analyze multiple videos and/or multiple options for the size of the video - by specifying a resizing factor or the number of pixels to use in the image (keeping aspect ratio constant). - Options to record inference times (to examine inference speed), - display keypoints to visually check the accuracy, - or save poses to an hdf5 file as in :function:`deeplabcut.benchmark_videos` and - create a labeled video as in :function:`deeplabcut.create_labeled_video`. - - Parameters - ---------- - model_path : str - path to exported DeepLabCut model - video_path : str or list - path to video file or list of paths to video files - output : str - path to directory to save results - tf_config : :class:`tensorflow.ConfigProto` - tensorflow session configuration - resize : int, optional - resize factor. Can only use one of resize or pixels. If both are provided, will use pixels. by default None - pixels : int, optional - downsize image to this number of pixels, maintaining aspect ratio. Can only use one of resize or pixels. If both are provided, will use pixels. by default None - cropping : list of int - cropping parameters in pixel number: [x1, x2, y1, y2] - dynamic: triple containing (state, detectiontreshold, margin) - If the state is true, then dynamic cropping will be performed. That means that if an object is detected (i.e. any body part > detectiontreshold), - then object boundaries are computed according to the smallest/largest x position and smallest/largest y position of all body parts. This window is - expanded by the margin and from then on only the posture within this crop is analyzed (until the object is lost, i.e. `, by default "bmy" - save_poses : bool, optional - flag to save poses to an hdf5 file. If True, operates similar to :function:`DeepLabCut.benchmark_videos`, by default False - save_video : bool, optional - flag to save a labeled video. If True, operates similar to :function:`DeepLabCut.create_labeled_video`, by default False - - Example - ------- - Return a vector of inference times for 10000 frames on one video or two videos: - dlclive.benchmark_videos('/my/exported/model', 'my_video.avi', n_frames=10000) - dlclive.benchmark_videos('/my/exported/model', ['my_video1.avi', 'my_video2.avi'], n_frames=10000) - - Return a vector of inference times, testing full size and resizing images to half the width and height for inference, for two videos - dlclive.benchmark_videos('/my/exported/model', ['my_video1.avi', 'my_video2.avi'], n_frames=10000, resize=[1.0, 0.5]) - - Display keypoints to check the accuracy of an exported model - dlclive.benchmark_videos('/my/exported/model', 'my_video.avi', display=True) - - Analyze a video (save poses to hdf5) and create a labeled video, similar to :function:`DeepLabCut.benchmark_videos` and :function:`create_labeled_video` - dlclive.benchmark_videos('/my/exported/model', 'my_video.avi', save_poses=True, save_video=True) - """ - - # convert video_paths to list - - video_path = video_path if type(video_path) is list else [video_path] - - # fix resize - - if pixels: - pixels = pixels if type(pixels) is list else [pixels] - resize = [None for p in pixels] - elif resize: - resize = resize if type(resize) is list else [resize] - pixels = [None for r in resize] - else: - resize = [None] - pixels = [None] - - # loop over videos - - for v in video_path: - # initialize full inference times - - inf_times = [] - im_size_out = [] - - for i in range(len(resize)): - print(f"\nRun {i+1} / {len(resize)}\n") - - this_inf_times, this_im_size, TFGPUinference, meta = benchmark( - model_path, - v, - tf_config=tf_config, - resize=resize[i], - pixels=pixels[i], - cropping=cropping, - dynamic=dynamic, - n_frames=n_frames, - print_rate=print_rate, - display=display, - pcutoff=pcutoff, - display_radius=display_radius, - cmap=cmap, - save_poses=save_poses, - save_video=save_video, - output=output, - ) - - inf_times.append(this_inf_times) - im_size_out.append(this_im_size) - - inf_times = np.array(inf_times) - im_size_out = np.array(im_size_out) - - # save results - - if output is not None: - sys_info = get_system_info() - save_inf_times( - sys_info, - inf_times, - im_size_out, - TFGPUinference, - model=os.path.basename(model_path), - meta=meta, - output=output, - ) - - -def main(): - """Provides a command line interface :function:`benchmark_videos`""" - - import argparse - - parser = argparse.ArgumentParser() - parser.add_argument("model_path", type=str) - parser.add_argument("video_path", type=str, nargs="+") - parser.add_argument("-o", "--output", type=str, default=None) - parser.add_argument("-n", "--n-frames", type=int, default=1000) - parser.add_argument("-r", "--resize", type=float, nargs="+") - parser.add_argument("-p", "--pixels", type=float, nargs="+") - parser.add_argument("-v", "--print-rate", default=False, action="store_true") - parser.add_argument("-d", "--display", default=False, action="store_true") - parser.add_argument("-l", "--pcutoff", default=0.5, type=float) - parser.add_argument("-s", "--display-radius", default=3, type=int) - parser.add_argument("-c", "--cmap", type=str, default="bmy") - parser.add_argument("--cropping", nargs="+", type=int, default=None) - parser.add_argument("--dynamic", nargs="+", type=float, default=[]) - parser.add_argument("--save-poses", action="store_true") - parser.add_argument("--save-video", action="store_true") - args = parser.parse_args() - - if (args.cropping) and (len(args.cropping) < 4): - raise Exception( - "Cropping not properly specified. Must provide 4 values: x1, x2, y1, y2" - ) - - if not args.dynamic: - args.dynamic = (False, 0.5, 10) - elif len(args.dynamic) < 3: - raise Exception( - "Dynamic cropping not properly specified. Must provide three values: 0 or 1 as boolean flag, pcutoff, and margin" - ) - else: - args.dynamic = (bool(args.dynamic[0]), args.dynamic[1], args.dynamic[2]) - - benchmark_videos( - args.model_path, - args.video_path, - output=args.output, - resize=args.resize, - pixels=args.pixels, - cropping=args.cropping, - dynamic=args.dynamic, - n_frames=args.n_frames, - print_rate=args.print_rate, - display=args.display, - pcutoff=args.pcutoff, - display_radius=args.display_radius, - cmap=args.cmap, - save_poses=args.save_poses, - save_video=args.save_video, - ) - - -if __name__ == "__main__": - main() diff --git a/dlclive/check_install/check_install.py b/dlclive/check_install/check_install.py index 30d6e79..2bc4e65 100755 --- a/dlclive/check_install/check_install.py +++ b/dlclive/check_install/check_install.py @@ -5,6 +5,8 @@ Licensed under GNU Lesser General Public License v3.0 """ +import os +import urllib.request import argparse import shutil import sys @@ -14,7 +16,8 @@ from dlclibrary.dlcmodelzoo.modelzoo_download import download_huggingface_model -from dlclive.benchmark_tf import benchmark_videos +from dlclive.benchmark import benchmark_videos +from dlclive.engine import Engine MODEL_NAME = "superanimal_quadruped" SNAPSHOT_NAME = "snapshot-700000.pb" @@ -43,8 +46,7 @@ def main(): if not display: print("Running without displaying video") - # make temporary directory in $HOME - # TODO: why create this temp directory in $HOME? + # make temporary directory in $current print("\nCreating temporary directory...\n") tmp_dir = Path().home() / "dlc-live-tmp" tmp_dir.mkdir(mode=0o775, exist_ok=True) @@ -53,16 +55,18 @@ def main(): 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) - print(f"Downloading Video to {video_file}") - url_link = "https://github.com/DeepLabCut/DeepLabCut-live/blob/master/check_install/dog_clip.avi?raw=True" - urllib.request.urlretrieve(url_link, video_file, reporthook=urllib_pbar) + if not os.path.exists(video_file): + print(f"Downloading Video to {video_file}") + url_link = "https://github.com/DeepLabCut/DeepLabCut-live/blob/main/check_install/dog_clip.avi?raw=True" + urllib.request.urlretrieve(url_link, video_file, reporthook=urllib_pbar) + else: + print(f"Video already exists at {video_file}") # download model from the DeepLabCut Model Zoo if Path(model_dir / SNAPSHOT_NAME).exists(): print("Model already downloaded, using cached version") else: - print("Downloading full_dog model from the DeepLabCut Model Zoo...") + print("Downloading a test model from the DeepLabCut Model Zoo...") download_huggingface_model(MODEL_NAME, model_dir) # assert these things exist so we can give informative error messages @@ -74,7 +78,12 @@ def main(): # run benchmark videos print("\n Running inference...\n") benchmark_videos( - str(model_dir), video_file, display=display, resize=0.5, pcutoff=0.25 + model_path=str(model_dir), + model_type="base" if Engine.from_model_path(model_dir) == Engine.TENSORFLOW else "pytorch", + video_path=video_file, + display=display, + resize=0.5, + pcutoff=0.25 ) # deleting temporary files diff --git a/dlclive/engine.py b/dlclive/engine.py new file mode 100644 index 0000000..eed0af1 --- /dev/null +++ b/dlclive/engine.py @@ -0,0 +1,33 @@ +from enum import Enum +from pathlib import Path + +class Engine(Enum): + TENSORFLOW = "tensorflow" + PYTORCH = "pytorch" + + @classmethod + def from_model_type(cls, model_type: str) -> "Engine": + if model_type.lower() == "pytorch": + return cls.PYTORCH + elif model_type.lower() in ("tensorflow", "base", "tensorrt", "lite"): + return cls.TENSORFLOW + else: + raise ValueError(f"Unknown model type: {model_type}") + + @classmethod + def from_model_path(cls, model_path: str | Path) -> "Engine": + path = Path(model_path) + + if not path.exists(): + raise FileNotFoundError(f"Model path does not exist: {model_path}") + + if path.is_dir(): + has_cfg = (path / "pose_cfg.yaml").is_file() + has_pb = any(p.suffix == ".pb" for p in path.glob("*.pb")) + if has_cfg and has_pb: + return cls.TENSORFLOW + elif path.is_file(): + if path.suffix == ".pt": + return cls.PYTORCH + + raise ValueError(f"Could not determine engine from model path: {model_path}") \ No newline at end of file diff --git a/dlclive/factory.py b/dlclive/factory.py index 0c22e22..df70029 100644 --- a/dlclive/factory.py +++ b/dlclive/factory.py @@ -5,6 +5,7 @@ from typing import Literal from dlclive.core.runner import BaseRunner +from dlclive.engine import Engine def build_runner( @@ -16,7 +17,7 @@ def build_runner( Parameters ---------- - model_type: str, optional + model_type: str Which model to use. For the PyTorch engine, options are [`pytorch`]. For the TensorFlow engine, options are [`base`, `tensorrt`, `lite`]. model_path: str, Path @@ -33,13 +34,13 @@ def build_runner( ------- """ - if model_type.lower() == "pytorch": + if Engine.from_model_type(model_type) == Engine.PYTORCH: from dlclive.pose_estimation_pytorch.runner import PyTorchRunner valid = {"device", "precision", "single_animal", "dynamic", "top_down_config"} return PyTorchRunner(model_path, **filter_keys(valid, kwargs)) - elif model_type.lower() in ("tensorflow", "base", "tensorrt", "lite"): + elif Engine.from_model_type(model_type) == Engine.TENSORFLOW: from dlclive.pose_estimation_tensorflow.runner import TensorFlowRunner if model_type.lower() == "tensorflow": diff --git a/pyproject.toml b/pyproject.toml index 4b14185..7597994 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,17 +25,19 @@ dlc-live-benchmark = "dlclive.benchmark:main" [tool.poetry.dependencies] python = ">=3.10,<3.12" -numpy = ">=1.20,<2" +numpy = ">=1.26,<2.0" "ruamel.yaml" = "^0.17.20" colorcet = "^3.0.0" einops = ">=0.6.1" Pillow = ">=8.0.0" +opencv-python-headless = ">=4.5.0,<5.0.0" py-cpuinfo = ">=5.0.0" tqdm = "^4.62.3" pandas = ">=1.0.1,!=1.5.0" tables = "^3.8" -opencv-python-headless = "^4.5" +pytest = "^8.0" dlclibrary = ">=0.0.6" + # PyTorch models scipy = ">=1.9" timm = { version = ">=1.0.7", optional = true } diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..d08b10e --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +markers = + functional: functional tests diff --git a/tests/test_benchmark_script.py b/tests/test_benchmark_script.py new file mode 100644 index 0000000..3aadb04 --- /dev/null +++ b/tests/test_benchmark_script.py @@ -0,0 +1,48 @@ +import os +import glob +import pathlib +import pytest +from dlclive import benchmark_videos, download_benchmarking_data +from dlclive.engine import Engine + +@pytest.mark.functional +def test_benchmark_script_runs(tmp_path): + datafolder = tmp_path / "Data-DLC-live-benchmark" + download_benchmarking_data(str(datafolder)) + + dog_models = glob.glob(str(datafolder / "dog" / "*[!avi]")) + dog_video = glob.glob(str(datafolder / "dog" / "*.avi"))[0] + mouse_models = glob.glob(str(datafolder / "mouse_lick" / "*[!avi]")) + mouse_video = glob.glob(str(datafolder / "mouse_lick" / "*.avi"))[0] + + out_dir = tmp_path / "results" + out_dir.mkdir(exist_ok=True) + + pixels = [100, 400] #[2500, 10000] + n_frames = 5 + + for model_path in dog_models: + print(f"Running dog model: {model_path}") + result = benchmark_videos( + model_path=model_path, + model_type="base" if Engine.from_model_path(model_path) == Engine.TENSORFLOW else "pytorch", + video_path=dog_video, + output=str(out_dir), + n_frames=n_frames, + pixels=pixels + ) + print("Dog model result:", result) + + for model_path in mouse_models: + print(f"Running mouse model: {model_path}") + result = benchmark_videos( + model_path=model_path, + model_type="base" if Engine.from_model_path(model_path) == Engine.TENSORFLOW else "pytorch", + video_path=mouse_video, + output=str(out_dir), + n_frames=n_frames, + pixels=pixels + ) + print("Mouse model result:", result) + + assert any(out_dir.iterdir())