Pre-trained models downloader

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",
+}

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)