cyberes
/
wmts-exfiltrator
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add arg for specifying number of tiff threads

This commit is contained in:
Cyberes 2023-11-05 20:34:45 -07:00
parent 62b8a17fe7
commit 7499f26481
2 changed files with 7 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
exfiltrate.py
View File

@ -22,6 +22,7 @@ if __name__ == '__main__':
parser.add_argument('--referer', help='The content of the Referer header to send.')
parser.add_argument('--output', default='wmts-output', help='Output directory path.')
parser.add_argument('--proxy', action='store_true', help='Enable using a proxy.')
parser.add_argument('--tiff-threads', default=None, help='Number of threads to use when building TIFF. Default: auto')
parser.add_argument('--output-tiff', help='Path for output GeoTIFF. Default: wmts-output/output.tiff')
parser.add_argument('--bbox', required=True, type=str, metavar='Bounding Box', nargs='+', default=(None, None, None, None), help='Bounding Box of the area to download. Separate each value with a space. (top left lat, top left lon, bottom right lat, bottom right lon)')
args = parser.parse_args()
@ -57,9 +58,10 @@ if __name__ == '__main__':
row_bar = tqdm(total=len(row_iter), desc=f'Row {row_i}', postfix={'new_files': total_downloaded, 'failures': len(retries)})
for row in row_iter:
row_i = row
col_bar = tqdm(total=len(range(min_col, max_col + 1)), leave=False, disable=True)
col_iter = range(min_col, max_col + 1)
col_bar = tqdm(total=len(col_iter), leave=False, disable=True)
with ThreadPoolExecutor(args.threads) as executor:
futures = [executor.submit(download_tile, (row, col, args.base_url, r_headers, tiles_output, args.proxy)) for col in range(min_col, max_col + 1)]
futures = [executor.submit(download_tile, (row, col, args.base_url, r_headers, tiles_output, args.proxy)) for col in col_iter]
for future in as_completed(futures):
result = future.result()
if result:
@ -73,8 +75,8 @@ if __name__ == '__main__':
elif new_image == 'failure':
retries.append((result_row, result_col))
row_bar.set_postfix({'new_files': total_downloaded, 'failures': len(retries)})
col_bar.update()
row_bar.refresh()
col_bar.update()
row_bar.refresh()
col_bar.close()
row_bar.set_postfix({'new_files': total_downloaded, 'failures': len(retries)})
row_bar.update()
@ -136,7 +138,7 @@ if __name__ == '__main__':
image_data[row_pos:row_pos + tile_size, col_pos:col_pos + tile_size] = tile_data
with ThreadPoolExecutor() as executor:
with ThreadPoolExecutor(max_workers=args.tiff_threads) as executor:
futures = {executor.submit(build_tiff_data, task) for task in tiles}
for future in tqdm(as_completed(futures), total=len(futures), desc='Building TIFF'):
pass

4
pkg/thread.py
View File

@ -33,10 +33,6 @@ def download_tile(task):
raise Exception(f'Response gave Content-Type: {response.headers.get("Content-Type")}')
with open(output_path, "wb") as f:
f.write(response.content)
# Don't check the PNG after downloading, only when loading.
# if not is_png(output_path):
# del_path(output_path)
# return row, col, 'failure'
return row, col, 'success'
else:
print(f"Failed to download tile {row}_{col}")