From 08ebec25f238d025bb4e69fa3e5bc20319e9d266 Mon Sep 17 00:00:00 2001
From: sneakers-the-rat <sneakers-the-rat@protonmail.com>
Date: Wed, 14 Feb 2024 00:53:56 -0800
Subject: [PATCH] fix #111: coerce floats and ints for pixels, replace string
 with float in readme

---
 README.md            |  2 +-
 dlclive/benchmark.py | 18 +++++++++---------
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index 311d22c..0e91ed5 100644
--- a/README.md
+++ b/README.md
@@ -93,7 +93,7 @@ DeepLabCut-live offers some analysis tools that allow users to peform the follow
 1. Test inference speed across a range of image sizes, downsizing images by specifying the `resize` or `pixels` parameter. Using the `pixels` parameter will resize images to the desired number of `pixels`, without changing the aspect ratio. Results will be saved (along with system info) to a pickle file if you specify an output directory.
 ##### python
 ```python
-dlclive.benchmark_videos('/path/to/exported/model', ['/path/to/video1', '/path/to/video2'], output='/path/to/output', resize=[1.0, 0.75, '0.5'])
+dlclive.benchmark_videos('/path/to/exported/model', ['/path/to/video1', '/path/to/video2'], output='/path/to/output', resize=[1.0, 0.75, 0.5])
 ```
 ##### command line
 ```
diff --git a/dlclive/benchmark.py b/dlclive/benchmark.py
index 4cb4fb1..b5a345c 100644
--- a/dlclive/benchmark.py
+++ b/dlclive/benchmark.py
@@ -12,7 +12,7 @@
 import sys
 import warnings
 import subprocess
-import typing
+from typing import List, Optional, Tuple, Union
 import pickle
 import colorcet as cc
 from PIL import ImageColor
@@ -151,8 +151,8 @@ def benchmark(
     model_path,
     video_path,
     tf_config=None,
-    resize=None,
-    pixels=None,
+    resize: Optional[float] = None,
+    pixels: Optional[int] = None,
     cropping=None,
     dynamic=(False, 0.5, 10),
     n_frames=1000,
@@ -164,7 +164,7 @@ def benchmark(
     save_poses=False,
     save_video=False,
     output=None,
-) -> typing.Tuple[np.ndarray, tuple, bool, dict]:
+) -> Tuple[np.ndarray, tuple, bool, dict]:
     """ Analyze DeepLabCut-live exported model on a video:
     Calculate inference time,
     display keypoints, or
@@ -521,8 +521,8 @@ def benchmark_videos(
     output=None,
     n_frames=1000,
     tf_config=None,
-    resize=None,
-    pixels=None,
+    resize: Optional[Union[float, List[float]]] = None,
+    pixels: Optional[Union[int, List[int]]] = None,
     cropping=None,
     dynamic=(False, 0.5, 10),
     print_rate=False,
@@ -551,7 +551,7 @@ def benchmark_videos(
         path to directory to save results
     tf_config : :class:`tensorflow.ConfigProto`
         tensorflow session configuration
-    resize : int, optional
+    resize : float, 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
@@ -603,10 +603,10 @@ def benchmark_videos(
     # fix resize
 
     if pixels:
-        pixels = pixels if type(pixels) is list else [pixels]
+        pixels = [int(p) for p in pixels] if type(pixels) is list else [int(pixels)]
         resize = [None for p in pixels]
     elif resize:
-        resize = resize if type(resize) is list else [resize]
+        resize = [float(r) for r in resize] if type(resize) is list else [float(resize)]
         pixels = [None for r in resize]
     else:
         resize = [None]