diff --git a/ikt/constants.py b/ikt/constants.py new file mode 100644 index 0000000..64d1110 --- /dev/null +++ b/ikt/constants.py @@ -0,0 +1,21 @@ +# flake8: noqa +pre_trained_models = { + # pretrained-yolov3.h5 : A pre-trained YOLOv3 model for transfer learning when training new detection models + "hololens.h5": "https://github.com/OlafenwaMoses/ImageAI/releases/download/essential-v4/hololens-ex-60--loss-2.76.h5", + # hololens-ex-60--loss-2.76.h5 : A YOLOv3 model trained with ImageAI on the Hololens dataset + "yolov3.h5": "https://github.com/OlafenwaMoses/ImageAI/releases/download/essential-v4/pretrained-yolov3.h5", + # The configuration JSON file for performing detection in images and video using the trained YOLOv3 model for Hololens. + "yolov3_detection_config.json": "https://github.com/OlafenwaMoses/ImageAI/releases/download/essential-v4/detection_config.json", + "idenprof_densenet.h5": "https://github.com/OlafenwaMoses/ImageAI/releases/download/models-v3/idenprof_densenet-0.763500.h5", + "idenprof_full_resnet_ex-001_acc-0.h5": "https://github.com/OlafenwaMoses/ImageAI/releases/download/models-v3/idenprof_full_resnet_ex-001_acc-0.119792.h5", + "idenprof_inception_0.h5": "https://github.com/OlafenwaMoses/ImageAI/releases/download/models-v3/idenprof_inception_0.719500.h5", + "idenprof_resnet.h5": "https://github.com/OlafenwaMoses/ImageAI/releases/download/models-v3/idenprof_resnet.h5", + "resnet_model_ex-020_acc-0.651714.h5": "https://github.com/OlafenwaMoses/ImageAI/releases/download/1.0.1/resnet_model_ex-020_acc-0.651714.h5", + "DenseNet-BC-121-32.h5": "https://github.com/OlafenwaMoses/ImageAI/releases/download/1.0/DenseNet-BC-121-32.h5", + "inception_v3_weights_tf_dim_ordering_tf_kernels.h5": "https://github.com/OlafenwaMoses/ImageAI/releases/download/1.0/inception_v3_weights_tf_dim_ordering_tf_kernels.h5", + "resnet50_coco_best_v2.0.1.h5": "https://github.com/OlafenwaMoses/ImageAI/releases/download/1.0/resnet50_coco_best_v2.0.1.h5", + "resnet50_weights_tf_dim_ordering_tf_kernels.h5": "https://github.com/OlafenwaMoses/ImageAI/releases/download/1.0/resnet50_weights_tf_dim_ordering_tf_kernels.h5", + "squeezenet_weights_tf_dim_ordering_tf_kernels.h5": "https://github.com/OlafenwaMoses/ImageAI/releases/download/1.0/squeezenet_weights_tf_dim_ordering_tf_kernels.h5", + "yolo-tiny.h5": "https://github.com/OlafenwaMoses/ImageAI/releases/download/1.0/yolo-tiny.h5", + "yolo.h5": "https://github.com/OlafenwaMoses/ImageAI/releases/download/1.0/yolo.h5", +} diff --git a/ikt/downloader.py b/ikt/downloader.py new file mode 100644 index 0000000..a43ab6e --- /dev/null +++ b/ikt/downloader.py @@ -0,0 +1,44 @@ +from pathlib import Path + +import requests + +from ikt import logger +from ikt.constants import pre_trained_models + +storage = Path("pre_trained_models") + + +def download_file(filename, url, storage_dir): + """Downloads file to storage directory.""" + logger.debug("Downloading [%s].", filename) + response = requests.get(url) + + if response.status_code == 200: + with open(Path(storage_dir, filename), 'wb') as fd: + for chunk in response.iter_content(chunk_size=128): + fd.write(chunk) + logger.debug("Download [%s] success.", filename) + return True + + logger.debug("Download [%s] failed.", filename) + return False + + +def download_models(): + """Downloads pre-trained models from ImageAI GitHub.""" + logger.info( + "Scheduling pre-trained models download to [%s]", Path(storage).absolute() + ) + if not Path(storage).exists(): + logger.debug("Path [%s] did not exist, creating new.", Path(storage).absolute()) + Path(storage).mkdir() + + for filename, url in pre_trained_models.items(): + if Path(storage, filename).exists(): + logger.debug( + "File [%s] exist, skipping.", Path(storage, filename).absolute() + ) + + continue + + download_file(filename, url, storage) diff --git a/tests/test_downloader.py b/tests/test_downloader.py new file mode 100644 index 0000000..e69de29