def from_model_bundle(model_bundle_uri: str, tmp_dir: str): """Create a Learner from a model bundle.""" model_bundle_path = download_if_needed(model_bundle_uri, tmp_dir) model_bundle_dir = join(tmp_dir, 'model-bundle') unzip(model_bundle_path, model_bundle_dir) config_path = join(model_bundle_dir, 'pipeline-config.json') model_path = join(model_bundle_dir, 'model.pth') config_dict = file_to_json(config_path) config_dict = upgrade_config(config_dict) cfg = build_config(config_dict) hub_dir = join(model_bundle_dir, MODULES_DIRNAME) model_def_path = None loss_def_path = None # retrieve existing model definition, if available ext_cfg = cfg.learner.model.external_def if ext_cfg is not None: model_def_path = get_hubconf_dir_from_cfg(ext_cfg, parent=hub_dir) log.info( f'Using model definition found in bundle: {model_def_path}') # retrieve existing loss function definition, if available ext_cfg = cfg.learner.solver.external_loss_def if ext_cfg is not None: loss_def_path = get_hubconf_dir_from_cfg(ext_cfg, parent=hub_dir) log.info(f'Using loss definition found in bundle: {loss_def_path}') return cfg.learner.build(tmp_dir=tmp_dir, model_path=model_path, model_def_path=model_def_path, loss_def_path=loss_def_path)
def from_model_bundle(model_bundle_uri: str, tmp_dir: str): """Create a Learner from a model bundle.""" model_bundle_path = download_if_needed(model_bundle_uri, tmp_dir) model_bundle_dir = join(tmp_dir, 'model-bundle') unzip(model_bundle_path, model_bundle_dir) config_path = join(model_bundle_dir, 'pipeline-config.json') model_path = join(model_bundle_dir, 'model.pth') config_dict = file_to_json(config_path) config_dict = upgrade_config(config_dict) cfg = build_config(config_dict) return cfg.learner.build(tmp_dir, model_path=model_path)
def __init__(self, model_bundle_uri, tmp_dir, update_stats=False, channel_order=None): """Creates a new Predictor. Args: model_bundle_uri: URI of the model bundle to use. Can be any type of URI that Raster Vision can read. tmp_dir: Temporary directory in which to store files that are used by the Predictor. This directory is not cleaned up by this class. channel_order: Option for a new channel order to use for the imagery being predicted against. If not present, the channel_order from the original configuration in the predict package will be used. """ self.tmp_dir = tmp_dir self.update_stats = update_stats self.model_loaded = False bundle_path = download_if_needed(model_bundle_uri, tmp_dir) bundle_dir = join(tmp_dir, 'bundle') make_dir(bundle_dir) with zipfile.ZipFile(bundle_path, 'r') as bundle_zip: bundle_zip.extractall(path=bundle_dir) config_path = join(bundle_dir, 'pipeline-config.json') config_dict = file_to_json(config_path) rv_config.set_everett_config( config_overrides=config_dict.get('rv_config')) config_dict = upgrade_config(config_dict) self.pipeline = build_config(config_dict).build(tmp_dir) self.scene = None if not hasattr(self.pipeline, 'predict'): raise Exception( 'pipeline in model bundle must have predict method') self.scene = self.pipeline.config.dataset.validation_scenes[0] if not hasattr(self.scene.raster_source, 'uris'): raise Exception( 'raster_source in model bundle must have uris as field') if not hasattr(self.scene.label_store, 'uri'): raise Exception( 'label_store in model bundle must have uri as field') for t in self.scene.raster_source.transformers: t.update_root(bundle_dir) if self.update_stats: stats_analyzer = StatsAnalyzerConfig( output_uri=join(bundle_dir, 'stats.json')) self.pipeline.config.analyzers = [stats_analyzer] self.scene.label_source = None self.scene.aoi_uris = None self.pipeline.config.dataset.train_scenes = [self.scene] self.pipeline.config.dataset.validation_scenes = [self.scene] self.pipeline.config.dataset.test_scenes = None self.pipeline.config.train_uri = bundle_dir if channel_order is not None: self.scene.raster_source.channel_order = channel_order
def test_upgrade(self): plugin_versions_v0 = dict(self.plugin_versions) plugin_versions_v0['rastervision.ab'] = 0 plugin_versions_v0['rastervision2.c'] = 0 # after upgrading: the y field in the root should get converted to x, and # the z field in the instances of a should get convert to x. c_dict_v0 = { 'plugin_versions': plugin_versions_v0, 'root_uri': None, 'rv_config': None, 'type_hint': 'c', 'a': { 'type_hint': 'asub1', 'z': 'x', 'y': 'y' }, 'al': [{ 'type_hint': 'a', 'z': 'x' }, { 'type_hint': 'asub1', 'z': 'x', 'y': 'y' }, { 'type_hint': 'asub2', 'z': 'x', 'y': 'y' }], 'b': { 'type_hint': 'b', 'x': 'x' }, 'bl': [{ 'type_hint': 'b', 'x': 'x' }], 'y': 'x' } c_dict_v1 = { 'plugin_versions': plugin_versions_v0, 'root_uri': None, 'rv_config': None, 'type_hint': 'c', 'a': { 'type_hint': 'asub1', 'x': 'x', 'y': 'y' }, 'al': [{ 'type_hint': 'a', 'x': 'x' }, { 'type_hint': 'asub1', 'x': 'x', 'y': 'y' }, { 'type_hint': 'asub2', 'x': 'x', 'y': 'y' }], 'b': { 'type_hint': 'b', 'x': 'x' }, 'bl': [{ 'type_hint': 'b', 'x': 'x' }], 'x': 'x' } upgraded_c_dict = upgrade_config(c_dict_v0) self.assertDictEqual(upgraded_c_dict, c_dict_v1)