|
7 | 7 |
|
8 | 8 |
|
9 | 9 | import os |
| 10 | +import sys |
10 | 11 | import shutil |
| 12 | +import warnings |
| 13 | + |
11 | 14 | from dlclive import benchmark_videos |
12 | 15 | import urllib.request |
| 16 | +import argparse |
| 17 | +from pathlib import Path |
| 18 | +import tarfile |
| 19 | + |
| 20 | + |
| 21 | +def urllib_pbar(count, blockSize, totalSize): |
| 22 | + percent = int(count * blockSize * 100 / totalSize) |
| 23 | + outstr = f"{round(percent)}%" |
| 24 | + sys.stdout.write(outstr) |
| 25 | + sys.stdout.write("\b"*len(outstr)) |
| 26 | + sys.stdout.flush() |
13 | 27 |
|
| 28 | +def main(display:bool=None): |
| 29 | + parser = argparse.ArgumentParser( |
| 30 | + description="Test DLC-Live installation by downloading and evaluating a demo DLC project!") |
| 31 | + parser.add_argument('--nodisplay', action='store_false', help="Run the test without displaying tracking") |
| 32 | + args = parser.parse_args() |
14 | 33 |
|
15 | | -def main(): |
| 34 | + if display is None: |
| 35 | + display = args.nodisplay |
| 36 | + |
| 37 | + if not display: |
| 38 | + print('Running without displaying video') |
16 | 39 |
|
17 | 40 | # make temporary directory in $HOME |
18 | 41 | print("\nCreating temporary directory...\n") |
19 | | - home = os.path.expanduser("~") |
20 | | - tmp_dir = os.path.normpath(f"{home}/dlc-live-tmp") |
21 | | - os.makedirs(tmp_dir, exist_ok=True) |
22 | | - os.chdir(tmp_dir) |
| 42 | + tmp_dir = Path().home() / 'dlc-live-tmp' |
| 43 | + tmp_dir.mkdir(mode=0o775,exist_ok=True) |
| 44 | + |
| 45 | + video_file = str(tmp_dir / 'dog_clip.avi') |
| 46 | + model_tarball = tmp_dir / 'DLC_Dog_resnet_50_iteration-0_shuffle-0.tar.gz' |
| 47 | + model_dir = model_tarball.with_suffix('').with_suffix('') # remove two suffixes (tar.gz) |
| 48 | + |
23 | 49 |
|
24 | 50 | # download dog test video from github: |
| 51 | + print(f"Downloading Video to {video_file}") |
25 | 52 | url_link = "https://github.com/DeepLabCut/DeepLabCut-live/blob/master/check_install/dog_clip.avi?raw=True" |
26 | | - urllib.request.urlretrieve(url_link, "dog_clip.avi") |
27 | | - video_file = os.path.join(url_link, "dog_clip.avi") |
| 53 | + urllib.request.urlretrieve(url_link, video_file, reporthook=urllib_pbar) |
28 | 54 |
|
29 | 55 | # download exported dog model from DeepLabCut Model Zoo |
30 | | - print("Downloading full_dog model from the DeepLabCut Model Zoo...") |
31 | | - model_url = "http://deeplabcut.rowland.harvard.edu/models/DLC_Dog_resnet_50_iteration-0_shuffle-0.tar.gz" |
32 | | - os.system(f"curl {model_url} | tar xvz") |
| 56 | + if Path(model_tarball).exists(): |
| 57 | + print('Tarball already downloaded, using cached version') |
| 58 | + else: |
| 59 | + print("Downloading full_dog model from the DeepLabCut Model Zoo...") |
| 60 | + model_url = "http://deeplabcut.rowland.harvard.edu/models/DLC_Dog_resnet_50_iteration-0_shuffle-0.tar.gz" |
| 61 | + urllib.request.urlretrieve(model_url, str(model_tarball), reporthook=urllib_pbar) |
| 62 | + |
| 63 | + print('Untarring compressed model') |
| 64 | + model_file = tarfile.open(str(model_tarball)) |
| 65 | + model_file.extractall(str(model_dir.parent)) |
| 66 | + model_file.close() |
| 67 | + |
| 68 | + # assert these things exist so we can give informative error messages |
| 69 | + assert Path(video_file).exists() |
| 70 | + assert Path(model_dir).exists() and Path(model_dir).is_dir() |
33 | 71 |
|
34 | 72 | # run benchmark videos |
35 | 73 | print("\n Running inference...\n") |
36 | | - model_dir = "DLC_Dog_resnet_50_iteration-0_shuffle-0" |
37 | | - print(video_file) |
38 | | - benchmark_videos(model_dir, video_file, display=True, resize=0.5, pcutoff=0.25) |
| 74 | + # model_dir = "DLC_Dog_resnet_50_iteration-0_shuffle-0" |
| 75 | + # print(video_file) |
| 76 | + benchmark_videos(str(model_dir), video_file, display=display, resize=0.5, pcutoff=0.25) |
39 | 77 |
|
40 | 78 | # deleting temporary files |
41 | 79 | print("\n Deleting temporary files...\n") |
42 | | - shutil.rmtree(tmp_dir) |
| 80 | + try: |
| 81 | + shutil.rmtree(tmp_dir) |
| 82 | + except PermissionError: |
| 83 | + warnings.warn(f'Could not delete temporary directory {str(tmp_dir)} due to a permissions error, but otherwise dlc-live seems to be working fine!') |
43 | 84 |
|
44 | 85 | print("\nDone!\n") |
45 | 86 |
|
46 | 87 |
|
47 | 88 | if __name__ == "__main__": |
48 | | - main() |
| 89 | + |
| 90 | + |
| 91 | + display = args.nodisplay |
| 92 | + |
| 93 | + |
| 94 | + main(display=args.nodisplay) |
0 commit comments