Skip to content

Commit

Permalink
Reformat with black.
Browse files Browse the repository at this point in the history
  • Loading branch information
@bcoltin
bcoltin committed Jul 31, 2024
1 parent 7a8623f commit 2df1f82
Show file tree
Hide file tree
Showing 3 changed files with 189 additions and 77 deletions.
89 changes: 59 additions & 30 deletions localization/sparse_mapping/scripts/colmap/colmap.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
import struct
import subprocess

_MAX_IMAGE_ID = 2 ** 31 - 1
_MAX_IMAGE_ID = 2**31 - 1

# Access colmap sqlite database with cameras, images and matches
class COLMAPDatabase(sqlite3.Connection):
Expand All @@ -29,7 +29,6 @@ def __image_ids_to_pair_id(image_id1, image_id2):
image_id1, image_id2 = image_id2, image_id1
return image_id1 * _MAX_IMAGE_ID + image_id2


def cameras(self):
cameras = []
rows = self.execute("SELECT * FROM cameras")
Expand All @@ -46,22 +45,36 @@ def images(self):

def image_id(self, image_name):
images = []
rows = self.execute("SELECT image_id FROM images WHERE images.name = '%s'" % (image_name))
rows = self.execute(
"SELECT image_id FROM images WHERE images.name = '%s'" % (image_name)
)
if rows == None:
return None
return rows.fetchall()[0][0]

def num_matches(self, image1, image2):
if image1 == image2:
return 0
rows = self.execute("SELECT rows FROM matches WHERE pair_id = %d" % (self.__image_ids_to_pair_id(image1, image2)))
rows = self.execute(
"SELECT rows FROM matches WHERE pair_id = %d"
% (self.__image_ids_to_pair_id(image1, image2))
)
return next(rows)[0]


_FILE_PATH = os.path.dirname(os.path.realpath(__file__))

# creates a colmap project ini file based on a template, filling in key arguments (use in a "with:" block)
class ColmapProjectConfig:
def __init__(self, database_path, image_path, output_path, image_list=None, ini_file='mapper.ini', input_path=None):
def __init__(
self,
database_path,
image_path,
output_path,
image_list=None,
ini_file="mapper.ini",
input_path=None,
):
self.database_path = database_path
self.image_path = image_path
self.output_path = output_path
Expand All @@ -76,35 +89,38 @@ def __enter__(self):
self.config = tempfile.NamedTemporaryFile(delete=False)
if self.image_list:
self.image_config = tempfile.NamedTemporaryFile(delete=False)
with open (self.image_config.name, 'w' ) as f:
with open(self.image_config.name, "w") as f:
for image in self.image_list:
f.write(image + '\n')
f.write(image + "\n")

with open (os.path.join(_FILE_PATH, self.ini_file), 'r' ) as f:
with open(os.path.join(_FILE_PATH, self.ini_file), "r") as f:
content = f.read()
content = re.sub('DATABASE_PATH', self.database_path, content)
content = re.sub('IMAGE_PATH', self.image_path, content)
content = re.sub('OUTPUT_PATH', self.output_path, content)
content = re.sub("DATABASE_PATH", self.database_path, content)
content = re.sub("IMAGE_PATH", self.image_path, content)
content = re.sub("OUTPUT_PATH", self.output_path, content)
if self.input_path != None:
content = re.sub('INPUT_PATH', self.input_path, content)
image_list_set = 'image_list_path = %s' % (self.image_config.name) if self.image_list else ''
content = re.sub('IMAGE_LIST_SET', image_list_set, content)
with open (self.config.name, 'w' ) as cfg:
content = re.sub("INPUT_PATH", self.input_path, content)
image_list_set = (
"image_list_path = %s" % (self.image_config.name)
if self.image_list
else ""
)
content = re.sub("IMAGE_LIST_SET", image_list_set, content)
with open(self.config.name, "w") as cfg:
cfg.write(content)

return self

def __exit__(self, exception_type, exception_value, exception_traceback):
os.remove(self.config.name)
if self.image_list:
os.remove(self.image_config.name)


CameraModel = collections.namedtuple(
"CameraModel", ["model_id", "model_name", "num_params"]
)
Camera = collections.namedtuple(
"Camera", ["id", "model", "width", "height", "params"]
)
Camera = collections.namedtuple("Camera", ["id", "model", "width", "height", "params"])
BaseImage = collections.namedtuple(
"Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"]
)
Expand All @@ -117,6 +133,7 @@ class Image(BaseImage):
def qvec2rotmat(self):
return qvec2rotmat(self.qvec)


def _read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):
"""Read and unpack the next bytes from a binary file.
:param fid:
Expand All @@ -128,19 +145,25 @@ def _read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<")
data = fid.read(num_bytes)
return struct.unpack(endian_character + format_char_sequence, data)


def _skip_next_bytes(fid, num_bytes):
fid.seek(num_bytes, 1)


# snippets taken from colmap scripts/python/read_write_model.py
# this class reads a colmap model model and exposes fields of interest (mainly just the images)
class Model:
def __init__(self, fname):
self.filename = fname
self.__read_images_binary(os.path.join(fname, 'images.bin'))
self.__read_images_binary(os.path.join(fname, "images.bin"))
self.__analyze_model()

def __str__(self):
return '%d images, %g mean track length, %g mean reprojection error' % (self.num_images, self.mean_track_length, self.mean_reprojection_error)
return "%d images, %g mean track length, %g mean reprojection error" % (
self.num_images,
self.mean_track_length,
self.mean_reprojection_error,
)

def __read_images_binary(self, path_to_model_file):
"""
Expand Down Expand Up @@ -168,21 +191,23 @@ def __read_images_binary(self, path_to_model_file):
num_points2D = _read_next_bytes(
fid, num_bytes=8, format_char_sequence="Q"
)[0]
_skip_next_bytes(fid, num_bytes=24 * num_points2D) # faster to skip, we don't care
_skip_next_bytes(
fid, num_bytes=24 * num_points2D
) # faster to skip, we don't care
xys = None
point3D_ids = None
#x_y_id_s = read_next_bytes(
# x_y_id_s = read_next_bytes(
# fid,
# num_bytes=24 * num_points2D,
# format_char_sequence="ddq" * num_points2D,
#)
#xys = np.column_stack(
# )
# xys = np.column_stack(
# [
# tuple(map(float, x_y_id_s[0::3])),
# tuple(map(float, x_y_id_s[1::3])),
# ]
#)
#point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3])))
# )
# point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3])))
images[image_id] = Image(
id=image_id,
qvec=qvec,
Expand All @@ -195,15 +220,19 @@ def __read_images_binary(self, path_to_model_file):
self.images = images

def __analyze_model(self):
cmd = subprocess.Popen('colmap model_analyzer --path %s' % (self.filename), shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
cmd = subprocess.Popen(
"colmap model_analyzer --path %s" % (self.filename),
shell=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
output = cmd.communicate()[1]
if cmd.returncode != 0:
raise Exception('Model %s could not be analyzed.' % (self.filename))
raise Exception("Model %s could not be analyzed." % (self.filename))
result = dict()
m = re.search("Registered images: (\d+)", str(output))
self.num_images = int(m.groups()[0])
m = re.search("Mean track length: ([-+]?(?:\d*\.*\d+))", str(output))
self.mean_track_length = float(m.groups()[0])
m = re.search("Mean reprojection error: (.+)px", str(output))
self.mean_reprojection_error = float(m.groups()[0])

Loading

0 comments on commit 2df1f82

Please sign in to comment.