|
import urllib.request |
|
import tarfile |
|
from tqdm import tqdm |
|
import os |
|
import yaml |
|
from ruamel.yaml import YAML |
|
|
|
def read_plainconfig(configname): |
|
if not os.path.exists(configname): |
|
raise FileNotFoundError( |
|
f"Config {configname} is not found. Please make sure that the file exists." |
|
) |
|
with open(configname) as file: |
|
return YAML().load(file) |
|
|
|
def DownloadModel(modelname, target_dir): |
|
""" |
|
Downloads a DeepLabCut Model Zoo Project |
|
""" |
|
|
|
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 |
|
|
|
neturls = read_plainconfig("./model/pretrained_model_urls.yaml") |
|
|
|
if modelname in neturls.keys(): |
|
url = neturls[modelname] |
|
print(url) |
|
response = urllib.request.urlopen(url) |
|
print( |
|
"Downloading the model 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)) |
|
else: |
|
models = [ |
|
fn |
|
for fn in neturls.keys() |
|
if "resnet_" not in fn and "mobilenet_" not in fn |
|
] |
|
print("Model does not exist: ", modelname) |
|
print("Pick one of the following: ", models) |
|
return target_dir |
|
|
