· 5 years ago · Jun 10, 2020, 07:28 PM
1import concurrent.futures
2from collections import defaultdict
3from typing import Tuple, List, DefaultDict
4
5import imageio
6import mercantile
7import numpy as np
8import rasterio as rio
9import requests
10from affine import Affine
11from rasterio import transform, MemoryFile
12from rasterio.coords import BoundingBox
13from rasterio.errors import RasterioIOError
14from rasterio.windows import Window
15
16import boto3
17s3 = boto3.client("s3")
18
19
20def get_tiles_from_field_metadata(rgb_key: str, channel: str, zooms: int) -> List[List[int]]:
21 # Tiles per defined zoom level with mercantile API.
22 response = requests.get(f"https://jgau9m9qpf.execute-api.us-east-1.amazonaws.com/bondville"
23 f"/metadata/{rgb_key}/{channel}/")
24 west, south, east, north = response.json()['bounds']
25 tiles = list(mercantile.tiles(west, south, east, north, zooms))
26 final_tiles = [[tile.x, tile.y, tile.z] for tile in tiles]
27 return final_tiles
28
29
30def get_resulting_image_shape(rgb_key: str, channel: str, zooms: int,
31 tile_size: tuple = (256, 256)) -> Tuple[int, int]:
32 # Final resulting image shape (height, width).
33 tiles = get_tiles_from_field_metadata(rgb_key, channel, zooms)
34 combined_arrays_by_x = combine_image_data_by_x(tiles)
35
36 image_width, image_height = len(combined_arrays_by_x.keys()) * tile_size[0], None
37 for k in combined_arrays_by_x:
38 image_height = len(combined_arrays_by_x[k]) * tile_size[1]
39 break
40 return image_width, image_height
41
42
43def combine_image_data_by_x(tiles: List[List[int]]) -> DefaultDict[int, List[List[int]]]:
44 # Get dict of image array combined by x value.
45 initial_x = tiles[0][0]
46 combined_arrays_by_x = defaultdict(list)
47 for tile in tiles:
48 local_x = tile[0]
49 if local_x == initial_x:
50 xyz_array = [tile[0], tile[1], tile[2]]
51 combined_arrays_by_x[local_x].append(xyz_array)
52 else:
53 initial_x = local_x
54 xyz_array = [tile[0], tile[1], tile[2]]
55 combined_arrays_by_x[initial_x].append(xyz_array)
56 return combined_arrays_by_x
57
58
59def get_tile_bounds(tile: list) -> BoundingBox:
60 # Returns BoundingBox object for each given tile.
61 return BoundingBox(*mercantile.xy_bounds(tile))
62
63
64def transform_from_bounds(bounds, tile_size) -> Affine:
65 return transform.from_bounds(*bounds, *tile_size)
66
67
68def get_resulting_image_bounds(rgb_key: str, channel: str, zooms: int) -> BoundingBox:
69 # Final resulting image boundary (x_min, y_min, x_max, y_max).
70 tiles = get_tiles_from_field_metadata(rgb_key, channel, zooms)
71 x_min, y_min, x_max, y_max = [], [], [], []
72 for index, tile in enumerate(tiles):
73 bbox = get_tile_bounds(tiles[index])
74 x_min.append(bbox.left)
75 y_min.append(bbox.bottom)
76 x_max.append(bbox.right)
77 y_max.append(bbox.top)
78
79 return BoundingBox(min(x_min), min(y_min), max(x_max), max(y_max))
80
81
82def construct_url(rgb_key: str, x: int, y: int, z: int):
83 # Makes api calls to rgb tile endpoint.
84 return f"https://jgau9m9qpf.execute-api.us-east-1.amazonaws.com/bondville/" \
85 f"{rgb_key}/{x}/{y}/{z}.png?&r=red&g=green&b=blue" \
86 f"&r_range=%5B2048.0%2C8735.0%5D&g_range=%5B2048.0%2C8735.0%5D&b_range=%5B2048.0%2C8735.0%5D"
87
88
89def get_image_data_array(rgb_key: str, x: int, y: int, z: int):
90 # Construct tile data each time calling api with respective x, y, z coordinates.
91 image_data = imageio.imread(construct_url(rgb_key, x, y, z))
92 return image_data
93
94
95def write_image_for_one_tile(destination, image_data, image_width, image_height, g_transform,
96 no_data, window, count=3, mode='r+', driver='GTiff', crs='epsg:3857'):
97 dst_profile = {
98 'driver': driver,
99 'count': count,
100 'height': image_height,
101 'width': image_width,
102 'dtype': image_data.dtype,
103 'nodata': no_data,
104 'crs': crs,
105 'transform': g_transform,
106 }
107 with MemoryFile() as memfile:
108 with memfile.open(**dst_profile) as dst:
109 dst.write(np.transpose(np.moveaxis(image_data, [0, 2, 1], [1, 2, 0])), window=window)
110 s3.upload_fileobj(memfile, destination, 'temp.tif')
111
112
113def thread_function(d_file: str, row_offset: int, column_offset: int, tile: list, width: int,
114 height: int, g_transform: Affine, no_data_v: int):
115 # Thread function in order to pass to ThreadPullExecutor
116 image_data_ = get_image_data_array("rgb/flight/9BINMT6XB", tile[2], tile[0], tile[1])
117 tile_window = Window(column_offset, row_offset, 256, 256)
118 write_image_for_one_tile(d_file, image_data_, width,
119 height, g_transform, no_data_v,
120 tile_window)
121
122
123if __name__ == "__main__":
124 resulting_tiles = get_tiles_from_field_metadata('flight/9BINMT6XB', 'red', 18)
125 res_image_width, res_image_height = get_resulting_image_shape('flight/9BINMT6XB', 'red', 18)
126 res_bounds = get_resulting_image_bounds('flight/9BINMT6XB', 'red', 18)
127 geo_transform = transform_from_bounds(res_bounds, (res_image_width, res_image_height))
128 tiles_map = combine_image_data_by_x(resulting_tiles)
129
130 # Please refer here to come with solution to "error setting certificate verify locations".
131 # https://github.com/mapbox/rasterio/issues/1289
132 s3_destination = 'intelinair-infra'
133
134 # s3_destination = 'temp.tif'
135 no_data_value = 0
136
137 col_off = -256
138
139 for key, x_tiles in tiles_map.items():
140 col_off += 256
141 row_off = -256
142
143 with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
144 futures = []
145 for x_tile in x_tiles:
146 row_off += 256
147 futures.append(executor.submit(thread_function, s3_destination, row_off, col_off,
148 x_tile, res_image_width, res_image_height,
149 geo_transform, no_data_value))
150 concurrent.futures.wait(futures)