From a6189195cb1ef095f0caed24ada03bb49e86a908 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Wed, 28 Feb 2024 16:46:33 +0000 Subject: [PATCH 1/6] speedup resize --- convert.py | 29 ++++++++++------------------- 1 file changed, 10 insertions(+), 19 deletions(-) diff --git a/convert.py b/convert.py index ec41299e8..06e196abe 100644 --- a/convert.py +++ b/convert.py @@ -211,26 +211,17 @@ def concat_alpha(seg_path): # Copy each file from the source directory to the destination directory for file in files: source_file = os.path.join(args.source_path, "images", file) - - destination_file = os.path.join(args.source_path, "images_2", file) - shutil.copy2(source_file, destination_file) - exit_code = os.system("mogrify -resize 50% " + destination_file) - if exit_code != 0: - logging.error(f"50% resize failed with code {exit_code}. Exiting.") - exit(EXIT_FAIL) - - destination_file = os.path.join(args.source_path, "images_4", file) - shutil.copy2(source_file, destination_file) - exit_code = os.system("mogrify -resize 25% " + destination_file) - if exit_code != 0: - logging.error(f"25% resize failed with code {exit_code}. Exiting.") - exit(EXIT_FAIL) - - destination_file = os.path.join(args.source_path, "images_8", file) - shutil.copy2(source_file, destination_file) - exit_code = os.system("mogrify -resize 12.5% " + destination_file) + output_file2 = os.path.join(args.source_path, "images_2", file) + output_file4 = os.path.join(args.source_path, "images_4", file) + output_file8 = os.path.join(args.source_path, "images_8", file) + + # generate the resized images in a single call + exit_code = os.system("convert " + source_file + + " \( +clone -resize 50% -write " + output_file2 + " +delete \)" + + " \( +clone -resize 25% -write " + output_file4 + " +delete \)" + + " -resize 12.5% " + output_file8) if exit_code != 0: - logging.error(f"12.5% resize failed with code {exit_code}. Exiting.") + logging.error(f"resize failed with code {exit_code}. Exiting.") exit(EXIT_FAIL) print("Done.") From d07fb2208dc5a654bc319734f140cf0509aada0b Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Thu, 29 Feb 2024 15:49:44 +0000 Subject: [PATCH 2/6] handle masks correctly --- convert.py | 196 ++++++++++++++++++++++++++--------------------------- 1 file changed, 95 insertions(+), 101 deletions(-) diff --git a/convert.py b/convert.py index 06e196abe..62947a810 100644 --- a/convert.py +++ b/convert.py @@ -10,6 +10,7 @@ # import os +import subprocess import logging from argparse import ArgumentParser import shutil @@ -37,7 +38,17 @@ magick_command = '"{}"'.format(args.magick_executable) if len(args.magick_executable) > 0 else "magick" use_gpu = 1 if not args.no_gpu else 0 -EXIT_FAIL = 1 +# configure logging +logging.basicConfig(level = logging.INFO) + +# execute a command after logging it and propagate failure correctly +def exec(cmd): + logging.info(f"Executing: {cmd}") + try: + subprocess.check_call(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT, shell=True) + except subprocess.CalledProcessError as e: + logging.error(f"Command failed with code {e.returncode}. Exiting.") + exit(e.returncode) if not args.skip_matching: os.makedirs(args.source_path + "/distorted/sparse", exist_ok=True) @@ -49,19 +60,13 @@ --ImageReader.single_camera 1 \ --ImageReader.camera_model " + args.camera + " \ --SiftExtraction.use_gpu " + str(use_gpu) - exit_code = os.system(feat_extracton_cmd) - if exit_code != 0: - logging.error(f"Feature extraction failed with code {exit_code}. Exiting.") - exit(EXIT_FAIL) + exec(feat_extracton_cmd) ## Feature matching feat_matching_cmd = colmap_command + " exhaustive_matcher \ --database_path " + args.source_path + "/distorted/database.db \ --SiftMatching.use_gpu " + str(use_gpu) - exit_code = os.system(feat_matching_cmd) - if exit_code != 0: - logging.error(f"Feature matching failed with code {exit_code}. Exiting.") - exit(EXIT_FAIL) + exec(feat_matching_cmd) ### Bundle adjustment # The default Mapper tolerance is unnecessarily large, @@ -71,11 +76,7 @@ --image_path " + args.source_path + "/input \ --output_path " + args.source_path + "/distorted/sparse \ --Mapper.ba_global_function_tolerance=0.000001") - exit_code = os.system(mapper_cmd) - if exit_code != 0: - logging.error(f"Mapper failed with code {exit_code}. Exiting.") - exit(EXIT_FAIL) - + exec(mapper_cmd) # select the largest submodel i = 0 @@ -107,97 +108,90 @@ --input_path " + distorted_sparse_path + " \ --output_path " + args.source_path + "\ --output_type COLMAP") - -exit_code = os.system(img_undist_cmd) -if exit_code != 0: - logging.error(f"image_undistorter failed with code {exit_code}. Exiting.") - exit(EXIT_FAIL) +exec(img_undist_cmd) -def remove_dir_if_exist(path): - if Path(path).exists(): - shutil.rmtree(path) - +# Handle masks if args.masks_path is not None: - remove_dir_if_exist(args.source_path + "/alpha_distorted_sparse_txt/") - Path(args.source_path + "/alpha_distorted_sparse_txt/").mkdir(exist_ok=True) - # We need to "hack" colmap to undistort segmentation maps modify paths + # We need to modify the colmap database to reference the mask images + # which are always in png format. + mask_model_path = args.masks_path + "/model" + Path(mask_model_path).mkdir(exist_ok=True) + # First convert model to text format model_converter_cmd = (colmap_command + " model_converter \ --input_path " + distorted_sparse_path + " \ - --output_path " + args.source_path + "/alpha_distorted_sparse_txt/ \ + --output_path " + mask_model_path + " \ --output_type TXT") - exit_code = os.system(model_converter_cmd) - if exit_code != 0: - logging.error(f"model_converter failed with code {exit_code}. Exiting.") - exit(EXIT_FAIL) - - # replace '.jpg' to '.png' - with open(args.source_path + "/alpha_distorted_sparse_txt/images.txt", "r+") as f: - images_txt = f.read() - images_txt = images_txt.replace('.jpg', '.png') - f.seek(0) - f.write(images_txt) - f.truncate() - - # Undistort alpha masks + exec(model_converter_cmd) + + # read images.txt + with open(mask_model_path + "/images.txt", 'r') as file: + lines = file.readlines() + + # replace image filenames with png extensions (and keep the list of renames for later) + filenames = [] + l = 0 + for i in range(len(lines)): + if lines[i].startswith("#"): + # skip comments + continue + if l % 2 == 0: + # handle every second line + words = lines[i].rstrip().split(" ") + filename = words[-1].split(".") + filename[-1] = "png" + new_filename = ".".join(filename) + filenames.append([words[-1], new_filename]) + words[-1] = new_filename + lines[i] = " ".join(words) + "\n" + l += 1 + + # write modified images.txt + with open(mask_model_path + "/images.txt", 'w') as file: + file.writelines(lines) + + # Undistort mask images seg_undist_cmd = (colmap_command + " image_undistorter \ --image_path " + args.masks_path + " \ - --input_path " + args.source_path + "/alpha_distorted_sparse_txt/ \ - --output_path " + args.source_path + "/alpha_undistorted_sparse \ + --input_path " + mask_model_path + " \ + --output_path " + args.masks_path + "/undistorted \ --output_type COLMAP") - exit_code = os.system(seg_undist_cmd) - if exit_code != 0: - logging.error(f"image_undistorter for segs failed with code {exit_code}. Exiting.") - exit(EXIT_FAIL) - - # switch images - remove_dir_if_exist(f'{args.source_path}/alpha_undistorted_sparse/alphas') - Path(f'{args.source_path}/alpha_undistorted_sparse/images').replace(f'{args.source_path}/alpha_undistorted_sparse/alphas') - remove_dir_if_exist(f'{args.source_path}/images_src/') - Path(f'{args.source_path}/images/').replace(f'{args.source_path}/images_src/') - - # concat undistorted images with undistorted alpha masks - TODO: make parallel - remove_dir_if_exist(f'{args.source_path}/images/') - Path(f'{args.source_path}/images/').mkdir() - - def concat_alpha(seg_path): - seg = Image.open(seg_path).convert('L') - img = Image.open(f'{args.source_path}/images_src/{Path(seg_path).stem}.jpg') - img.putalpha(seg) - img.save(f'{args.source_path}/images/{Path(seg_path).stem}.png') - - all_masks_paths = glob(args.source_path + "/alpha_undistorted_sparse/alphas/*.png") - with mp.Pool() as pool: - list(tqdm(pool.imap_unordered(concat_alpha, all_masks_paths), total=len(all_masks_paths))) - - # switch models - remove_dir_if_exist(f'{args.source_path}/sparse_src/') - Path(f'{args.source_path}/sparse').replace(f'{args.source_path}/sparse_src/') - Path(f'{args.source_path}/alpha_undistorted_sparse/sparse').replace(f'{args.source_path}/sparse/') - -if args.generate_text_model: - ### Convert model to text format so we can read cameras - convert_cmd = (colmap_command + " model_converter \ - --input_path " + args.source_path + "/sparse" + " \ - --output_path " + args.source_path + "/sparse" + " \ - --output_type TXT") - exit_code = os.system(convert_cmd) - if exit_code != 0: - logging.error(f"Convert failed with code {exit_code}. Exiting.") - exit(exit_code) - -# move all files from sparse into sparse/0, as train.py expects it -files = os.listdir(args.source_path + "/sparse") -os.makedirs(args.source_path + "/sparse/0", exist_ok=True) -# Copy each file from the source directory to the destination directory -for file in files: - if file == "0": - continue - source_file = os.path.join(args.source_path, "sparse", file) - destination_file = os.path.join(args.source_path, "sparse", "0", file) - shutil.move(source_file, destination_file) + exec(seg_undist_cmd) + + # combine undistorted color and mask images + def combine(color_path, alpha_path, output_path): + alpha = Image.open(alpha_path).convert('L') + clr = Image.open(color_path) + clr.putalpha(alpha) + clr.save(output_path) + + for i in range(len(filenames)): + color_image = args.source_path + "/images/" + filenames[i][0] + mask_image = args.masks_path + "/undistorted/images/" + filenames[i][1] + output_image = args.source_path + "/images/" + filenames[i][1] + combine(color_image, mask_image, color_image) + + # copy the modified database to final location for use in next step + target_path = args.source_path + "/sparse/0" + Path(target_path).mkdir(exist_ok=True) + files = os.listdir(mask_model_path) + for file in files: + source_file = os.path.join(mask_model_path, file) + destination_file = os.path.join(target_path, file) + shutil.move(source_file, destination_file) +else: + # move all files from sparse into sparse/0, as train.py expects it + files = os.listdir(args.source_path + "/sparse") + os.makedirs(args.source_path + "/sparse/0", exist_ok=True) + # Copy each file from the source directory to the destination directory + for file in files: + if file == "0": + continue + source_file = os.path.join(args.source_path, "sparse", file) + destination_file = os.path.join(args.source_path, "sparse", "0", file) + shutil.move(source_file, destination_file) if(args.resize): print("Copying and resizing...") @@ -216,12 +210,12 @@ def concat_alpha(seg_path): output_file8 = os.path.join(args.source_path, "images_8", file) # generate the resized images in a single call - exit_code = os.system("convert " + source_file + - " \( +clone -resize 50% -write " + output_file2 + " +delete \)" + - " \( +clone -resize 25% -write " + output_file4 + " +delete \)" + - " -resize 12.5% " + output_file8) - if exit_code != 0: - logging.error(f"resize failed with code {exit_code}. Exiting.") - exit(EXIT_FAIL) + generate_thumbnails_cmd = ("convert " + # resize input file, uses less memory + f"{source_file}[50%]" + f" -write mpr:thumb -write {output_file2} +delete" + f" mpr:thumb -resize 50% -write mpr:thumb -write {output_file4} +delete" + f" mpr:thumb -resize 50% {output_file8}") + exec(generate_thumbnails_cmd) print("Done.") From ede073dda6ae362e671644dad2d4e74db73c468d Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Thu, 29 Feb 2024 16:12:23 +0000 Subject: [PATCH 3/6] fix --- convert.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/convert.py b/convert.py index 62947a810..c67bc2039 100644 --- a/convert.py +++ b/convert.py @@ -173,12 +173,14 @@ def combine(color_path, alpha_path, output_path): output_image = args.source_path + "/images/" + filenames[i][1] combine(color_image, mask_image, color_image) - # copy the modified database to final location for use in next step + # copy the modified database to final location for use in training target_path = args.source_path + "/sparse/0" Path(target_path).mkdir(exist_ok=True) - files = os.listdir(mask_model_path) + + source_path = args.masks_path + "/undistorted/sparse" + files = os.listdir(source_path) for file in files: - source_file = os.path.join(mask_model_path, file) + source_file = os.path.join(source_path, file) destination_file = os.path.join(target_path, file) shutil.move(source_file, destination_file) else: @@ -193,7 +195,7 @@ def combine(color_path, alpha_path, output_path): destination_file = os.path.join(args.source_path, "sparse", "0", file) shutil.move(source_file, destination_file) -if(args.resize): +if (args.resize): print("Copying and resizing...") # Resize images. From cf48b503767dd8761797aadeb24755ed85639283 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Thu, 29 Feb 2024 17:00:33 +0000 Subject: [PATCH 4/6] small --- convert.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/convert.py b/convert.py index c67bc2039..bfb1ea16e 100644 --- a/convert.py +++ b/convert.py @@ -171,7 +171,7 @@ def combine(color_path, alpha_path, output_path): color_image = args.source_path + "/images/" + filenames[i][0] mask_image = args.masks_path + "/undistorted/images/" + filenames[i][1] output_image = args.source_path + "/images/" + filenames[i][1] - combine(color_image, mask_image, color_image) + combine(color_image, mask_image, output_image) # copy the modified database to final location for use in training target_path = args.source_path + "/sparse/0" From 542d3f9cb26c5b59645a8d8d20f547867a783b29 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Fri, 1 Mar 2024 11:59:49 +0000 Subject: [PATCH 5/6] pipe err to err --- convert.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/convert.py b/convert.py index bfb1ea16e..7b2435067 100644 --- a/convert.py +++ b/convert.py @@ -45,7 +45,7 @@ def exec(cmd): logging.info(f"Executing: {cmd}") try: - subprocess.check_call(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT, shell=True) + subprocess.check_call(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.STDERR, shell=True) except subprocess.CalledProcessError as e: logging.error(f"Command failed with code {e.returncode}. Exiting.") exit(e.returncode) From 487d0d28dc0311c65cdffc4846eed4684aa0126f Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Fri, 1 Mar 2024 12:06:22 +0000 Subject: [PATCH 6/6] revert --- convert.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/convert.py b/convert.py index 7b2435067..bfb1ea16e 100644 --- a/convert.py +++ b/convert.py @@ -45,7 +45,7 @@ def exec(cmd): logging.info(f"Executing: {cmd}") try: - subprocess.check_call(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.STDERR, shell=True) + subprocess.check_call(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT, shell=True) except subprocess.CalledProcessError as e: logging.error(f"Command failed with code {e.returncode}. Exiting.") exit(e.returncode)