break tiffs into chunks
This commit is contained in:
parent
65953c9bde
commit
6acbdf3117
|
|
@ -120,65 +120,70 @@ if __name__ == '__main__':
|
|||
* Affine.scale((bottom_right_lon - top_left_lon) / (num_cols * tile_size),
|
||||
(bottom_right_lat - top_left_lat) / (num_rows * tile_size)))
|
||||
|
||||
# Divide the tiles into n groups
|
||||
tile_groups = np.array_split(tiles, args.tiff_threads)
|
||||
|
||||
def worker(pbar):
|
||||
while True:
|
||||
row, col = q.get()
|
||||
if row is None:
|
||||
break
|
||||
|
||||
tile_file = tiles_output / f"{row}_{col}.png"
|
||||
if not tile_file.is_file():
|
||||
raise Exception(f'Tile does not exist: {tile_file}')
|
||||
def worker(pbar, output_tiff, tile_group):
|
||||
with rasterio.open(output_tiff, "w", driver="GTiff", height=num_rows * tile_size, width=num_cols * tile_size, count=3, dtype='uint8', crs='EPSG:4326', transform=transform, compress="DEFLATE", nodata=0) as dst:
|
||||
while True:
|
||||
task = q.get()
|
||||
if task is None:
|
||||
break
|
||||
|
||||
with Image.open(tile_file) as img:
|
||||
tile_data = np.array(img, dtype=np.uint8)
|
||||
row, col = task
|
||||
tile_file = tiles_output / f"{row}_{col}.png"
|
||||
if not tile_file.is_file():
|
||||
raise Exception(f'Tile does not exist: {tile_file}')
|
||||
|
||||
# Remove the alpha channel
|
||||
tile_data = tile_data[:, :, :3]
|
||||
with Image.open(tile_file) as img:
|
||||
tile_data = np.array(img, dtype=np.uint8)
|
||||
|
||||
# Replace white pixels with NODATA
|
||||
tile_data[np.all(tile_data == [255, 255, 255], axis=-1)] = [0, 0, 0]
|
||||
# Remove the alpha channel
|
||||
tile_data = tile_data[:, :, :3]
|
||||
|
||||
# ArcGIS does not like pixels that have zeros in them, eg. (255, 0, 0). We need to convert the zeros to ones, eg. (255, 1, 1).
|
||||
mask = np.any(tile_data == 0, axis=-1) & np.any(tile_data != 0, axis=-1) # Identify pixels where not all bands are zero and at least one band is zero.
|
||||
for i in range(3): # Iterate over each band.
|
||||
# For these pixels, set zero bands to one.
|
||||
tile_data[mask & (tile_data[:, :, i] == 0), i] = 1
|
||||
# Replace white pixels with NODATA
|
||||
tile_data[np.all(tile_data == [255, 255, 255], axis=-1)] = [0, 0, 0]
|
||||
|
||||
# Calculate the position of the tile in the image data array.
|
||||
row_pos = (row - min_row) * tile_size
|
||||
col_pos = (col - min_col) * tile_size
|
||||
# ArcGIS does not like pixels that have zeros in them, eg. (255, 0, 0). We need to convert the zeros to ones, eg. (255, 1, 1).
|
||||
mask = np.any(tile_data == 0, axis=-1) & np.any(tile_data != 0, axis=-1) # Identify pixels where not all bands are zero and at least one band is zero.
|
||||
for i in range(3): # Iterate over each band.
|
||||
# For these pixels, set zero bands to one.
|
||||
tile_data[mask & (tile_data[:, :, i] == 0), i] = 1
|
||||
|
||||
tile_data = np.transpose(tile_data, (2, 0, 1))
|
||||
# Calculate the position of the tile in the image data array.
|
||||
row_pos = (row - min_row) * tile_size
|
||||
col_pos = (col - min_col) * tile_size
|
||||
|
||||
# Write the tile data to the GeoTIFF file
|
||||
with lock:
|
||||
tile_data = np.transpose(tile_data, (2, 0, 1))
|
||||
|
||||
# Write the tile data to the GeoTIFF file
|
||||
dst.write(tile_data, window=rasterio.windows.Window(col_pos, row_pos, tile_size, tile_size), indexes=[1, 2, 3])
|
||||
|
||||
q.task_done()
|
||||
pbar.update()
|
||||
q.task_done()
|
||||
pbar.update()
|
||||
|
||||
|
||||
q = Queue()
|
||||
lock = threading.Lock()
|
||||
|
||||
with rasterio.open(output_tiff, "w", driver="GTiff", height=num_rows * tile_size, width=num_cols * tile_size, count=3, dtype='uint8', crs='EPSG:4326', transform=transform, compress="DEFLATE", nodata=0) as dst:
|
||||
with tqdm(total=len(tiles), desc='Building GeoTIFF') as pbar:
|
||||
threads = []
|
||||
for i in range(args.tiff_threads):
|
||||
t = threading.Thread(target=worker, args=(pbar,))
|
||||
t.start()
|
||||
threads.append(t)
|
||||
with tqdm(total=len(tiles), desc='Building GeoTIFF') as pbar:
|
||||
threads = []
|
||||
for i in range(args.tiff_threads):
|
||||
output_tiff_thread = output_tiff.with_stem(output_tiff.stem + f'_part{i}')
|
||||
t = threading.Thread(target=worker, args=(pbar, output_tiff_thread, tile_groups[i]))
|
||||
t.start()
|
||||
threads.append(t)
|
||||
|
||||
for row, col in tiles:
|
||||
for i, tile_group in enumerate(tile_groups):
|
||||
for row, col in tile_group:
|
||||
q.put((row, col))
|
||||
|
||||
# block until all tasks are done
|
||||
q.join()
|
||||
# block until all tasks are done
|
||||
q.join()
|
||||
|
||||
# stop workers
|
||||
for i in range(args.tiff_threads):
|
||||
q.put((None, None))
|
||||
for t in threads:
|
||||
t.join()
|
||||
# stop workers
|
||||
for i in range(args.tiff_threads):
|
||||
q.put(None)
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
|
|
|||
Loading…
Reference in New Issue