Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Logging and mask updates, resize speedup #8

Merged
merged 6 commits into from
Mar 1, 2024
Merged

Logging and mask updates, resize speedup #8

Changes from 4 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
a618919
speedup resize
slimbuck Feb 28, 2024
d07fb22
handle masks correctly
slimbuck Feb 29, 2024
ede073d
fix
slimbuck Feb 29, 2024
cf48b50
small
slimbuck Feb 29, 2024
542d3f9
pipe err to err
slimbuck Mar 1, 2024
487d0d2
revert
slimbuck Mar 1, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
221 changes: 104 additions & 117 deletions convert.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
#

import os
import subprocess
import logging
from argparse import ArgumentParser
import shutil
Expand Down Expand Up @@ -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)
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is this less than 256?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yep. The problem before was actually os.system which returns a 16-bit number on unix "whose low byte is the signal number that killed the process, and whose high byte is the exit status". This was causing the strange numbers.

subprocess.check_call doesn't do this, but supplies returncode directly in the exception.


if not args.skip_matching:
os.makedirs(args.source_path + "/distorted/sparse", exist_ok=True)
Expand All @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -107,99 +108,94 @@
--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)

if(args.resize):
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, output_image)

# 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)

source_path = args.masks_path + "/undistorted/sparse"
files = os.listdir(source_path)
for file in files:
source_file = os.path.join(source_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...")

# Resize images.
Expand All @@ -211,26 +207,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)
if exit_code != 0:
logging.error(f"12.5% resize failed with code {exit_code}. Exiting.")
exit(EXIT_FAIL)
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
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.")