From 8c7f6abb1cd2211560ee3c4100b6ca93588f9df1 Mon Sep 17 00:00:00 2001 From: Pax1601 Date: Sun, 10 Mar 2024 13:43:09 +0100 Subject: [PATCH] Added more configurability to map generator --- scripts/python/map_generator/main.py | 21 ++++-- scripts/python/map_generator/map_generator.py | 73 ++++++++++--------- 2 files changed, 50 insertions(+), 44 deletions(-) diff --git a/scripts/python/map_generator/main.py b/scripts/python/map_generator/main.py index 630e415a..ccb60c1f 100644 --- a/scripts/python/map_generator/main.py +++ b/scripts/python/map_generator/main.py @@ -1,4 +1,4 @@ -import sys +import os import yaml import json import requests @@ -16,9 +16,12 @@ parser = argparse.ArgumentParser( epilog='Hit the DCS Olympus Discord for more information') 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('-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('-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('-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() @@ -36,6 +39,10 @@ with open('configs/screen_properties.yml', 'r') as sp: map_config = yaml.safe_load(cp) 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(f"-> Screen width: {screen_config['width']}px") @@ -43,6 +50,8 @@ with open('configs/screen_properties.yml', 'r') as sp: print("Map parameters:") 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"-> Zoom factor: {map_config['zoom_factor']}") @@ -88,8 +97,4 @@ with open('configs/screen_properties.yml', 'r') as sp: print(f"Estimated number of tiles: {tiles_num}") print(f"Estimated number of screenshots: {screenshots_num}") - map_generator.run(map_config, port) - - - - + map_generator.run(map_config, port) \ No newline at end of file diff --git a/scripts/python/map_generator/map_generator.py b/scripts/python/map_generator/map_generator.py index a0f38ae2..048dfce9 100644 --- a/scripts/python/map_generator/map_generator.py +++ b/scripts/python/map_generator/map_generator.py @@ -58,14 +58,15 @@ def done_callback(fut): def extract_tiles(n, screenshots_XY, params): f = params['f'] zoom = params['zoom'] - output_directory = params['output_directory'] n_width = params['n_width'] n_height = params['n_height'] + screenshots_folder = params['screenshots_folder'] + tiles_folder = params['tiles_folder'] 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 - 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 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 # 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: - 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: # Ignore this error, it means one other thread has already created the folder pass except Exception as 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 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): X = tile[0] @@ -128,10 +129,10 @@ def merge_tiles(base_path, zoom, tile): # Save the image 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_screenshot(XY, 0, 0, map_config, zoom, output_directory, "calib", "ref", port) - calib_ref = Image.open(os.path.join(output_directory, "screenshots", f"calib_ref_{zoom}.jpg")) + take_screenshot(XY, 0, 0, map_config, zoom, screenshots_folder, "calib", "ref", port) + 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 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 # 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], XY[1] + n_height), 0, 0, map_config, zoom, output_directory, "calib", "lat", port) - calib_lat = Image.open(os.path.join(output_directory, "screenshots", f"calib_lat_{zoom}.jpg")) - calib_lng = Image.open(os.path.join(output_directory, "screenshots", f"calib_lng_{zoom}.jpg")) + 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, screenshots_folder, "calib", "lat", port) + calib_lat = Image.open(os.path.join(screenshots_folder, f"calib_lat_{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 best_err = None @@ -189,7 +190,7 @@ def compute_variation(imageA): max = numpy.max((numpy.array(imageA))) 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 # 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) @@ -236,7 +237,7 @@ def take_screenshot(XY, n_width, n_height, map_config, zoom, output_directory, f sy = s_height / m_height # 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): global tot_futs, fut_counter @@ -246,26 +247,22 @@ def run(map_config, port): screen_config = yaml.safe_load(sp) # Create output folders - output_directory = map_config['output_directory'] - if not os.path.exists(output_directory): - os.mkdir(output_directory) + if not os.path.exists(map_config['tiles_folder']): + os.makedirs(map_config['tiles_folder']) - 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 replace_screenshots = True - os.mkdir(os.path.join(output_directory, "screenshots")) + os.makedirs(os.path.join(map_config['screenshots_folder'])) else: skip_screenshots = map_config['skip_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 usable_width = screen_config['width'] - 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: final_level = 1 else: @@ -331,27 +328,31 @@ def run(map_config, port): print(f"Feature {f} of {len(features)}, taking screenshots...") n = 0 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: - 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: 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)) n += 1 + if map_config["screenshots_only"]: + return + ########### Extract the tiles print("Tiles extraction starting at: ", datetime.datetime.now()) - if not os.path.exists(os.path.join(output_directory, "tiles", str(zoom))): - os.mkdir(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(map_config["tiles_folder"], str(zoom))) params = { "f": f, "zoom": zoom, - "output_directory": output_directory, "n_width": n_width, "n_height": n_height, + "screenshots_folder": map_config['screenshots_folder'], + "tiles_folder": map_config['tiles_folder'] } # Extract the tiles with parallel thread execution @@ -373,10 +374,10 @@ def run(map_config, port): ########### Assemble tiles to get lower zoom levels print("Tiles merging start time: ", datetime.datetime.now()) 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 = [] 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: existing_tiles.append((X, Y)) @@ -389,10 +390,10 @@ def run(map_config, port): with futures.ThreadPoolExecutor() as executor: 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))): - os.mkdir(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(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) fut_counter = 0 [fut.add_done_callback(done_callback) for fut in futs]