diff --git a/ml_export/ml_tools/mlbase.py b/ml_export/ml_tools/mlbase.py index 291aba2..28fb080 100644 --- a/ml_export/ml_tools/mlbase.py +++ b/ml_export/ml_tools/mlbase.py @@ -80,7 +80,6 @@ class MLTFServing(): predict: This should receive a np array of 3 x 1024 x 1024 and return a numpy array of 1x1024x1024 predict_batch: This should receive a list of np arrays of [np(3,1024,1024)] and return a list of [np(1,1024,1024] - model_dictionary = {'model_file': "test.hdf5", "model_description": "Passthrough Model", "model_version": "0.1", @@ -89,7 +88,8 @@ class MLTFServing(): """ - def __init__(self, api_location, output_num_channels=1, debug=False): + def __init__(self, api_location, output_num_channels=1, preproc_func=None, + postproc_func=None, debug=False): ''' Inititialize model ''' self.logger = logging.getLogger(__name__) @@ -115,6 +115,8 @@ def __init__(self, api_location, output_num_channels=1, debug=False): self.predict_api_loc = api_location self.num_channels = output_num_channels self.model_speed = 1 + self.preproc_func = preproc_func + self.postproc_func = postproc_func # Load Model Into Memory self.load_model_dict() @@ -132,14 +134,19 @@ def predict(self, np_array): return np_array[None, 0, :, :] def predict_batch(self, super_res_tile_batch): - inputs = np.moveaxis(super_res_tile_batch, 1, 3).astype(np.float32)/255 + if self.preproc_func is not None: + inputs = self.preproc_func(super_res_tile_batch) + else: + inputs = super_res_tile_batch + inputs = np.moveaxis(inputs, 1, 3).astype(np.float32)/255 payload = {'inputs': inputs.tolist()} # Send prediction request r = requests.post(self.predict_api_loc, json=payload) content = json.loads(r.content) - all_image_preds = np.asarray(content['outputs']).reshape(len(inputs), - 256, 256) - all_image_preds = all_image_preds[:, np.newaxis, :, :] - - return all_image_preds + preds = np.asarray(content['outputs']).reshape(len(inputs), + 256, 256) + preds = preds[:, np.newaxis, :, :] + if self.postproc_func is not None: + preds = self.postproc_func(preds) + return preds diff --git a/ml_export/tile_aggregator.py b/ml_export/tile_aggregator.py index 471ed9b..2c7c465 100644 --- a/ml_export/tile_aggregator.py +++ b/ml_export/tile_aggregator.py @@ -1,5 +1,4 @@ import mercantile -from rio_tiler import main import numpy as np from affine import Affine import rasterio @@ -12,7 +11,8 @@ logging.basicConfig(format='%(levelname)s:%(asctime)s %(message)s', datefmt='%m/%d/%Y %I:%M:%S %p', level=logging.INFO) - +logger = logging.getLogger(__name__) +logger.setLevel(logging.DEBUG) def calculate_webmercator_meters_per_pixel(zoom_level): """Calculate webmercator pixel size based on zoom level.""" @@ -239,7 +239,7 @@ def build_cog_from_tiles_gen(file_name, large_tile_object, with rasterio.open(file_name, 'w', **large_cog_profile) as dst_dataset: - tile_dataset = tile_dataset_class( + tile_dataset = TileClassDataset( root_tile_obj=large_tile_object, raster_location=raster_tile_server_template, desired_zoom_level=desired_small_tile_zoom_level, @@ -258,8 +258,12 @@ def build_cog_from_tiles_gen(file_name, large_tile_object, else: super_res_tile_results = super_res_tile_np for super_res_tile, small_tile_object_tensor in zip( - super_res_tile_results, zip(small_tile_obj_batch[i].numpy() - for i in range(batch_size))): + super_res_tile_results, + zip(small_tile_obj_batch[0].numpy(), + small_tile_obj_batch[1].numpy(), + small_tile_obj_batch[2].numpy(), + small_tile_obj_batch[3].numpy() + )): left, bottom, right, top = small_tile_object_tensor dst_window = rasterio.windows.from_bounds( @@ -268,4 +272,6 @@ def build_cog_from_tiles_gen(file_name, large_tile_object, dst_dataset.write(super_res_tile.astype( large_cog_profile['dtype']), window=dst_window) + dst_dataset.close() + return file_name diff --git a/ml_export/tile_class_generator.py b/ml_export/tile_class_generator.py index cb31a22..dd73337 100644 --- a/ml_export/tile_class_generator.py +++ b/ml_export/tile_class_generator.py @@ -1,6 +1,5 @@ """tests ml_export.tile_class_generator""" import os -import pytest ## Note, for mac osx compatability import something from shapely.geometry before importing fiona or geopandas ## https://github.com/Toblerity/Shapely/issues/553 * Import shapely before rasterio or fioana from shapely import geometry diff --git a/ml_export/tile_generator.py b/ml_export/tile_generator.py index ff2e477..539f6ec 100644 --- a/ml_export/tile_generator.py +++ b/ml_export/tile_generator.py @@ -130,7 +130,6 @@ def create_super_tile_image(tile_object, address, desired_zoom_level=19, (len(indexes,(2**zoom_level)*tile_size,(2**zoom_level)*tile_size) """ - if cog: return create_super_tile_image_cog( tile_object, address, desired_zoom_level=desired_zoom_level, diff --git a/nw_notes.md b/nw_notes.md index 2c1664b..4912f07 100644 --- a/nw_notes.md +++ b/nw_notes.md @@ -57,3 +57,11 @@ This module appears to be intended for creating Torch datasets for OpenCV models ## Notes on `ml_export/utils.py` - Utilities for generating overviews, managing S3 storage of imagery, and creating STAC items for inference outputs. + + +## Notes on implementation +It looks like dave was trying to use `tile_aggregator.build_cog_from_tiles_gen` to + +# steps to the query that need to be executed +## 1. TM User triggers request with API +## 2. ml-export-tool initiates