Merge pull request #95 from DeepLabCut/mwm/fixModelDownload

Update check_install.py

Author: AlexEMG

CITATION.cff

@@ -52,4 +52,4 @@ authors:
 date-released: 2020-08-05
 doi: "10.7554/eLife.61909"
 license: "AGPL-3.0-or-later"
-version:  "1.0.2"
+version:  "1.0.3"

dlclive/check_install/check_install.py

@@ -6,7 +6,6 @@
 """
 
 
-import os
 import sys
 import shutil
 import warnings
@@ -15,7 +14,9 @@
 import urllib.request
 import argparse
 from pathlib import Path
-import tarfile
+from dlclibrary.dlcmodelzoo.modelzoo_download import (
+    download_huggingface_model,
+)
 
 
 def urllib_pbar(count, blockSize, totalSize):
@@ -25,6 +26,7 @@ def urllib_pbar(count, blockSize, totalSize):
     sys.stdout.write("\b"*len(outstr))
     sys.stdout.flush()
 
+
 def main(display:bool=None):
     parser = argparse.ArgumentParser(
         description="Test DLC-Live installation by downloading and evaluating a demo DLC project!")
@@ -43,31 +45,23 @@ def main(display:bool=None):
     tmp_dir.mkdir(mode=0o775,exist_ok=True)
 
     video_file = str(tmp_dir / 'dog_clip.avi')
-    model_tarball = tmp_dir / 'DLC_Dog_resnet_50_iteration-0_shuffle-0.tar.gz'
-    model_dir = model_tarball.with_suffix('').with_suffix('') # remove two suffixes (tar.gz)
-
+    model_dir = tmp_dir / 'DLC_Dog_resnet_50_iteration-0_shuffle-0'
 
     # download dog test video from github:
     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)
 
     # download exported dog model from DeepLabCut Model Zoo
-    if Path(model_tarball).exists():
-        print('Tarball already downloaded, using cached version')
+    if Path(model_dir / 'snapshot-75000.pb').exists():
+        print('Model already downloaded, using cached version')
     else:
         print("Downloading full_dog model from the DeepLabCut Model Zoo...")
-        model_url = "http://deeplabcut.rowland.harvard.edu/models/DLC_Dog_resnet_50_iteration-0_shuffle-0.tar.gz"
-        urllib.request.urlretrieve(model_url, str(model_tarball), reporthook=urllib_pbar)
-
-    print('Untarring compressed model')
-    model_file = tarfile.open(str(model_tarball))
-    model_file.extractall(str(model_dir.parent))
-    model_file.close()
+        download_huggingface_model("full_dog", model_dir)
 
     # assert these things exist so we can give informative error messages
     assert Path(video_file).exists()
-    assert Path(model_dir).exists() and Path(model_dir).is_dir()
+    assert Path(model_dir / 'snapshot-75000.pb').exists()
 
     # run benchmark videos
     print("\n Running inference...\n")

dlclive/version.py

@@ -7,5 +7,5 @@
 """
 
 
-__version__ = "1.0.2"
+__version__ = "1.0.3"
 VERSION = __version__