From feb0582439c6a692ffd15827efc7a9b3134db841 Mon Sep 17 00:00:00 2001 From: Lucas Vandroux Date: Tue, 9 Apr 2019 13:42:46 +0200 Subject: [PATCH] Add rockethub package --- detect.py | 2 +- requirements.txt | 1 + rockethub/DiscoveryAPI.py | 116 -------------------------------------- rockethub/Rocket.py | 89 ----------------------------- rockethub/__init__.py | 0 5 files changed, 2 insertions(+), 206 deletions(-) delete mode 100644 rockethub/DiscoveryAPI.py delete mode 100644 rockethub/Rocket.py delete mode 100644 rockethub/__init__.py diff --git a/detect.py b/detect.py index 7d6a0ff..76b5586 100644 --- a/detect.py +++ b/detect.py @@ -1,6 +1,6 @@ import torch from torch.autograd import Variable -from rockethub.Rocket import Rocket +from rockethub import Rocket from PIL import Image # --- LOAD IMAGE --- diff --git a/requirements.txt b/requirements.txt index 3806485..62fdcf0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -25,6 +25,7 @@ pyparsing==2.3.1 python-dateutil==2.8.0 PyWavelets==1.0.2 requests==2.21.0 +rockethub==0.1.1 scikit-image==0.15.0 scipy==1.2.1 setuptools==40.8.0 diff --git a/rockethub/DiscoveryAPI.py b/rockethub/DiscoveryAPI.py deleted file mode 100644 index ba6fc51..0000000 --- a/rockethub/DiscoveryAPI.py +++ /dev/null @@ -1,116 +0,0 @@ -import itertools - -class DiscoveryAPI: - def __init__(self): - self.hangar = { - 'igor':{ - 'retinanet':{ - 'folder_name': 'RETINANET', - 'v1a': 'https://storage.googleapis.com/rockets-hangar/retinanet_v1a.tar' - }, - }, - 'lucas':{ - 'yolov3':{ - 'folder_name': 'YOLOv3', - 'v1a': 'https://storage.googleapis.com/rockets-hangar/yolov3_v1a.tar', - 'v1b': 'https://storage.googleapis.com/rockets-hangar/yolov3_v1b.tar', - 'v1c': 'https://storage.googleapis.com/rockets-hangar/yolov3_v1c.tar', - }, - 'ssd':{ - 'folder_name': 'SSD', - 'v1a': 'https://storage.googleapis.com/rockets-hangar/ssd_v1a.tar', - 'v1b': 'https://storage.googleapis.com/rockets-hangar/ssd_v1b.tar', - 'v1c': 'https://storage.googleapis.com/rockets-hangar/ssd_v1c.tar', - }, - } - } - - def get_rocket_info(self, rocket: str): - """ Parse the Rocket identifier. - - The Rocket identifier is a String following this pattern: rocket_author/rocket_name/rocket_version - The rocket_version is not mandatory and if not provided by the user the newest version will be returned. - - Args: - rocket (str): Rocket identifier that needs to be parsed - """ - assert len(rocket) > 0, 'Please specify the rocket you want to get.' - - # Parse the Rocket url - rocket_parsed = rocket.split('/') - assert len(rocket_parsed) > 1, 'Please provide more information about the rocket you want to get.' - - rocket_author = rocket_parsed[0] - rocket_name = rocket_parsed[1] - - print('Looking for the Rocket ' + rocket_name + ' made by ' + rocket_author + '...') - - # Test that the rocket exists - assert rocket_author in self.hangar.keys(), rocket_author + ' can\'t be found as an author.' - assert rocket_name in self.hangar[rocket_author].keys(), rocket_name + ' can\'t be found as a rocket from ' + rocket_author - - if len(rocket_parsed) <= 2: - rocket_version = self.get_rocket_last_version(rocket_author, rocket_name) - print('You didn\'t say which version you wanted so we selected for you the newest one: ' + rocket_version) - - else: - rocket_version = rocket_parsed[2] - assert rocket_version in self.hangar[rocket_author][rocket_name].keys(), rocket_version + ' can\'t be found as a version for the rocket ' + rocket_name + ' from author ' + rocket_author - print('Version ' + rocket_version + ' exists and is waiting for you.') - - return rocket_author, rocket_name, rocket_version - - def get_rocket_url(self, rocket_author: str, rocket_name: str, rocket_version: str): - """ Get the url from which to download the Rocket. - - Args: - rocket_author (str): Username of the author of the Rocket - rocket_name (str): Name of the rocket - rocket_version (str): Version of the Rocket - """ - return self.hangar[rocket_author][rocket_name][rocket_version] - - def get_rocket_folder(self, rocket_author: str, rocket_name: str, rocket_version: str): - """ Get the name of the folder where the Rocket is unpacked. - - Args: - rocket_author (str): Username of the author of the Rocket - rocket_name (str): Name of the rocket - rocket_version (str): Version of the Rocket - """ - return self.hangar[rocket_author][rocket_name]['folder_name'] - - def get_rocket_last_version(self, rocket_author: str, rocket_name: str): - """Get the last version of a Rocket. - - Args: - rocket_author (str): Username of the author of the Rocket - rocket_name (str): Name of the rocket - """ - # Verify the rocket exist - assert rocket_author in self.hangar.keys(), rocket_author + ' can\'t be found as an author.' - assert rocket_name in self.hangar[rocket_author].keys(), rocket_name + ' can\'t be found as a rocket from ' + rocket_author - - # Get list of versions for a specific Rockets - list_versions = [v[1:] for v in self.hangar[rocket_author][rocket_name].keys() if v.startswith('v')] - - mainVersion = 0 - minorVersion = 'a' - - for version in list_versions: - v = ["".join(x) for _, x in itertools.groupby(version, key=str.isdigit)] - temp_mainVersion = int(v[0]) - temp_minorVersion = v[1] - - assert len(temp_minorVersion) == 1, 'Automatic selection of the newest version doesn\'t support minor version made of more than 1 char.' - - if temp_mainVersion == mainVersion: - if temp_minorVersion > minorVersion: - minorVersion = temp_minorVersion - elif temp_mainVersion > mainVersion: - mainVersion = temp_mainVersion - minorVersion = 'a' - if temp_minorVersion > minorVersion: - minorVersion = temp_minorVersion - - return 'v' + str(mainVersion) + minorVersion diff --git a/rockethub/Rocket.py b/rockethub/Rocket.py deleted file mode 100644 index dc5d6d7..0000000 --- a/rockethub/Rocket.py +++ /dev/null @@ -1,89 +0,0 @@ -import tarfile -import os -import importlib -import types -import requests -from tqdm import tqdm -from .DiscoveryAPI import DiscoveryAPI -from datetime import datetime - -def unpack_archive(path: str): - ensure_dir('rockets') - t = tarfile.open(path, 'r') - model_name = os.path.commonprefix(t.getnames()) - t.extractall('rockets') - os.remove(path) - - return model_name - - -def ensure_dir(dir_name: str): - """Creates folder if not exists. - """ - if not os.path.exists(dir_name): - os.makedirs(dir_name) - -class Rocket: - - @staticmethod - def land(rocket: str, folder_path = 'rockets', chunk_size = 512): - """ Download or check that the Rocket is ready locally - - Download the Rocket if it is not yet locally here. - - Args: - rocket (str): Rocket identifier (author/name/(version)) - folder_path (str): folder where to check if the Rocket is here already or where to download it - chunk_size (int): size of the chunk when downloading the Rocket - """ - - # Check if the rocket exists and get the last version if not precised - rocket_author, rocket_name, rocket_version = DiscoveryAPI().get_rocket_info(rocket) - - # Check if folder to download the rockets exists - ensure_dir(folder_path) - - # Check if the rocket has already been downloaded - rocket_folder = DiscoveryAPI().get_rocket_folder(rocket_author, rocket_name, rocket_version) - - if rocket_folder in os.listdir(folder_path): - model_name = rocket_folder - print('Rocket is already here.') - - else: - # Get the rocket's url - url = DiscoveryAPI().get_rocket_url(rocket_author, rocket_name, rocket_version) - path_to_landing_rocket = os.path.join(folder_path, 'landing_rocket.tar') - - #Download URL - print('Rocket in approach...') - h = requests.head(url, allow_redirects=True) - headers = h.headers - content_type = headers.get('content-type') - content_length = int(headers.get('content-length', None)) - # print('content-type', content_type) - - response = requests.get(url, stream=True) - - pbar = tqdm(total=content_length, ascii=True, desc='Rocket Landing') - with open(path_to_landing_rocket, 'wb') as handle: - for data in response.iter_content(chunk_size = chunk_size): - handle.write(data) - # update the progress bar - pbar.update(chunk_size) - - pbar.close() - - model_name = unpack_archive(path_to_landing_rocket) - print('It is a sucess! The Rocket has landed!') - - print("Let's build the Rocket...") - #Build the model - module = importlib.import_module('rockets.{}.rocket_builder'.format(model_name)) - build_func = getattr(module, 'build') - model = build_func() - print("The Rocket is ready for use!") - return model - - - diff --git a/rockethub/__init__.py b/rockethub/__init__.py deleted file mode 100644 index e69de29..0000000