From af4cdad04e97503f3977b3cb44ac4e7d85c6bacc Mon Sep 17 00:00:00 2001 From: Joshua Chu Date: Sun, 5 May 2019 14:05:09 -0500 Subject: [PATCH 1/4] Download pre-trained models in a more robust way --- .../trainingsetmanipulation.py | 41 ++++++++++--------- .../models/pretrained/models.txt | 2 + setup.py | 1 - 3 files changed, 23 insertions(+), 21 deletions(-) create mode 100644 deeplabcut/pose_estimation_tensorflow/models/pretrained/models.txt diff --git a/deeplabcut/generate_training_dataset/trainingsetmanipulation.py b/deeplabcut/generate_training_dataset/trainingsetmanipulation.py index 81346812a..ad458fe31 100644 --- a/deeplabcut/generate_training_dataset/trainingsetmanipulation.py +++ b/deeplabcut/generate_training_dataset/trainingsetmanipulation.py @@ -22,7 +22,9 @@ mpl.use('TkAgg') import matplotlib.pyplot as plt - +import urllib +import tarfile +from io import BytesIO #if os.environ.get('DLClight', default=False) == 'True': # mpl.use('AGG') #anti-grain geometry engine #https://matplotlib.org/faq/usage_faq.html @@ -265,7 +267,7 @@ def MakeLabeledPlots(folder,DataCombined,cfg,Labels,Colorscheme,cc,scale): tmpfolder = str(folder) + '_labeled' auxiliaryfunctions.attempttomakefolder(tmpfolder) for index, imagename in enumerate(DataCombined.index.values): - image = io.imread(os.path.join(cfg['project_path'],imagename)) + image = skio.imread(os.path.join(cfg['project_path'],imagename)) plt.axis('off') if np.ndim(image)==2: @@ -434,11 +436,11 @@ def mergeandsplit(config,trainindex=0,uniform=True,windows2linux=False): >>> deeplabcut.create_training_dataset(config,Shuffles=[3],trainIndexes=trainIndexes,testIndexes=testIndexes) To freeze a (uniform) split: - >>> trainIndices, testIndices=deeplabcut.mergeandsplit(config,trainindex=0,uniform=True) + >>> trainIndexes, testIndexes=deeplabcut.mergeandsplit(config,trainindex=0,uniform=True) You can then create two model instances that have the identical trainingset. Thereby you can assess the role of various parameters on the performance of DLC. - >>> deeplabcut.create_training_dataset(config,Shuffles=[0],trainIndices=trainIndices,testIndices=testIndices) - >>> deeplabcut.create_training_dataset(config,Shuffles=[1],trainIndices=trainIndices,testIndices=testIndices) + >>> deeplabcut.create_training_dataset(config,Shuffles=[0],trainIndexes=trainIndexes,testIndexes=testIndexes) + >>> deeplabcut.create_training_dataset(config,Shuffles=[1],trainIndexes=trainIndexes,testIndexes=testIndexes) -------- """ @@ -479,7 +481,7 @@ def mergeandsplit(config,trainindex=0,uniform=True,windows2linux=False): return trainIndexes, testIndexes -def create_training_dataset(config,num_shuffles=1,Shuffles=None,windows2linux=False,trainIndices=None,testIndices=None): +def create_training_dataset(config,num_shuffles=1,Shuffles=None,windows2linux=False,trainIndexes=None,testIndexes=None): """ Creates a training dataset. Labels from all the extracted frames are merged into a single .h5 file.\n Only the videos included in the config file are used to create this dataset.\n @@ -501,9 +503,6 @@ def create_training_dataset(config,num_shuffles=1,Shuffles=None,windows2linux=Fa The annotation files contain path formated according to your operating system. If you label on windows but train & evaluate on a unix system (e.g. ubunt, colab, Mac) set this variable to True to convert the paths. - trainIndices and testIndices: list of indices for traininng and testing. Use mergeandsplit(config,trainindex=0,uniform=True,windows2linux=False) to create them - See help for deeplabcut.mergeandsplit? - Example -------- >>> deeplabcut.create_training_dataset('/analysis/project/reaching-task/config.yaml',num_shuffles=1) @@ -528,12 +527,13 @@ def create_training_dataset(config,num_shuffles=1,Shuffles=None,windows2linux=Fa Data = Data[scorer] #extract labeled data #set model type. we will allow more in the future. + parent_path = os.path.dirname(deeplabcut.__file__) if cfg['resnet']==50: net_type ='resnet_'+str(cfg['resnet']) - resnet_path = str(Path(deeplabcut.__file__).parents[0] / 'pose_estimation_tensorflow/models/pretrained/resnet_v1_50.ckpt') + resnet_path = os.path.join(parent_path, 'pose_estimation_tensorflow/models/pretrained/resnet_v1_50.ckpt') elif cfg['resnet']==101: net_type ='resnet_'+str(cfg['resnet']) - resnet_path = str(Path(deeplabcut.__file__).parents[0] / 'pose_estimation_tensorflow/models/pretrained/resnet_v1_101.ckpt') + resnet_path = os.path.join(parent_path, 'pose_estimation_tensorflow/models/pretrained/resnet_v1_101.ckpt') else: print("Currently only ResNet 50 or 101 supported, please change 'resnet' entry in config.yaml!") num_shuffles=-1 #thus the loop below is empty... @@ -542,11 +542,14 @@ def create_training_dataset(config,num_shuffles=1,Shuffles=None,windows2linux=Fa """ Downloads the ImageNet pretrained weights for ResNet. """ - start = os.getcwd() - os.chdir(str(Path(resnet_path).parents[0])) - print("Downloading the pretrained model (ResNets)....") - subprocess.call("download.sh", shell=True) - os.chdir(start) + target_dir = os.path.dirname(resnet_path) + with open(os.path.join(target_dir, 'models.txt'), 'r') as f: + for line in f: + url = line.rstrip() + print("Downloading a pretrained model (ResNet) from {}....".format(url)) + response = urllib.request.urlopen(url) + with tarfile.open(fileobj=BytesIO(response.read()), mode='r:gz') as tar: + tar.extractall(path=target_dir) if Shuffles==None: Shuffles=range(1,num_shuffles+1,1) @@ -558,11 +561,9 @@ def create_training_dataset(config,num_shuffles=1,Shuffles=None,windows2linux=Fa for shuffle in Shuffles: # Creating shuffles starting from 1 for trainFraction in TrainingFraction: #trainIndexes, testIndexes = SplitTrials(range(len(Data.index)), trainFraction) - if trainIndices is None and testIndices is None: + if trainIndexes is None and testIndexes is None: trainIndexes, testIndexes = SplitTrials(range(len(Data.index)), trainFraction) - else: # set to passed values... - trainIndexes=trainIndices - testIndexes=testIndices + else: print("You passed a split with the following fraction:", len(trainIndexes)*1./(len(testIndexes)+len(trainIndexes))*100) #################################################### diff --git a/deeplabcut/pose_estimation_tensorflow/models/pretrained/models.txt b/deeplabcut/pose_estimation_tensorflow/models/pretrained/models.txt new file mode 100644 index 000000000..f77450faa --- /dev/null +++ b/deeplabcut/pose_estimation_tensorflow/models/pretrained/models.txt @@ -0,0 +1,2 @@ +http://download.tensorflow.org/models/resnet_v1_50_2016_08_28.tar.gz +http://download.tensorflow.org/models/resnet_v1_101_2016_08_28.tar.gz diff --git a/setup.py b/setup.py index 13115813f..2cacf7b4a 100644 --- a/setup.py +++ b/setup.py @@ -30,7 +30,6 @@ 'ruamel.yaml==0.15','setuptools','scikit-image~=0.14.0','scikit-learn~=0.19.2', 'scipy~=1.1.0','six==1.11.0','statsmodels==0.9.0','tables', 'tqdm>4','wheel==0.31.1'], - scripts=['deeplabcut/pose_estimation_tensorflow/models/pretrained/download.sh'], packages=setuptools.find_packages(), data_files=[('deeplabcut',['deeplabcut/pose_cfg.yaml'])], include_package_data=True, From b1fa32ad9c784c196aee2152af1caafe5312e89e Mon Sep 17 00:00:00 2001 From: Joshua Chu Date: Sun, 5 May 2019 14:08:14 -0500 Subject: [PATCH 2/4] Minor change --- deeplabcut/generate_training_dataset/trainingsetmanipulation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deeplabcut/generate_training_dataset/trainingsetmanipulation.py b/deeplabcut/generate_training_dataset/trainingsetmanipulation.py index ad458fe31..921b714f5 100644 --- a/deeplabcut/generate_training_dataset/trainingsetmanipulation.py +++ b/deeplabcut/generate_training_dataset/trainingsetmanipulation.py @@ -267,7 +267,7 @@ def MakeLabeledPlots(folder,DataCombined,cfg,Labels,Colorscheme,cc,scale): tmpfolder = str(folder) + '_labeled' auxiliaryfunctions.attempttomakefolder(tmpfolder) for index, imagename in enumerate(DataCombined.index.values): - image = skio.imread(os.path.join(cfg['project_path'],imagename)) + image = io.imread(os.path.join(cfg['project_path'],imagename)) plt.axis('off') if np.ndim(image)==2: From 3a4cd0b429f2afe5461eccb718d1e7125c3c0442 Mon Sep 17 00:00:00 2001 From: Joshua Chu Date: Sun, 5 May 2019 15:18:04 -0500 Subject: [PATCH 3/4] Platform indepedent paths --- .../trainingsetmanipulation.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/deeplabcut/generate_training_dataset/trainingsetmanipulation.py b/deeplabcut/generate_training_dataset/trainingsetmanipulation.py index 921b714f5..b472f0267 100644 --- a/deeplabcut/generate_training_dataset/trainingsetmanipulation.py +++ b/deeplabcut/generate_training_dataset/trainingsetmanipulation.py @@ -527,23 +527,23 @@ def create_training_dataset(config,num_shuffles=1,Shuffles=None,windows2linux=Fa Data = Data[scorer] #extract labeled data #set model type. we will allow more in the future. - parent_path = os.path.dirname(deeplabcut.__file__) + parent_path = Path(os.path.dirname(deeplabcut.__file__)) if cfg['resnet']==50: net_type ='resnet_'+str(cfg['resnet']) - resnet_path = os.path.join(parent_path, 'pose_estimation_tensorflow/models/pretrained/resnet_v1_50.ckpt') + resnet_path = parent_path / 'pose_estimation_tensorflow/models/pretrained/resnet_v1_50.ckpt' elif cfg['resnet']==101: net_type ='resnet_'+str(cfg['resnet']) - resnet_path = os.path.join(parent_path, 'pose_estimation_tensorflow/models/pretrained/resnet_v1_101.ckpt') + resnet_path = parent_path / 'pose_estimation_tensorflow/models/pretrained/resnet_v1_101.ckpt' else: print("Currently only ResNet 50 or 101 supported, please change 'resnet' entry in config.yaml!") num_shuffles=-1 #thus the loop below is empty... - if not Path(resnet_path).is_file(): + if not resnet_path.is_file(): """ Downloads the ImageNet pretrained weights for ResNet. """ - target_dir = os.path.dirname(resnet_path) - with open(os.path.join(target_dir, 'models.txt'), 'r') as f: + target_dir = resnet_path.parents[0] + with open(target_dir / 'models.txt', 'r') as f: for line in f: url = line.rstrip() print("Downloading a pretrained model (ResNet) from {}....".format(url)) From d6bbe488ebfe811f987fb022fc106be71fcf4cd5 Mon Sep 17 00:00:00 2001 From: AlexEMG Date: Thu, 6 Jun 2019 17:26:55 -0400 Subject: [PATCH 4/4] Streamlined --- .../trainingsetmanipulation.py | 57 ++++--------------- .../models/pretrained/download.sh | 1 + .../models/pretrained/models.txt | 2 - .../pretrained/pretrained_model_urls.yaml | 2 + deeplabcut/utils/__init__.py | 2 + deeplabcut/utils/auxfun_models.py | 46 +++++++++++++++ setup.py | 3 +- 7 files changed, 65 insertions(+), 48 deletions(-) delete mode 100644 deeplabcut/pose_estimation_tensorflow/models/pretrained/models.txt create mode 100644 deeplabcut/pose_estimation_tensorflow/models/pretrained/pretrained_model_urls.yaml create mode 100644 deeplabcut/utils/auxfun_models.py diff --git a/deeplabcut/generate_training_dataset/trainingsetmanipulation.py b/deeplabcut/generate_training_dataset/trainingsetmanipulation.py index b472f0267..03391a791 100644 --- a/deeplabcut/generate_training_dataset/trainingsetmanipulation.py +++ b/deeplabcut/generate_training_dataset/trainingsetmanipulation.py @@ -21,22 +21,11 @@ else: mpl.use('TkAgg') import matplotlib.pyplot as plt - -import urllib -import tarfile -from io import BytesIO - -#if os.environ.get('DLClight', default=False) == 'True': -# mpl.use('AGG') #anti-grain geometry engine #https://matplotlib.org/faq/usage_faq.html -# pass -#else: -# mpl.use('TkAgg') -#import matplotlib.pyplot as plt from skimage import io import yaml from deeplabcut import DEBUG -from deeplabcut.utils import auxiliaryfunctions, conversioncode +from deeplabcut.utils import auxiliaryfunctions, conversioncode, auxfun_models #matplotlib.use('Agg') @@ -444,7 +433,6 @@ def mergeandsplit(config,trainindex=0,uniform=True,windows2linux=False): -------- """ - # Loading metadata from config file: cfg = auxiliaryfunctions.read_config(config) scorer = cfg['scorer'] @@ -512,9 +500,7 @@ def create_training_dataset(config,num_shuffles=1,Shuffles=None,windows2linux=Fa """ from skimage import io import scipy.io as sio - import deeplabcut - import subprocess - + # Loading metadata from config file: cfg = auxiliaryfunctions.read_config(config) scorer = cfg['scorer'] @@ -525,32 +511,16 @@ def create_training_dataset(config,num_shuffles=1,Shuffles=None,windows2linux=Fa Data = merge_annotateddatasets(cfg,project_path,Path(os.path.join(project_path,trainingsetfolder)),windows2linux) Data = Data[scorer] #extract labeled data - - #set model type. we will allow more in the future. + + + #loading & linking pretrained models + net_type ='resnet_'+str(cfg['resnet']) + import deeplabcut parent_path = Path(os.path.dirname(deeplabcut.__file__)) - if cfg['resnet']==50: - net_type ='resnet_'+str(cfg['resnet']) - resnet_path = parent_path / 'pose_estimation_tensorflow/models/pretrained/resnet_v1_50.ckpt' - elif cfg['resnet']==101: - net_type ='resnet_'+str(cfg['resnet']) - resnet_path = parent_path / 'pose_estimation_tensorflow/models/pretrained/resnet_v1_101.ckpt' - else: - print("Currently only ResNet 50 or 101 supported, please change 'resnet' entry in config.yaml!") - num_shuffles=-1 #thus the loop below is empty... - - if not resnet_path.is_file(): - """ - Downloads the ImageNet pretrained weights for ResNet. - """ - target_dir = resnet_path.parents[0] - with open(target_dir / 'models.txt', 'r') as f: - for line in f: - url = line.rstrip() - print("Downloading a pretrained model (ResNet) from {}....".format(url)) - response = urllib.request.urlopen(url) - with tarfile.open(fileobj=BytesIO(response.read()), mode='r:gz') as tar: - tar.extractall(path=target_dir) - + defaultconfigfile = str(parent_path / 'pose_cfg.yaml') + + model_path,num_shuffles=auxfun_models.Check4weights(net_type,parent_path,num_shuffles) + if Shuffles==None: Shuffles=range(1,num_shuffles+1,1) else: @@ -650,13 +620,10 @@ def create_training_dataset(config,num_shuffles=1,Shuffles=None,windows2linux=Fa "num_joints": len(bodyparts), "all_joints": [[i] for i in range(len(bodyparts))], "all_joints_names": [str(bpt) for bpt in bodyparts], - "init_weights": resnet_path, + "init_weights": model_path, "project_path": str(cfg['project_path']), "net_type": net_type } - - defaultconfigfile = str(Path(deeplabcut.__file__).parents[0] / 'pose_cfg.yaml') - trainingdata = MakeTrain_pose_yaml(items2change,path_train_config,defaultconfigfile) keys2save = [ "dataset", "num_joints", "all_joints", "all_joints_names", diff --git a/deeplabcut/pose_estimation_tensorflow/models/pretrained/download.sh b/deeplabcut/pose_estimation_tensorflow/models/pretrained/download.sh index 6c76636a0..0d5f31ae1 100644 --- a/deeplabcut/pose_estimation_tensorflow/models/pretrained/download.sh +++ b/deeplabcut/pose_estimation_tensorflow/models/pretrained/download.sh @@ -1,3 +1,4 @@ +# legacy. #!/bin/sh curl http://download.tensorflow.org/models/resnet_v1_50_2016_08_28.tar.gz | tar xvz diff --git a/deeplabcut/pose_estimation_tensorflow/models/pretrained/models.txt b/deeplabcut/pose_estimation_tensorflow/models/pretrained/models.txt deleted file mode 100644 index f77450faa..000000000 --- a/deeplabcut/pose_estimation_tensorflow/models/pretrained/models.txt +++ /dev/null @@ -1,2 +0,0 @@ -http://download.tensorflow.org/models/resnet_v1_50_2016_08_28.tar.gz -http://download.tensorflow.org/models/resnet_v1_101_2016_08_28.tar.gz diff --git a/deeplabcut/pose_estimation_tensorflow/models/pretrained/pretrained_model_urls.yaml b/deeplabcut/pose_estimation_tensorflow/models/pretrained/pretrained_model_urls.yaml new file mode 100644 index 000000000..a76f10e05 --- /dev/null +++ b/deeplabcut/pose_estimation_tensorflow/models/pretrained/pretrained_model_urls.yaml @@ -0,0 +1,2 @@ +resnet_50: http://download.tensorflow.org/models/resnet_v1_50_2016_08_28.tar.gz +resnet_101: http://download.tensorflow.org/models/resnet_v1_101_2016_08_28.tar.gz diff --git a/deeplabcut/utils/__init__.py b/deeplabcut/utils/__init__.py index fd99ccd10..dc2d37f93 100644 --- a/deeplabcut/utils/__init__.py +++ b/deeplabcut/utils/__init__.py @@ -1,5 +1,7 @@ from deeplabcut.utils.make_labeled_video import * from deeplabcut.utils.auxiliaryfunctions import * +from deeplabcut.utils.auxfun_models import * + from deeplabcut.utils.video_processor import * from deeplabcut.utils.plotting import * diff --git a/deeplabcut/utils/auxfun_models.py b/deeplabcut/utils/auxfun_models.py new file mode 100644 index 000000000..e9f1167b8 --- /dev/null +++ b/deeplabcut/utils/auxfun_models.py @@ -0,0 +1,46 @@ +""" +DeepLabCut Toolbox +https://github.com/AlexEMG/DeepLabCut +A Mathis, alexander.mathis@bethgelab.org +M Mathis, mackenzie@post.harvard.edu + +""" + +from deeplabcut.utils import auxiliaryfunctions + +def Check4weights(modeltype,parent_path,num_shuffles): + ''' gets local path to network weights and checks if they are present. If not, downloads them from tensorflow.org ''' + if 'resnet_50' == modeltype: + model_path = parent_path / 'pose_estimation_tensorflow/models/pretrained/resnet_v1_50.ckpt' + elif 'resnet_101' == modeltype: + model_path = parent_path / 'pose_estimation_tensorflow/models/pretrained/resnet_v1_101.ckpt' + else: + print("Currently only ResNet 50 or 101 supported, please change 'resnet' entry in config.yaml!") + num_shuffles=-1 #thus the loop below is empty... + model_path=parent_path + + if num_shuffles>0: + if not model_path.is_file(): + Downloadweights(modeltype,model_path) + + return str(model_path),num_shuffles + +def Downloadweights(modeltype,model_path): + """ + Downloads the ImageNet pretrained weights for ResNet. + """ + + import urllib + import tarfile + from io import BytesIO + + target_dir = model_path.parents[0] + neturls=auxiliaryfunctions.read_plainconfig(target_dir / 'pretrained_model_urls.yaml') + try: + url = neturls[modeltype] + print("Downloading a ImageNet-pretrained model from {}....".format(url)) + response = urllib.request.urlopen(url) + with tarfile.open(fileobj=BytesIO(response.read()), mode='r:gz') as tar: + tar.extractall(path=target_dir) + except KeyError: + print("Model does not exist", modeltype) \ No newline at end of file diff --git a/setup.py b/setup.py index 2cacf7b4a..10de813b8 100644 --- a/setup.py +++ b/setup.py @@ -30,8 +30,9 @@ 'ruamel.yaml==0.15','setuptools','scikit-image~=0.14.0','scikit-learn~=0.19.2', 'scipy~=1.1.0','six==1.11.0','statsmodels==0.9.0','tables', 'tqdm>4','wheel==0.31.1'], + scripts=['deeplabcut/pose_estimation_tensorflow/models/pretrained/download.sh'], packages=setuptools.find_packages(), - data_files=[('deeplabcut',['deeplabcut/pose_cfg.yaml'])], + data_files=[('deeplabcut',['deeplabcut/pose_cfg.yaml','deeplabcut/pose_estimation_tensorflow/models/pretrained/pretrained_model_urls.yaml'])], include_package_data=True, classifiers=( "Programming Language :: Python :: 3",