diff --git a/README.md b/README.md
index 6f3c668..8915b47 100644
--- a/README.md
+++ b/README.md
@@ -1,9 +1,16 @@
# DeepLabCut-live SDK
+
-[](https://github.com/DeepLabCut/deeplabcutlive/raw/master/LICENSE)
+[](https://github.com/DeepLabCut/DeepLabCut-live)
+[](https://github.com/DeepLabCut/DeepLabCut-live)
+
+[](https://github.com/DeepLabCut/deeplabcutlive/raw/master/LICENSE)
+[](https://forum.image.sc/tags/deeplabcut)
+[](https://gitter.im/DeepLabCut/community?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge)
+[](https://twitter.com/DeepLabCut)
This package contains a [DeepLabCut](http://www.mousemotorlab.org/deeplabcut) inference pipeline for real-time applications that has minimal (software) dependencies. Thus, it is as easy to install as possible (in particular, on atypical systems like [NVIDIA Jetson boards](https://developer.nvidia.com/buy-jetson)).
@@ -31,7 +38,7 @@ Note, you can test your installation by running:
`python /check_install/check_install.py`
-This will download
+If installed properly, this script will i) create a temporary folder ii) download the full_dog model from the [DeepLabCut Model Zoo](http://www.mousemotorlab.org/dlc-modelzoo), iii) download a short video clip of a dog, and iv) run inference while displaying keypoints. v) remove the temporary folder.
### Quick Start: instructions for use:
diff --git a/benchmarking/run.py b/benchmarking/run.py
index 3a3104e..db9d222 100644
--- a/benchmarking/run.py
+++ b/benchmarking/run.py
@@ -13,26 +13,35 @@
from dlclive import benchmark_model_by_size
# Update the datafolder to where the data is:
-datafolder='/your/path/to/data/here'
+datafolder = "/your/path/to/data/here"
-n_frames = 1000 #change to 10000 for testing on a GPU!
+n_frames = 1000 # change to 10000 for testing on a GPU!
pixels = [2500, 10000, 40000, 160000, 320000, 640000]
-dog_models = glob.glob(datafolder + '/dog/*[!avi]')
-dog_video = glob.glob(datafolder + '/dog/*.avi')[0]
-mouse_models = glob.glob(datafolder + '/mouse_lick/*[!avi]')
-mouse_video = glob.glob(datafolder + '/mouse_lick/*.avi')[0]
+dog_models = glob.glob(datafolder + "/dog/*[!avi]")
+dog_video = glob.glob(datafolder + "/dog/*.avi")[0]
+mouse_models = glob.glob(datafolder + "/mouse_lick/*[!avi]")
+mouse_video = glob.glob(datafolder + "/mouse_lick/*.avi")[0]
this_dir = os.path.dirname(os.path.realpath(__file__))
-#storing results in /benchmarking/results: (for your PR)
-out_dir = os.path.normpath(this_dir + '/results')
+# storing results in /benchmarking/results: (for your PR)
+out_dir = os.path.normpath(this_dir + "/results")
for ind_m, m in enumerate(dog_models):
print("\n\nMODEL {:d} / 8".format(ind_m))
- benchmark_model_by_size(m, dog_video, ind_m, out_dir=out_dir, n_frames=n_frames, pixels=pixels)
+ benchmark_model_by_size(
+ m, dog_video, ind_m, out_dir=out_dir, n_frames=n_frames, pixels=pixels
+ )
-offset=ind_m+1
+offset = ind_m + 1
for ind_m, m in enumerate(mouse_models):
print("\n\nMODEL {:d} / 8".format(ind_m))
- benchmark_model_by_size(m, mouse_video, ind_m + offset, out_dir=out_dir, n_frames=n_frames, pixels=pixels)
+ benchmark_model_by_size(
+ m,
+ mouse_video,
+ ind_m + offset,
+ out_dir=out_dir,
+ n_frames=n_frames,
+ pixels=pixels,
+ )
diff --git a/check_install/check_install.py b/check_install/check_install.py
index 88694bc..2c3c60d 100755
--- a/check_install/check_install.py
+++ b/check_install/check_install.py
@@ -27,7 +27,7 @@ def main():
# download dog video clip from github
print("\nDownloading dog video clip...\n")
- #video_url = "https://github.com/DeepLabCut/DeepLabCut-live/raw/master/check_install/dog_clip.avi"
+ # video_url = "https://github.com/DeepLabCut/DeepLabCut-live/raw/master/check_install/dog_clip.avi"
video_url = '"https://docs.google.com/uc?export=download&id=1W_5AOl1SewXR2q5QC1K5Chm71I9LAmld"'
os.system(f"curl -L {video_url} -o dog_clip.avi")
@@ -45,4 +45,4 @@ def main():
if __name__ == "__main__":
- main()
\ No newline at end of file
+ main()
diff --git a/dlclive/bench.py b/dlclive/bench.py
index 4ba0b91..f8b1f55 100644
--- a/dlclive/bench.py
+++ b/dlclive/bench.py
@@ -33,6 +33,7 @@
from dlclive.utils import decode_fourcc
+
def get_system_info() -> dict:
""" Return summary info for system running benchmark
@@ -51,14 +52,13 @@ def get_system_info() -> dict:
* ``dlclive_version`` (str): dlclive version from :data:`dlclive.VERSION`
"""
-
### get os
op_sys = platform.platform()
- host_name = platform.node().replace(' ', '')
+ 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':
+ if platform.system() == "Windows":
host_python = sys.executable.split(os.path.sep)[-2]
else:
host_python = sys.executable.split(os.path.sep)[-3]
@@ -67,8 +67,10 @@ def get_system_info() -> dict:
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')
+ 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
@@ -80,31 +82,46 @@ def get_system_info() -> dict:
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_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() #[get_cpu_info()['brand']]
+
+ dev = get_cpu_info() # [get_cpu_info()['brand']]
dev_type = "CPU"
# return a dictionary rather than a tuple for inspectability's sake
return {
- 'host_name': host_name,
- 'op_sys' : op_sys,
- 'python': host_python,
- 'device_type': dev_type,
- 'device': dev,
- 'freeze': list(freeze.freeze()), # pip freeze to get versions of all packages
- 'python_version': sys.version,
- 'git_hash': git_hash,
- 'dlclive_version': VERSION
+ "host_name": host_name,
+ "op_sys": op_sys,
+ "python": host_python,
+ "device_type": dev_type,
+ "device": dev,
+ "freeze": list(freeze.freeze()), # pip freeze to get versions of all packages
+ "python_version": sys.version,
+ "git_hash": git_hash,
+ "dlclive_version": VERSION,
}
-def run_benchmark(model_path, video_path, tf_config=None,
- resize=None, pixels=None, n_frames=10000,
- print_rate=False, display=False, pcutoff=0.0,
- display_radius=3) -> typing.Tuple[np.ndarray, int, bool, dict]:
+
+def run_benchmark(
+ model_path,
+ video_path,
+ tf_config=None,
+ resize=None,
+ pixels=None,
+ n_frames=10000,
+ print_rate=False,
+ display=False,
+ pcutoff=0.0,
+ display_radius=3,
+) -> typing.Tuple[np.ndarray, int, bool, dict]:
""" Benchmark on inference times for a given DLC model and video
Parameters
@@ -145,7 +162,14 @@ def run_benchmark(model_path, video_path, tf_config=None,
### initialize live object
- live = DLCLive(model_path, tf_config=tf_config, resize=resize, display=display, pcutoff=pcutoff, display_radius=display_radius)
+ live = DLCLive(
+ model_path,
+ tf_config=tf_config,
+ resize=resize,
+ display=display,
+ pcutoff=pcutoff,
+ display_radius=display_radius,
+ )
live.init_inference(frame)
TFGPUinference = True if len(live.outputs) == 1 else False
@@ -159,9 +183,13 @@ def run_benchmark(model_path, video_path, tf_config=None,
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))
+ 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()
live.get_pose(frame)
inf_times[i] = time.time() - start_pose
@@ -170,7 +198,7 @@ def run_benchmark(model_path, video_path, tf_config=None,
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))))
+ print("mean pose rate = {:d}".format(int(np.mean(1 / inf_times))))
### close video and tensorflow session
@@ -198,41 +226,45 @@ def run_benchmark(model_path, video_path, tf_config=None,
except:
frame_count = None
-
-
meta = {
- 'video_path': video_path,
- 'video_codec': fourcc,
- 'video_pixel_format': pix_fmt,
- 'video_fps': fps,
- 'video_total_frames': frame_count,
- 'resize': resize,
- 'original_frame_size': im_size,
- 'resized_frame_size': (im_size[0]*resize, im_size[1]*resize),
- 'pixels': pixels,
- 'dlclive_params': live.parameterization
+ "video_path": video_path,
+ "video_codec": fourcc,
+ "video_pixel_format": pix_fmt,
+ "video_fps": fps,
+ "video_total_frames": frame_count,
+ "resize": resize,
+ "original_frame_size": im_size,
+ "resized_frame_size": (im_size[0] * resize, im_size[1] * resize),
+ "pixels": pixels,
+ "dlclive_params": live.parameterization,
}
cap.release()
live.close()
- return inf_times, resize*im_size[0] * resize*im_size[1], TFGPUinference, meta
+ return inf_times, resize * im_size[0] * resize * im_size[1], TFGPUinference, meta
-def get_savebenchmarkfn(sys_info ,i, fn_ind, out_dir=None):
- ''' get filename to save data (definitions see save_benchmark)'''
+
+def get_savebenchmarkfn(sys_info, i, fn_ind, out_dir=None):
+ """ get filename to save data (definitions see save_benchmark)"""
out_dir = out_dir if out_dir is not None else os.getcwd()
- base_name = "benchmark_{}_{}_{}_{}.pickle".format(sys_info['host_name'], sys_info['device_type'], fn_ind, i)
- datafilename = out_dir + '/' + base_name
+ base_name = "benchmark_{}_{}_{}_{}.pickle".format(
+ sys_info["host_name"], sys_info["device_type"], fn_ind, i
+ )
+ datafilename = out_dir + "/" + base_name
return datafilename
-def save_benchmark(sys_info: dict,
- inf_times: np.ndarray,
- pixels: typing.Union[np.ndarray, float],
- iter: int,
- TFGPUinference: bool = None,
- model: str = None,
- out_dir: str = None,
- meta: dict=None):
+
+def save_benchmark(
+ sys_info: dict,
+ inf_times: np.ndarray,
+ pixels: typing.Union[np.ndarray, float],
+ iter: int,
+ TFGPUinference: bool = None,
+ model: str = None,
+ out_dir: str = None,
+ meta: dict = None,
+):
""" Save benchmarking data with system information to a pickle file
Parameters
@@ -264,30 +296,30 @@ def save_benchmark(sys_info: dict,
model_type = None
if model is not None:
- if 'resnet' in model:
- model_type = 'resnet'
- elif 'mobilenet' in model:
- model_type = 'mobilenet'
+ if "resnet" in model:
+ model_type = "resnet"
+ elif "mobilenet" in model:
+ model_type = "mobilenet"
else:
model_type = None
fn_ind = 0
- base_name = "benchmark_{}_{}_{}_{}.pickle".format(sys_info['host_name'],
- sys_info['device'][0],
- fn_ind,
- iter)
- while os.path.isfile(os.path.normpath(out_dir + '/' + base_name)):
+ base_name = "benchmark_{}_{}_{}_{}.pickle".format(
+ sys_info["host_name"], sys_info["device"][0], fn_ind, iter
+ )
+ while os.path.isfile(os.path.normpath(out_dir + "/" + base_name)):
fn_ind += 1
- base_name = "benchmark_{}_{}_{}_{}.pickle".format(sys_info['host_name'],
- sys_info['device'][0],
- fn_ind,
- iter)
-
- data = {'model': model,
- 'model_type': model_type,
- 'TFGPUinference': TFGPUinference,
- 'pixels': pixels,
- 'inference_times': inf_times}
+ base_name = "benchmark_{}_{}_{}_{}.pickle".format(
+ sys_info["host_name"], sys_info["device"][0], fn_ind, iter
+ )
+
+ data = {
+ "model": model,
+ "model_type": model_type,
+ "TFGPUinference": TFGPUinference,
+ "pixels": pixels,
+ "inference_times": inf_times,
+ }
data.update(sys_info)
@@ -297,16 +329,30 @@ def save_benchmark(sys_info: dict,
data.update(sys_info)
datafilename = os.path.normpath(f"{out_dir}/{base_name}")
- pickle.dump(data, open(os.path.normpath(datafilename), 'wb'))
+ pickle.dump(data, open(os.path.normpath(datafilename), "wb"))
return True
+
def read_pickle(filename):
""" Read the pickle file """
with open(filename, "rb") as handle:
return pickle.load(handle)
-def benchmark_model_by_size(model_path, video_path, output=None, n_frames=10000, tf_config=None, resize=None, pixels=None, print_rate=False, display=False, pcutoff=0.5, display_radius=3):
+
+def benchmark_model_by_size(
+ model_path,
+ video_path,
+ output=None,
+ n_frames=10000,
+ tf_config=None,
+ resize=None,
+ pixels=None,
+ print_rate=False,
+ display=False,
+ pcutoff=0.5,
+ display_radius=3,
+):
"""Benchmark DLC model by image size
Parameters
@@ -351,7 +397,7 @@ def benchmark_model_by_size(model_path, video_path, output=None, n_frames=10000,
for i in range(len(resize)):
- print("\nRun {:d} / {:d}\n".format(i+1, len(resize)))
+ print("\nRun {:d} / {:d}\n".format(i + 1, len(resize)))
inf_times, pixels_out, TFGPUinference, benchmark_meta = run_benchmark(
model_path,
@@ -363,42 +409,51 @@ def benchmark_model_by_size(model_path, video_path, output=None, n_frames=10000,
print_rate=print_rate,
display=display,
pcutoff=pcutoff,
- display_radius=display_radius)
+ display_radius=display_radius,
+ )
- #TODO: check if a part has already been complted?
+ # TODO: check if a part has already been complted?
### saving results intermediately
- save_benchmark(sys_info, inf_times, pixels_out, i, TFGPUinference,
- model=os.path.basename(model_path),
- out_dir = output,
- meta=benchmark_meta)
+ save_benchmark(
+ sys_info,
+ inf_times,
+ pixels_out,
+ i,
+ TFGPUinference,
+ model=os.path.basename(model_path),
+ out_dir=output,
+ meta=benchmark_meta,
+ )
+
def main():
parser = argparse.ArgumentParser()
- parser.add_argument('model_path', type=str)
- parser.add_argument('video_path', type=str)
- parser.add_argument('-o', '--output', type=str, default=os.getcwd())
- parser.add_argument('-n', '--n-frames', type=int, default=10000)
- 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("model_path", type=str)
+ parser.add_argument("video_path", type=str)
+ parser.add_argument("-o", "--output", type=str, default=os.getcwd())
+ parser.add_argument("-n", "--n-frames", type=int, default=10000)
+ 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)
args = parser.parse_args()
-
- benchmark_model_by_size(args.model_path,
- args.video_path,
- output=args.output,
- resize=args.resize,
- pixels=args.pixels,
- n_frames=args.n_frames,
- print_rate=args.print_rate,
- display=args.display,
- pcutoff=args.pcutoff,
- display_radius=args.display_radius)
+ benchmark_model_by_size(
+ args.model_path,
+ args.video_path,
+ output=args.output,
+ resize=args.resize,
+ pixels=args.pixels,
+ n_frames=args.n_frames,
+ print_rate=args.print_rate,
+ display=args.display,
+ pcutoff=args.pcutoff,
+ display_radius=args.display_radius,
+ )
if __name__ == "__main__":
diff --git a/dlclive/benchmark.py b/dlclive/benchmark.py
index f6a682c..7379084 100644
--- a/dlclive/benchmark.py
+++ b/dlclive/benchmark.py
@@ -197,7 +197,7 @@ def benchmark(
)
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:
@@ -226,12 +226,14 @@ def benchmark(
if save_poses:
try:
- import pandas as pd
+ 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]).")
-
+ 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
@@ -372,7 +374,7 @@ def benchmark(
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)
@@ -381,7 +383,6 @@ def benchmark(
)
np.save(out_dlc_file, poses)
-
return inf_times, im_size, TFGPUinference, meta
@@ -452,6 +453,7 @@ def save_inf_times(
return True
+
def benchmark_videos(
model_path,
video_path,
@@ -589,6 +591,7 @@ def benchmark_videos(
output=output,
)
+
def main():
"""Provides a command line interface :function:`benchmark_videos`
"""
@@ -627,5 +630,6 @@ def main():
save_video=args.save_video,
)
+
if __name__ == "__main__":
main()
diff --git a/dlclive/display.py b/dlclive/display.py
index 534cb90..647d0ef 100644
--- a/dlclive/display.py
+++ b/dlclive/display.py
@@ -12,7 +12,7 @@
class Display(object):
- '''
+ """
Simple object to display frames with DLC labels.
Parameters
@@ -21,10 +21,9 @@ class Display(object):
string indicating the Matoplotlib colormap to use.
pcutoff : float
likelihood threshold to display points
- '''
+ """
-
- def __init__(self, cmap='bmy', radius=3, pcutoff=0.5):
+ def __init__(self, cmap="bmy", radius=3, pcutoff=0.5):
""" Constructor method
"""
@@ -34,7 +33,6 @@ def __init__(self, cmap='bmy', radius=3, pcutoff=0.5):
self.pcutoff = pcutoff
self.window = None
-
def set_display(self, im_size, bodyparts):
""" Create tkinter window to display image
@@ -52,11 +50,10 @@ def set_display(self, im_size, bodyparts):
self.lab.pack()
all_colors = getattr(cc, self.cmap)
- self.colors = all_colors[::int(len(all_colors)/bodyparts)]
-
+ self.colors = all_colors[:: int(len(all_colors) / bodyparts)]
def display_frame(self, frame, pose=None):
- '''
+ """
Display the image with DeepLabCut labels using opencv imshow
Parameters
@@ -66,7 +63,7 @@ def display_frame(self, frame, pose=None):
pose :class:`numpy.ndarray`
the pose estimated by DeepLabCut for the image
- '''
+ """
im_size = (frame.shape[1], frame.shape[0])
@@ -79,25 +76,42 @@ def display_frame(self, frame, pose=None):
draw = ImageDraw.Draw(img)
for i in range(pose.shape[0]):
- if pose[i,2] > self.pcutoff:
+ if pose[i, 2] > self.pcutoff:
try:
- x0 = pose[i,0] - self.radius if pose[i,0] - self.radius > 0 else 0
- x1 = pose[i,0] + self.radius if pose[i,0] + self.radius < im_size[1] else im_size[1]
- y0 = pose[i,1] - self.radius if pose[i,1] - self.radius > 0 else 0
- y1 = pose[i,1] + self.radius if pose[i,1] + self.radius < im_size[0] else im_size[0]
+ x0 = (
+ pose[i, 0] - self.radius
+ if pose[i, 0] - self.radius > 0
+ else 0
+ )
+ x1 = (
+ pose[i, 0] + self.radius
+ if pose[i, 0] + self.radius < im_size[1]
+ else im_size[1]
+ )
+ y0 = (
+ pose[i, 1] - self.radius
+ if pose[i, 1] - self.radius > 0
+ else 0
+ )
+ y1 = (
+ pose[i, 1] + self.radius
+ if pose[i, 1] + self.radius < im_size[0]
+ else im_size[0]
+ )
coords = [x0, y0, x1, y1]
- draw.ellipse(coords, fill=self.colors[i], outline=self.colors[i])
+ draw.ellipse(
+ coords, fill=self.colors[i], outline=self.colors[i]
+ )
except Exception as e:
print(e)
-
+
img_tk = ImageTk.PhotoImage(image=img, master=self.window)
self.lab.configure(image=img_tk)
self.window.update()
-
def destroy(self):
- '''
+ """
Destroys the opencv image window
- '''
+ """
- self.window.destroy()
\ No newline at end of file
+ self.window.destroy()
diff --git a/dlclive/dlclive.py b/dlclive/dlclive.py
index 3fdcdf1..3f1990c 100644
--- a/dlclive/dlclive.py
+++ b/dlclive/dlclive.py
@@ -257,7 +257,7 @@ def init_inference(self, frame=None, **kwargs):
# process frame
- if frame is None and (self.model_type == 'tflite'):
+ if frame is None and (self.model_type == "tflite"):
raise DLCLiveError(
"No image was passed to initialize inference. An image must be passed to the init_inference method"
)
diff --git a/dlclive/exceptions.py b/dlclive/exceptions.py
index 46ef17c..5d7a1aa 100644
--- a/dlclive/exceptions.py
+++ b/dlclive/exceptions.py
@@ -5,11 +5,14 @@
Licensed under GNU Lesser General Public License v3.0
"""
+
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 '''
- pass
\ No newline at end of file
+ """ Generic warning for incorrect use of the DLCLive class """
+
+ pass
diff --git a/dlclive/graph.py b/dlclive/graph.py
index d2d16e3..dfe3da7 100644
--- a/dlclive/graph.py
+++ b/dlclive/graph.py
@@ -10,7 +10,7 @@
def read_graph(file):
- '''
+ """
Loads the graph from a protobuf file
Parameters
@@ -22,16 +22,16 @@ def read_graph(file):
--------
graph_def :class:`tensorflow.tf.compat.v1.GraphDef`
The graph definition of the DeepLabCut model found at the object's path
- '''
+ """
- with tf.io.gfile.GFile(file, 'rb') as f:
+ with tf.io.gfile.GFile(file, "rb") as f:
graph_def = tf.compat.v1.GraphDef()
graph_def.ParseFromString(f.read())
return graph_def
def finalize_graph(graph_def):
- '''
+ """
Finalize the graph and get inputs to model
Parameters
@@ -45,7 +45,7 @@ def finalize_graph(graph_def):
The finalized graph of the DeepLabCut model
inputs :class:`tensorflow.Tensor`
Input tensor(s) for the model
- '''
+ """
graph = tf.Graph()
with graph.as_default():
@@ -56,7 +56,7 @@ def finalize_graph(graph_def):
def get_output_nodes(graph):
- '''
+ """
Get the output node names from a graph
Parameters
@@ -68,10 +68,10 @@ def get_output_nodes(graph):
--------
output : list
the output node names as a list of strings
- '''
+ """
op_names = [str(op.name) for op in graph.get_operations()]
- if 'concat_1' in op_names[-1]:
+ if "concat_1" in op_names[-1]:
output = [op_names[-1]]
else:
output = [op_names[-1], op_names[-2]]
@@ -80,7 +80,7 @@ def get_output_nodes(graph):
def get_output_tensors(graph):
- '''
+ """
Get the names of the output tensors from a graph
Parameters
@@ -92,10 +92,10 @@ def get_output_tensors(graph):
--------
output : list
the output tensor names as a list of strings
- '''
+ """
output_nodes = get_output_nodes(graph)
- output_tensor = [out+":0" for out in output_nodes]
+ output_tensor = [out + ":0" for out in output_nodes]
return output_tensor
@@ -106,7 +106,7 @@ def get_input_tensor(graph):
def extract_graph(graph, tf_config=None):
- '''
+ """
Initializes a tensorflow session with the specified graph and extracts the model's inputs and outputs
Parameters
@@ -121,7 +121,7 @@ def extract_graph(graph, tf_config=None):
a tensorflow session with the specified graph definition
outputs :class:`tensorflow.Tensor`
the output tensor(s) for the model
- '''
+ """
input_tensor = get_input_tensor(graph)
output_tensor = get_output_tensors(graph)
diff --git a/dlclive/pose.py b/dlclive/pose.py
index 3db557e..5c61f04 100644
--- a/dlclive/pose.py
+++ b/dlclive/pose.py
@@ -10,7 +10,7 @@
def extract_cnn_output(outputs, cfg):
- '''
+ """
Extract location refinement and score map from DeepLabCut network
Parameters
@@ -31,23 +31,23 @@ def extract_cnn_output(outputs, cfg):
locref : ?
location refinement
- '''
+ """
scmap = outputs[0]
scmap = np.squeeze(scmap)
locref = None
- if cfg['location_refinement']:
+ if cfg["location_refinement"]:
locref = np.squeeze(outputs[1])
shape = locref.shape
locref = np.reshape(locref, (shape[0], shape[1], -1, 2))
- locref *= cfg['locref_stdev']
- if len(scmap.shape)==2: #for single body part!
- scmap=np.expand_dims(scmap,axis=2)
+ locref *= cfg["locref_stdev"]
+ if len(scmap.shape) == 2: # for single body part!
+ scmap = np.expand_dims(scmap, axis=2)
return scmap, locref
def argmax_pose_predict(scmap, offmat, stride):
- '''
+ """
Combines score map and offsets to the final pose
Parameters
@@ -65,18 +65,17 @@ def argmax_pose_predict(scmap, offmat, stride):
--------
pose :class:`numpy.ndarray`
pose as a numpy array
- '''
+ """
num_joints = scmap.shape[2]
pose = []
for joint_idx in range(num_joints):
- maxloc = np.unravel_index(np.argmax(scmap[:, :, joint_idx]),
- scmap[:, :, joint_idx].shape)
+ maxloc = np.unravel_index(
+ np.argmax(scmap[:, :, joint_idx]), scmap[:, :, joint_idx].shape
+ )
offset = np.array(offmat[maxloc][joint_idx])[::-1]
- pos_f8 = (np.array(maxloc).astype('float') * stride + 0.5 * stride +
- offset)
- pose.append(np.hstack((pos_f8[::-1],
- [scmap[maxloc][joint_idx]])))
+ pos_f8 = np.array(maxloc).astype("float") * stride + 0.5 * stride + offset
+ pose.append(np.hstack((pos_f8[::-1], [scmap[maxloc][joint_idx]])))
return np.array(pose)
@@ -84,21 +83,21 @@ def multi_pose_predict(scmap, locref, stride, num_outputs):
Y, X = get_top_values(scmap[None], num_outputs)
Y, X = Y[:, 0], X[:, 0]
num_joints = scmap.shape[2]
- DZ=np.zeros((num_outputs,num_joints,3))
+ DZ = np.zeros((num_outputs, num_joints, 3))
for m in range(num_outputs):
for k in range(num_joints):
x = X[m, k]
y = Y[m, k]
- DZ[m,k,:2]=locref[y,x,k,:]
- DZ[m,k,2]=scmap[y,x,k]
+ DZ[m, k, :2] = locref[y, x, k, :]
+ DZ[m, k, 2] = scmap[y, x, k]
- X = X.astype('float32')*stride + .5*stride + DZ[:,:,0]
- Y = Y.astype('float32')*stride + .5*stride + DZ[:,:,1]
+ X = X.astype("float32") * stride + 0.5 * stride + DZ[:, :, 0]
+ Y = Y.astype("float32") * stride + 0.5 * stride + DZ[:, :, 1]
P = DZ[:, :, 2]
- pose = np.empty((num_joints, num_outputs*3), dtype='float32')
- pose[:,0::3] = X.T
- pose[:,1::3] = Y.T
- pose[:,2::3] = P.T
+ pose = np.empty((num_joints, num_outputs * 3), dtype="float32")
+ pose[:, 0::3] = X.T
+ pose[:, 1::3] = Y.T
+ pose[:, 2::3] = P.T
return pose
diff --git a/dlclive/processor/processor.py b/dlclive/processor/processor.py
index 3bdf702..8a52f5f 100644
--- a/dlclive/processor/processor.py
+++ b/dlclive/processor/processor.py
@@ -5,19 +5,19 @@
Licensed under GNU Lesser General Public License v3.0
"""
-'''
+"""
Default processor class. Processors must contain two methods:
i) process: takes in a pose, performs operations, and returns a pose
ii) save: saves any internal data generated by the processor (such as timestamps for commands to external hardware)
-'''
+"""
-class Processor(object):
+class Processor(object):
def __init__(self, **kwargs):
pass
def process(self, pose, **kwargs):
return pose
- def save(self, file=''):
+ def save(self, file=""):
return 0
diff --git a/dlclive/utils.py b/dlclive/utils.py
index 07176e2..4b0deaa 100644
--- a/dlclive/utils.py
+++ b/dlclive/utils.py
@@ -12,18 +12,23 @@
try:
import skimage
+
SK_IM = True
except Exception:
SK_IM = False
try:
import cv2
+
OPEN_CV = True
except Exception:
from PIL import Image
+
OPEN_CV = False
warnings.warn(
- "OpenCV is not installed. Using pillow for image processing, which is slower.", DLCLiveWarning)
+ "OpenCV is not installed. Using pillow for image processing, which is slower.",
+ DLCLiveWarning,
+ )
def convert_to_ubyte(frame):
@@ -95,7 +100,9 @@ def img_to_rgb(frame):
else:
warnings.warn(
- f"Image has {frame.ndim} dimensions. Must be 2 or 3 dimensions to convert to RGB", DLCLiveWarning)
+ f"Image has {frame.ndim} dimensions. Must be 2 or 3 dimensions to convert to RGB",
+ DLCLiveWarning,
+ )
return frame
@@ -115,7 +122,7 @@ def gray_to_rgb(frame):
else:
img = Image.fromarray(frame)
- img = img.convert('RGB')
+ img = img.convert("RGB")
return np.asarray(img)
@@ -135,7 +142,7 @@ def bgr_to_rgb(frame):
else:
img = Image.fromarray(frame)
- img = img.convert('RGB')
+ img = img.convert("RGB")
return np.asarray(img)
@@ -177,14 +184,15 @@ def _img_as_ubyte_np(frame):
elif np.issubdtype(im_type, np.integer):
im_type_info = np.iinfo(im_type)
- frame *= 255/im_type_info.max
+ frame *= 255 / im_type_info.max
frame[frame < 0] = 0
return frame.astype(np.uint8)
else:
raise TypeError(
- "image of type {} could not be converted to ubyte".format(im_type))
+ "image of type {} could not be converted to ubyte".format(im_type)
+ )
def decode_fourcc(cc):
diff --git a/dlclive/version.py b/dlclive/version.py
index 0b4a5dc..a28c9db 100644
--- a/dlclive/version.py
+++ b/dlclive/version.py
@@ -7,5 +7,5 @@
"""
-__version__ = '0.0.1'
+__version__ = "0.0.1"
VERSION = __version__
diff --git a/example_processors/DogJumpLED/izzy_jump.py b/example_processors/DogJumpLED/izzy_jump.py
index d843e84..f24356d 100644
--- a/example_processors/DogJumpLED/izzy_jump.py
+++ b/example_processors/DogJumpLED/izzy_jump.py
@@ -141,4 +141,3 @@ def process(self, pose, **kwargs):
def save(self, filename):
return IzzyJump.save(self, filename)
-
diff --git a/example_processors/DogJumpLED/izzy_jump_offline.py b/example_processors/DogJumpLED/izzy_jump_offline.py
index a92c3b8..9e574ab 100644
--- a/example_processors/DogJumpLED/izzy_jump_offline.py
+++ b/example_processors/DogJumpLED/izzy_jump_offline.py
@@ -121,4 +121,3 @@ def process(self, pose, **kwargs):
def save(self, filename):
return IzzyJumpOffline.save(self, filename)
-
diff --git a/setup.py b/setup.py
index 2257ed6..02b6ebd 100755
--- a/setup.py
+++ b/setup.py
@@ -26,7 +26,9 @@
]
if "tegra" in platform.platform():
- warnings.warn("Not installing the following packages:\nopencv-python\ntensorflow\npandas\ntables\nPlease follow instructions on github to install opencv and tensorflow. If you want to use the benchmark_videos function to save poses from a video, then please install pandas and tables (pip install pandas tables)")
+ warnings.warn(
+ "Not installing the following packages:\nopencv-python\ntensorflow\npandas\ntables\nPlease follow instructions on github to install opencv and tensorflow. If you want to use the benchmark_videos function to save poses from a video, then please install pandas and tables (pip install pandas tables)"
+ )
else:
install_requires.append("opencv-python")
install_requires.append("tensorflow==1.13.1")