Skip to content

Add write_model functions for python scripts #753

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

Merged
merged 2 commits into from
Dec 9, 2019
Merged

Add write_model functions for python scripts #753

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d2d5225
Add write_model function
ClementPinard Dec 5, 2019
70b3ed4
Put arg parser inside main
ClementPinard Dec 9, 2019
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
183 changes: 177 additions & 6 deletions scripts/python/read_model.py → scripts/python/read_write_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
import collections
import numpy as np
import struct
import argparse


CameraModel = collections.namedtuple(
Expand All @@ -45,6 +46,7 @@
Point3D = collections.namedtuple(
"Point3D", ["id", "xyz", "rgb", "error", "image_ids", "point2D_idxs"])


class Image(BaseImage):
def qvec2rotmat(self):
return qvec2rotmat(self.qvec)
Expand All @@ -63,8 +65,10 @@ def qvec2rotmat(self):
CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5),
CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12)
}
CAMERA_MODEL_IDS = dict([(camera_model.model_id, camera_model) \
CAMERA_MODEL_IDS = dict([(camera_model.model_id, camera_model)
for camera_model in CAMERA_MODELS])
CAMERA_MODEL_NAMES = dict([(camera_model.model_name, camera_model)
for camera_model in CAMERA_MODELS])


def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):
Expand All @@ -79,6 +83,22 @@ def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):
return struct.unpack(endian_character + format_char_sequence, data)


def write_next_bytes(fid, data, format_char_sequence, endian_character="<"):
"""pack and write to a binary file.
:param fid:
:param data: data to send, if multiple elements are sent at the same time,
they should be encapsuled either in a list or a tuple
:param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.
should be the same length as the data list or tuple
:param endian_character: Any of {@, =, <, >, !}
"""
if isinstance(data, (list, tuple)):
bytes = struct.pack(endian_character + format_char_sequence, *data)
else:
bytes = struct.pack(endian_character + format_char_sequence, data)
fid.write(bytes)


def read_cameras_text(path):
"""
see: src/base/reconstruction.cc
Expand Down Expand Up @@ -134,6 +154,43 @@ def read_cameras_binary(path_to_model_file):
return cameras


def write_cameras_text(cameras, path):
"""
see: src/base/reconstruction.cc
void Reconstruction::WriteCamerasText(const std::string& path)
void Reconstruction::ReadCamerasText(const std::string& path)
"""
HEADER = '# Camera list with one line of data per camera:\n'
'# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n'
'# Number of cameras: {}\n'.format(len(cameras))
with open(path, "w") as fid:
fid.write(HEADER)
for _, cam in cameras.items():
to_write = [cam.id, cam.model, cam.width, cam.height, *cam.params]
line = " ".join([str(elem) for elem in to_write])
fid.write(line + "\n")


def write_cameras_binary(cameras, path_to_model_file):
"""
see: src/base/reconstruction.cc
void Reconstruction::WriteCamerasBinary(const std::string& path)
void Reconstruction::ReadCamerasBinary(const std::string& path)
"""
with open(path_to_model_file, "wb") as fid:
write_next_bytes(fid, len(cameras), "Q")
for _, cam in cameras.items():
model_id = CAMERA_MODEL_NAMES[cam.model].model_id
camera_properties = [cam.id,
model_id,
cam.width,
cam.height]
write_next_bytes(fid, camera_properties, "iiQQ")
for p in cam.params:
write_next_bytes(fid, float(p), "d")
return cameras


def read_images_text(path):
"""
see: src/base/reconstruction.cc
Expand Down Expand Up @@ -200,6 +257,55 @@ def read_images_binary(path_to_model_file):
return images


def write_images_text(images, path):
"""
see: src/base/reconstruction.cc
void Reconstruction::ReadImagesText(const std::string& path)
void Reconstruction::WriteImagesText(const std::string& path)
"""
if len(images) == 0:
mean_observations = 0
else:
mean_observations = sum((len(img.point3D_ids) for _, img in images.items()))/len(images)
HEADER = '# Image list with two lines of data per image:\n'
'# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n'
'# POINTS2D[] as (X, Y, POINT3D_ID)\n'
'# Number of images: {}, mean observations per image: {}\n'.format(len(images), mean_observations)

with open(path, "w") as fid:
fid.write(HEADER)
for _, img in images.items():
image_header = [img.id, *img.qvec, *img.tvec, img.camera_id, img.name]
first_line = " ".join(map(str, image_header))
fid.write(first_line + "\n")

points_strings = []
for xy, point3D_id in zip(img.xys, img.point3D_ids):
points_strings.append(" ".join(map(str, [*xy, point3D_id])))
fid.write(" ".join(points_strings) + "\n")


def write_images_binary(images, path_to_model_file):
"""
see: src/base/reconstruction.cc
void Reconstruction::ReadImagesBinary(const std::string& path)
void Reconstruction::WriteImagesBinary(const std::string& path)
"""
with open(path_to_model_file, "wb") as fid:
write_next_bytes(fid, len(images), "Q")
for _, img in images.items():
write_next_bytes(fid, img.id, "i")
write_next_bytes(fid, img.qvec.tolist(), "dddd")
write_next_bytes(fid, img.tvec.tolist(), "ddd")
write_next_bytes(fid, img.camera_id, "i")
for char in img.name:
write_next_bytes(fid, char.encode("utf-8"), "c")
write_next_bytes(fid, b"\x00", "c")
write_next_bytes(fid, len(img.point3D_ids), "Q")
for xy, p3d_id in zip(img.xys, img.point3D_ids):
write_next_bytes(fid, [*xy, p3d_id], "ddq")


def read_points3D_text(path):
"""
see: src/base/reconstruction.cc
Expand Down Expand Up @@ -257,6 +363,50 @@ def read_points3d_binary(path_to_model_file):
return points3D


def write_points3D_text(points3D, path):
"""
see: src/base/reconstruction.cc
void Reconstruction::ReadPoints3DText(const std::string& path)
void Reconstruction::WritePoints3DText(const std::string& path)
"""
if len(points3D) == 0:
mean_track_length = 0
else:
mean_track_length = sum((len(pt.image_ids) for _, pt in points3D.items()))/len(points3D)
HEADER = '# 3D point list with one line of data per point:\n'
'# POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n'
'# Number of points: {}, mean track length: {}\n'.format(len(points3D), mean_track_length)

with open(path, "w") as fid:
fid.write(HEADER)
for _, pt in points3D.items():
point_header = [pt.id, *pt.xyz, *pt.rgb, pt.error]
fid.write(" ".join(map(str, point_header)) + " ")
track_strings = []
for image_id, point2D in zip(pt.image_ids, pt.point2D_idxs):
track_strings.append(" ".join(map(str, [image_id, point2D])))
fid.write(" ".join(track_strings) + "\n")


def write_points3d_binary(points3D, path_to_model_file):
"""
see: src/base/reconstruction.cc
void Reconstruction::ReadPoints3DBinary(const std::string& path)
void Reconstruction::WritePoints3DBinary(const std::string& path)
"""
with open(path_to_model_file, "wb") as fid:
write_next_bytes(fid, len(points3D), "Q")
for _, pt in points3D.items():
write_next_bytes(fid, pt.id, "Q")
write_next_bytes(fid, pt.xyz.tolist(), "ddd")
write_next_bytes(fid, pt.rgb.tolist(), "BBB")
write_next_bytes(fid, pt.error, "d")
track_length = pt.image_ids.shape[0]
write_next_bytes(fid, track_length, "Q")
for image_id, point2D_id in zip(pt.image_ids, pt.point2D_idxs):
write_next_bytes(fid, [image_id, point2D_id], "ii")


def read_model(path, ext):
if ext == ".txt":
cameras = read_cameras_text(os.path.join(path, "cameras" + ext))
Expand All @@ -269,6 +419,18 @@ def read_model(path, ext):
return cameras, images, points3D


def write_model(cameras, images, points3D, path, ext):
if ext == ".txt":
write_cameras_text(cameras, os.path.join(path, "cameras" + ext))
write_images_text(images, os.path.join(path, "images" + ext))
write_points3D_text(points3D, os.path.join(path, "points3D") + ext)
else:
write_cameras_binary(cameras, os.path.join(path, "cameras" + ext))
write_images_binary(images, os.path.join(path, "images" + ext))
write_points3d_binary(points3D, os.path.join(path, "points3D") + ext)
return cameras, images, points3D


def qvec2rotmat(qvec):
return np.array([
[1 - 2 * qvec[2]**2 - 2 * qvec[3]**2,
Expand Down Expand Up @@ -297,16 +459,25 @@ def rotmat2qvec(R):


def main():
if len(sys.argv) != 3:
print("Usage: python read_model.py path/to/model/folder [.txt,.bin]")
return

cameras, images, points3D = read_model(path=sys.argv[1], ext=sys.argv[2])
parser = argparse.ArgumentParser(description='Read and write COLMAP binary and text models')
parser.add_argument('input_model', help='path to input model folder')
parser.add_argument('input_format', choices=['.bin', '.txt'],
help='input model format')
parser.add_argument('--output_model', metavar='PATH',
help='path to output model folder')
parser.add_argument('--output_format', choices=['.bin', '.txt'],
help='outut model format', default='.txt')
args = parser.parse_args()

cameras, images, points3D = read_model(path=args.input_model, ext=args.input_format)

print("num_cameras:", len(cameras))
print("num_images:", len(images))
print("num_points3D:", len(points3D))

if args.output_model is not None:
write_model(cameras, images, points3D, path=args.output_model, ext=args.output_format)


if __name__ == "__main__":
main()
24 changes: 23 additions & 1 deletion scripts/python/test_read_model.py → scripts/python/test_read_write_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,8 @@
# Author: Johannes L. Schoenberger (jsch-at-demuc-dot-de)

import numpy as np
from read_model import read_model
from read_write_model import read_model, write_model
from tempfile import mkdtemp


def compare_cameras(cameras1, cameras2):
Expand Down Expand Up @@ -90,6 +91,27 @@ def main():
compare_points(points3D_txt, points3D_bin)

print("... text and binary models are equal.")
print("Saving text model and reloading it ...")

tmpdir = mkdtemp()
write_model(cameras_bin, images_bin, points3D_bin, tmpdir, ext='.txt')
cameras_txt, images_txt, points3D_txt = \
read_model(tmpdir, ext=".txt")
compare_cameras(cameras_txt, cameras_bin)
compare_images(images_txt, images_bin)
compare_points(points3D_txt, points3D_bin)

print("... saved text and loaded models are equal.")
print("Saving binary model and reloading it ...")

write_model(cameras_bin, images_bin, points3D_bin, tmpdir, ext='.bin')
cameras_bin, images_bin, points3D_bin = \
read_model(tmpdir, ext=".bin")
compare_cameras(cameras_txt, cameras_bin)
compare_images(images_txt, images_bin)
compare_points(points3D_txt, points3D_bin)

print("... saved binary and loaded models are equal.")


if __name__ == "__main__":
Expand Down