Added more configurability to map generator

This commit is contained in:
Pax1601
2024-03-10 13:43:09 +01:00
parent 5cc42dd9cf
commit 8c7f6abb1c
2 changed files with 50 additions and 44 deletions

View File

@@ -1,4 +1,4 @@
import sys import os
import yaml import yaml
import json import json
import requests import requests
@@ -16,9 +16,12 @@ parser = argparse.ArgumentParser(
epilog='Hit the DCS Olympus Discord for more information') epilog='Hit the DCS Olympus Discord for more information')
parser.add_argument('config', help='map configuration yaml file') parser.add_argument('config', help='map configuration yaml file')
parser.add_argument('-s', '--skip_screenshots', action='store_true', help='if screenshots are already present, this flag will cause the script to completely skip the screenshot loop') parser.add_argument('-s', '--skip_screenshots', action='store_true', help='if screenshots are already present, this flag will cause the script to completely skip the screenshot loop.')
parser.add_argument('-r', '--replace_screenshots', action='store_true', help='if screenshots are already present, this flag will cause the script to replace all screenshots, even those that already exist. Has no effect if -s or --skip_screenshots is present') parser.add_argument('-r', '--replace_screenshots', action='store_true', help='if screenshots are already present, this flag will cause the script to replace all screenshots, even those that already exist. Has no effect if -s or --skip_screenshots is present.')
parser.add_argument('-l', '--final_level', type=int, default=1, help='if tiles are already present for the zoom level that the script will output, this number will instruct up to which zoom level tile merging will be run. Defaults to 1.') parser.add_argument('-l', '--final_level', type=int, default=1, help='if tiles are already present for the zoom level that the script will output, this number will instruct up to which zoom level tile merging will be run. Defaults to 1.')
parser.add_argument('-f', '--screenshots_folder', help='if provided, will force the script to save the screenshots here. Defaults to output_directory/screenshots.')
parser.add_argument('-t', '--tiles_folder', help='if provided, will force the script to save the tiles here. Defaults to output_directory/tiles.')
parser.add_argument('-o', '--screenshots_only', action='store_true', help='if provided, the script will only run the screenshot acquisition algorithm.')
args = parser.parse_args() args = parser.parse_args()
@@ -36,6 +39,10 @@ with open('configs/screen_properties.yml', 'r') as sp:
map_config = yaml.safe_load(cp) map_config = yaml.safe_load(cp)
map_config.update(vars(args)) map_config.update(vars(args))
if map_config["screenshots_folder"] is None:
map_config["screenshots_folder"] = os.path.join(map_config['output_directory'], "screenshots")
if map_config["tiles_folder"] is None:
map_config["tiles_folder"] = os.path.join(map_config['output_directory'], "tiles")
print("Screen parameters:") print("Screen parameters:")
print(f"-> Screen width: {screen_config['width']}px") print(f"-> Screen width: {screen_config['width']}px")
@@ -43,6 +50,8 @@ with open('configs/screen_properties.yml', 'r') as sp:
print("Map parameters:") print("Map parameters:")
print(f"-> Output directory: {map_config['output_directory']}") print(f"-> Output directory: {map_config['output_directory']}")
print(f"-> Screenshots folder: {map_config['screenshots_folder']}")
print(f"-> Tiles folder: {map_config['tiles_folder']}")
print(f"-> Boundary file: {map_config['boundary_file']}") print(f"-> Boundary file: {map_config['boundary_file']}")
print(f"-> Zoom factor: {map_config['zoom_factor']}") print(f"-> Zoom factor: {map_config['zoom_factor']}")
@@ -89,7 +98,3 @@ with open('configs/screen_properties.yml', 'r') as sp:
print(f"Estimated number of screenshots: {screenshots_num}") print(f"Estimated number of screenshots: {screenshots_num}")
map_generator.run(map_config, port) map_generator.run(map_config, port)

View File

@@ -58,14 +58,15 @@ def done_callback(fut):
def extract_tiles(n, screenshots_XY, params): def extract_tiles(n, screenshots_XY, params):
f = params['f'] f = params['f']
zoom = params['zoom'] zoom = params['zoom']
output_directory = params['output_directory']
n_width = params['n_width'] n_width = params['n_width']
n_height = params['n_height'] n_height = params['n_height']
screenshots_folder = params['screenshots_folder']
tiles_folder = params['tiles_folder']
XY = screenshots_XY[n] XY = screenshots_XY[n]
if (os.path.exists(os.path.join(output_directory, "screenshots", f"{f}_{n}_{zoom}.jpg"))): if (os.path.exists(os.path.join(screenshots_folder, f"{f}_{n}_{zoom}.jpg"))):
# Open the source screenshot # Open the source screenshot
img = Image.open(os.path.join(output_directory, "screenshots", f"{f}_{n}_{zoom}.jpg")) img = Image.open(os.path.join(screenshots_folder, f"{f}_{n}_{zoom}.jpg"))
# Compute the Web Mercator Projection position of the top left corner of the most centered tile # Compute the Web Mercator Projection position of the top left corner of the most centered tile
X_center, Y_center = XY[0], XY[1] X_center, Y_center = XY[0], XY[1]
@@ -83,19 +84,19 @@ def extract_tiles(n, screenshots_XY, params):
Y = Y_center - math.floor(n_height / 2) + row Y = Y_center - math.floor(n_height / 2) + row
# Save the tile # Save the tile
if not os.path.exists(os.path.join(output_directory, "tiles", str(zoom), str(X))): if not os.path.exists(os.path.join(tiles_folder, str(zoom), str(X))):
try: try:
os.mkdir(os.path.join(output_directory, "tiles", str(zoom), str(X))) os.mkdir(os.path.join(tiles_folder, str(zoom), str(X)))
except FileExistsError: except FileExistsError:
# Ignore this error, it means one other thread has already created the folder # Ignore this error, it means one other thread has already created the folder
pass pass
except Exception as e: except Exception as e:
raise e raise e
img.crop(box).convert('RGBA').save(os.path.join(output_directory, "tiles", str(zoom), str(X), f"{Y}.png")) img.crop(box).convert('RGBA').save(os.path.join(tiles_folder, str(zoom), str(X), f"{Y}.png"))
n += 1 n += 1
else: else:
raise Exception(f"{os.path.join(output_directory, 'screenshots', f'{f}_{n}_{zoom}.jpg')} missing") raise Exception(f"{os.path.join(screenshots_folder, f'{f}_{n}_{zoom}.jpg')} missing")
def merge_tiles(base_path, zoom, tile): def merge_tiles(base_path, zoom, tile):
X = tile[0] X = tile[0]
@@ -128,10 +129,10 @@ def merge_tiles(base_path, zoom, tile):
# Save the image # Save the image
dst.save(os.path.join(base_path, str(zoom - 1), str(X), f"{Y}.png"), quality=98) dst.save(os.path.join(base_path, str(zoom - 1), str(X), f"{Y}.png"), quality=98)
def compute_correction_factor(XY, n_width, n_height, map_config, zoom, output_directory, port): def compute_correction_factor(XY, n_width, n_height, map_config, zoom, screenshots_folder, port):
# Take screenshots at the given position # Take screenshots at the given position
take_screenshot(XY, 0, 0, map_config, zoom, output_directory, "calib", "ref", port) take_screenshot(XY, 0, 0, map_config, zoom, screenshots_folder, "calib", "ref", port)
calib_ref = Image.open(os.path.join(output_directory, "screenshots", f"calib_ref_{zoom}.jpg")) calib_ref = Image.open(os.path.join(screenshots_folder, f"calib_ref_{zoom}.jpg"))
# These calibration boxes are located at the edge of the interest region # These calibration boxes are located at the edge of the interest region
box1 = (calib_ref.width / 2 + n_width / 2 * 256 - 50, calib_ref.height / 2 - n_height / 2 * 256 + 10, box1 = (calib_ref.width / 2 + n_width / 2 * 256 - 50, calib_ref.height / 2 - n_height / 2 * 256 + 10,
@@ -150,10 +151,10 @@ def compute_correction_factor(XY, n_width, n_height, map_config, zoom, output_di
return None # Not enough variation return None # Not enough variation
# Take screenshot east and south of it # Take screenshot east and south of it
take_screenshot((XY[0] + n_width, XY[1]), 0, 0, map_config, zoom, output_directory, "calib", "lng", port) take_screenshot((XY[0] + n_width, XY[1]), 0, 0, map_config, zoom, screenshots_folder, "calib", "lng", port)
take_screenshot((XY[0], XY[1] + n_height), 0, 0, map_config, zoom, output_directory, "calib", "lat", port) take_screenshot((XY[0], XY[1] + n_height), 0, 0, map_config, zoom, screenshots_folder, "calib", "lat", port)
calib_lat = Image.open(os.path.join(output_directory, "screenshots", f"calib_lat_{zoom}.jpg")) calib_lat = Image.open(os.path.join(screenshots_folder, f"calib_lat_{zoom}.jpg"))
calib_lng = Image.open(os.path.join(output_directory, "screenshots", f"calib_lng_{zoom}.jpg")) calib_lng = Image.open(os.path.join(screenshots_folder, f"calib_lng_{zoom}.jpg"))
# Find the best correction factor to bring the two images to be equal on the longitude direction # Find the best correction factor to bring the two images to be equal on the longitude direction
best_err = None best_err = None
@@ -189,7 +190,7 @@ def compute_variation(imageA):
max = numpy.max((numpy.array(imageA))) max = numpy.max((numpy.array(imageA)))
return max - min return max - min
def take_screenshot(XY, n_width, n_height, map_config, zoom, output_directory, f, n, port, correction = (0, 0)): def take_screenshot(XY, n_width, n_height, map_config, zoom, screenshots_folder, f, n, port, correction = (0, 0)):
# Making PUT request # Making PUT request
# If the number of rows or columns is odd, we need to take the picture at the CENTER of the tile! # If the number of rows or columns is odd, we need to take the picture at the CENTER of the tile!
lat, lng = num_to_deg(XY[0] + (n_width % 2) / 2, XY[1] + (n_height % 2) / 2, zoom) lat, lng = num_to_deg(XY[0] + (n_width % 2) / 2, XY[1] + (n_height % 2) / 2, zoom)
@@ -236,7 +237,7 @@ def take_screenshot(XY, n_width, n_height, map_config, zoom, output_directory, f
sy = s_height / m_height sy = s_height / m_height
# Rotate, resize and save the screenshot # Rotate, resize and save the screenshot
screenshot.rotate(math.degrees(geo_data['northRotation'])).resize((int(sx * screenshot.width) + correction[0], int(sy * screenshot.height)+ correction[1] )).save(os.path.join(output_directory, "screenshots", f"{f}_{n}_{zoom}.jpg"), quality=98) screenshot.rotate(math.degrees(geo_data['northRotation'])).resize((int(sx * screenshot.width) + correction[0], int(sy * screenshot.height)+ correction[1] )).save(os.path.join(screenshots_folder, f"{f}_{n}_{zoom}.jpg"), quality=98)
def run(map_config, port): def run(map_config, port):
global tot_futs, fut_counter global tot_futs, fut_counter
@@ -246,26 +247,22 @@ def run(map_config, port):
screen_config = yaml.safe_load(sp) screen_config = yaml.safe_load(sp)
# Create output folders # Create output folders
output_directory = map_config['output_directory'] if not os.path.exists(map_config['tiles_folder']):
if not os.path.exists(output_directory): os.makedirs(map_config['tiles_folder'])
os.mkdir(output_directory)
if not os.path.exists(os.path.join(output_directory, "screenshots")): if not os.path.exists(os.path.join(map_config['screenshots_folder'])):
skip_screenshots = False skip_screenshots = False
replace_screenshots = True replace_screenshots = True
os.mkdir(os.path.join(output_directory, "screenshots")) os.makedirs(os.path.join(map_config['screenshots_folder']))
else: else:
skip_screenshots = map_config['skip_screenshots'] skip_screenshots = map_config['skip_screenshots']
replace_screenshots = map_config['replace_screenshots'] replace_screenshots = map_config['replace_screenshots']
if not os.path.exists(os.path.join(output_directory, "tiles")):
os.mkdir(os.path.join(output_directory, "tiles"))
# Compute the optimal zoom level # Compute the optimal zoom level
usable_width = screen_config['width'] - 400 # Keep a margin around the center usable_width = screen_config['width'] - 400 # Keep a margin around the center
usable_height = screen_config['height'] - 400 # Keep a margin around the center usable_height = screen_config['height'] - 400 # Keep a margin around the center
existing_zoom_levels = [int(f) for f in listdir(os.path.join(output_directory, "tiles")) if isdir(join(output_directory, "tiles", f))] existing_zoom_levels = [int(f) for f in listdir(os.path.join(map_config["tiles_folder"])) if isdir(join(map_config["tiles_folder"], f))]
if len(existing_zoom_levels) == 0: if len(existing_zoom_levels) == 0:
final_level = 1 final_level = 1
else: else:
@@ -331,27 +328,31 @@ def run(map_config, port):
print(f"Feature {f} of {len(features)}, taking screenshots...") print(f"Feature {f} of {len(features)}, taking screenshots...")
n = 0 n = 0
for XY in screenshots_XY: for XY in screenshots_XY:
if not os.path.exists(os.path.join(output_directory, "screenshots", f"{f}_{n}_{zoom}.jpg")) or replace_screenshots: if not os.path.exists(os.path.join(map_config['screenshots_folder'], f"{f}_{n}_{zoom}.jpg")) or replace_screenshots:
if n % 10 == 0 or correction is None: if n % 10 == 0 or correction is None:
new_correction = compute_correction_factor(XY, n_width, n_height, map_config, zoom, output_directory, port) new_correction = compute_correction_factor(XY, n_width, n_height, map_config, zoom, map_config['screenshots_folder'], port)
if new_correction is not None: if new_correction is not None:
correction = new_correction correction = new_correction
take_screenshot(XY, n_width, n_height, map_config, zoom, output_directory, f, n, port, correction if correction is not None else (0, 0)) take_screenshot(XY, n_width, n_height, map_config, zoom, map_config['screenshots_folder'], f, n, port, correction if correction is not None else (0, 0))
print_progress_bar(n + 1, len(screenshots_XY)) print_progress_bar(n + 1, len(screenshots_XY))
n += 1 n += 1
if map_config["screenshots_only"]:
return
########### Extract the tiles ########### Extract the tiles
print("Tiles extraction starting at: ", datetime.datetime.now()) print("Tiles extraction starting at: ", datetime.datetime.now())
if not os.path.exists(os.path.join(output_directory, "tiles", str(zoom))): if not os.path.exists(os.path.join(map_config["tiles_folder"], str(zoom))):
os.mkdir(os.path.join(output_directory, "tiles", str(zoom))) os.mkdir(os.path.join(map_config["tiles_folder"], str(zoom)))
params = { params = {
"f": f, "f": f,
"zoom": zoom, "zoom": zoom,
"output_directory": output_directory,
"n_width": n_width, "n_width": n_width,
"n_height": n_height, "n_height": n_height,
"screenshots_folder": map_config['screenshots_folder'],
"tiles_folder": map_config['tiles_folder']
} }
# Extract the tiles with parallel thread execution # Extract the tiles with parallel thread execution
@@ -373,10 +374,10 @@ def run(map_config, port):
########### Assemble tiles to get lower zoom levels ########### Assemble tiles to get lower zoom levels
print("Tiles merging start time: ", datetime.datetime.now()) print("Tiles merging start time: ", datetime.datetime.now())
for current_zoom in range(zoom, final_level, -1): for current_zoom in range(zoom, final_level, -1):
Xs = [int(d) for d in listdir(os.path.join(output_directory, "tiles", str(current_zoom))) if isdir(join(output_directory, "tiles", str(current_zoom), d))] Xs = [int(d) for d in listdir(os.path.join(map_config["tiles_folder"], str(current_zoom))) if isdir(join(map_config["tiles_folder"], str(current_zoom), d))]
existing_tiles = [] existing_tiles = []
for X in Xs: for X in Xs:
Ys = [int(f.removesuffix(".png")) for f in listdir(os.path.join(output_directory, "tiles", str(current_zoom), str(X))) if isfile(join(output_directory, "tiles", str(current_zoom), str(X), f))] Ys = [int(f.removesuffix(".png")) for f in listdir(os.path.join(map_config["tiles_folder"], str(current_zoom), str(X))) if isfile(join(map_config["tiles_folder"], str(current_zoom), str(X), f))]
for Y in Ys: for Y in Ys:
existing_tiles.append((X, Y)) existing_tiles.append((X, Y))
@@ -389,10 +390,10 @@ def run(map_config, port):
with futures.ThreadPoolExecutor() as executor: with futures.ThreadPoolExecutor() as executor:
print(f"Merging tiles for zoom level {current_zoom - 1}...") print(f"Merging tiles for zoom level {current_zoom - 1}...")
if not os.path.exists(os.path.join(output_directory, "tiles", str(current_zoom - 1))): if not os.path.exists(os.path.join(map_config["tiles_folder"], str(current_zoom - 1))):
os.mkdir(os.path.join(output_directory, "tiles", str(current_zoom - 1))) os.mkdir(os.path.join(map_config["tiles_folder"], str(current_zoom - 1)))
futs = [executor.submit(merge_tiles, os.path.join(output_directory, "tiles"), current_zoom, tile) for tile in tiles_to_produce] futs = [executor.submit(merge_tiles, os.path.join(map_config["tiles_folder"]), current_zoom, tile) for tile in tiles_to_produce]
tot_futs = len(futs) tot_futs = len(futs)
fut_counter = 0 fut_counter = 0
[fut.add_done_callback(done_callback) for fut in futs] [fut.add_done_callback(done_callback) for fut in futs]