Add scripts for colmap input/output and training dataset conversion (…

…DTU and ETH 3D)
FangjinhuaWang · Sep 15, 2021 · 01461a1 · 01461a1
1 parent 7fbe780
commit 01461a1
Show file tree

Hide file tree

Showing 9 changed files with 706 additions and 292 deletions.
diff --git a/colmap_input.py b/colmap_input.py
diff --git a/colmap_output.py b/colmap_output.py
@@ -0,0 +1,144 @@
+import argparse
+import numpy as np
+import os
+import shutil
+from colmap_input import Camera, Image
+from datasets.data_io import read_cam_file, read_map, read_pair_file, save_map
+from typing import Dict, List, Tuple
+from PIL import Image as PilImage
+
+
+def rotation_matrix_to_quaternion(rot: np.ndarray) -> List[float]:
+    rxx, ryx, rzx, rxy, ryy, rzy, rxz, ryz, rzz = rot.flat
+    k = np.array([
+        [rxx - ryy - rzz, 0, 0, 0],
+        [ryx + rxy, ryy - rxx - rzz, 0, 0],
+        [rzx + rxz, rzy + ryz, rzz - rxx - ryy, 0],
+        [ryz - rzy, rzx - rxz, rxy - ryx, rxx + ryy + rzz]]) / 3.0
+    eigenvalues, eigenvectors = np.linalg.eigh(k)
+    qvec = eigenvectors[[3, 0, 1, 2], np.argmax(eigenvalues)]
+    if qvec[0] < 0:
+        qvec *= -1
+    return [qvec[0], qvec[1], qvec[2], qvec[3]]
+
+
+def create_output_dirs(path: str):
+    os.makedirs(path, exist_ok=True)
+    os.makedirs(os.path.join(path, "images"), exist_ok=True)
+    os.makedirs(os.path.join(path, "sparse"), exist_ok=True)
+    os.makedirs(os.path.join(path, "stereo"), exist_ok=True)
+    os.makedirs(os.path.join(path, "stereo", "confidence_maps"), exist_ok=True)
+    os.makedirs(os.path.join(path, "stereo", "consistency_graphs"), exist_ok=True)
+    os.makedirs(os.path.join(path, "stereo", "depth_maps"), exist_ok=True)
+    os.makedirs(os.path.join(path, "stereo", "normal_maps"), exist_ok=True)
+
+
+def copy_maps(input_path: str, results_path: str, output_path: str):
+    shutil.copytree(os.path.join(input_path, "images"), os.path.join(output_path, "images"), dirs_exist_ok=True)
+    output_type = os.path.splitext(os.listdir(os.path.join(results_path, "depth_est"))[0])[1]
+    if output_type == ".bin":
+        shutil.copytree(os.path.join(results_path, "depth_est"), os.path.join(output_path, "stereo/depth_maps"),
+                        dirs_exist_ok=True)
+        shutil.copytree(os.path.join(results_path, "confidence"), os.path.join(output_path, "stereo/confidence_maps"),
+                        dirs_exist_ok=True)
+    else:
+        # If output is .pfm we need to read and convert to colmap .bin
+        for map_file in os.listdir(os.path.join(results_path, "depth_est")):
+            name, ext = os.path.splitext(map_file)
+            save_map(os.path.join(output_path, "stereo/depth_maps", name + ".bin"),
+                     read_map(os.path.join(results_path, "depth_est", map_file)))
+            save_map(os.path.join(output_path, "stereo/confidence_maps", name + ".bin"),
+                     read_map(os.path.join(results_path, "confidence", map_file)))
+
+
+def read_reconstruction(path: str) -> Tuple[List[Camera], List[Image], List[Tuple[int, List[int]]]]:
+    cameras = []
+    images = []
+    for cam_file in os.listdir(os.path.join(path, "cams")):
+        im_id = int(cam_file.split("_")[0])
+        im_file = cam_file.split("_")[0] + ".jpg"
+        image = PilImage.open(os.path.join(path, "images", im_file), "r")
+        intrinsics, extrinsics, _ = read_cam_file(os.path.join(path, "cams", cam_file))
+        cameras.append(Camera(im_id, "PINHOLE", image.width, image.height,
+                              [intrinsics[0, 0], intrinsics[1, 1], intrinsics[0, 2], intrinsics[1, 2]]))
+        qvec = rotation_matrix_to_quaternion(extrinsics[0:3, 0:3])
+        tvec = list(extrinsics[0:3, 3])
+        images.append(Image(im_id, qvec, tvec, im_id, im_file))
+
+    return cameras, images, read_pair_file(os.path.join(path, "pair.txt"))
+
+
+def write_patch_match_config(path: str, images: List[Image], pairs: List[Tuple[int, List[int]]]):
+    image_names: Dict[int, str] = {image.id: image.name for image in images}
+    with open(path, "w") as f:
+        for ref_id, src_ids in pairs:
+            f.write(image_names[ref_id] + "\n")
+            f.write(", ".join([image_names[src_id] for src_id in src_ids]) + "\n")
+
+
+def write_fusion_config(path: str, images: List[Image], pairs: List[Tuple[int, List[int]]]):
+    image_names: Dict[int, str] = {image.id: image.name for image in images}
+    with open(path, "w") as f:
+        f.writelines([",".join([image_names[view_id] for view_id in [pair[0]] + pair[1]]) + "\n" for pair in pairs])
+
+
+def write_sparse(path: str, cameras: List[Camera], images: List[Image]):
+    write_cameras(os.path.join(path, "cameras.txt"), cameras)
+    write_images(os.path.join(path, "images.txt"), images)
+    write_points_3d(os.path.join(path, "points3D.txt"))
+
+
+# Write cameras in colmap text format
+def write_cameras(path: str, cameras: List[Camera]):
+    cam_list = ["{} {} {} {} {} {} {} {}\n".format(c.id, c.model, c.width, c.height, c.params[0], c.params[1],
+                                                   c.params[2], c.params[3]) for c in cameras]
+    with open(path, "w") as f:
+        f.write("# Camera list with one line of data per camera:\n")
+        f.write("#   CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n")
+        f.write("# Number of cameras: {}\n".format(len(cameras)))
+        f.writelines(cam_list)
+
+
+# Write images in colmap text format
+def write_images(path: str, images: List[Image]):
+    image_list = ["{} {} {} {} {} {} {} {} {} {}\n\n".format(
+        i.id, i.qvec[0], i.qvec[1], i.qvec[2], i.qvec[3], i.tvec[0], i.tvec[1], i.tvec[2], i.camera_id, i.name
+    ) for i in images]
+
+    with open(path, "w") as f:
+        f.write("# Image list with two lines of data per image:\n")
+        f.write("#   IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n")
+        f.write("#   POINTS2D[] as (X, Y, POINT3D_ID)\n")
+        f.write("# Number of images: {}, mean observations per image: 0\n".format(len(image_list)))
+        f.writelines(image_list)
+
+
+# This is a dummy write of an empty points3D file since we do not have actual sparse 3D points to use
+def write_points_3d(path: str):
+    with open(path, "w") as f:
+        f.write("# 3D point list with one line of data per point:\n")
+        f.write("#   POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n")
+        f.write("# Number of points: 0, mean track length: 0")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert colmap results into input for PatchMatchNet")
+
+    parser.add_argument("--input_folder", type=str, help="Input PatchMatchNet reconstruction dir")
+    parser.add_argument("--results_folder", type=str, help="Input PatchMatchNet results dir")
+    parser.add_argument("--output_folder", type=str, default='', help="Output ColMap MVS workspace")
+
+    args = parser.parse_args()
+
+    if not args.results_folder:
+        args.results_folder = args.input_folder
+
+    if not args.output_folder:
+        args.output_folder = args.input_folder
+
+    create_output_dirs(args.output_folder)
+    copy_maps(args.input_folder, args.results_folder, args.output_folder)
+    cams, ims, im_pairs = read_reconstruction(args.input_folder)
+    write_patch_match_config(os.path.join(args.output_folder, "stereo/patch-match.cfg"), ims, im_pairs)
+    write_fusion_config(os.path.join(args.output_folder, "stereo/fusion.cfg"), ims, im_pairs)
+    write_sparse(os.path.join(args.output_folder, "sparse"), cams, ims)
diff --git a/convert_dtu_dataset.py b/convert_dtu_dataset.py
@@ -0,0 +1,80 @@
+import argparse
+import numpy as np
+import os
+import shutil
+from PIL import Image
+
+from datasets.data_io import read_image, read_map, save_image, save_map
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert DTU training dataset to standard input format")
+    parser.add_argument("--input_folder", type=str, help="Input training data")
+    parser.add_argument("--output_folder", type=str, help="Output converted training data")
+    parser.add_argument("--scan_list", type=str, help="Input scan list for conversion")
+    args = parser.parse_args()
+
+    if args.input_folder is None or not os.path.isdir(args.input_folder):
+        raise Exception("Invalid input folder")
+
+    if args.output_folder is None or not os.path.isdir(args.output_folder):
+        raise Exception("Invalid output folder")
+
+    if args.scan_list is None or not os.path.isfile(args.scan_list):
+        raise Exception("Invalid input scan list")
+
+    # Read input scan list and create output scan list
+    with open(args.scan_list) as f:
+        scans = [line.rstrip() for line in f.readlines()]
+
+    # Process each input scan and copy files to output folders
+    for scan in scans:
+        # Create output folders
+        scan_path = os.path.join(args.output_folder, scan)
+        cam_path = os.path.join(scan_path, "cams")
+        depth_path = os.path.join(scan_path, "depth_gt")
+        image_path = os.path.join(scan_path, "images")
+        mask_path = os.path.join(scan_path, "masks")
+        os.makedirs(scan_path, exist_ok=True)
+        os.makedirs(cam_path, exist_ok=True)
+        os.makedirs(depth_path, exist_ok=True)
+        os.makedirs(image_path, exist_ok=True)
+        os.makedirs(mask_path, exist_ok=True)
+
+        # Copy pair file
+        shutil.copy(os.path.join(args.input_folder, "Cameras_1/pair.txt"), os.path.join(scan_path, "pair.txt"))
+
+        for cam_file in os.listdir(os.path.join(args.input_folder, "Cameras_1/train")):
+            # Extract view ID and write cam file
+            view_id = int(cam_file.split("_")[0])
+
+            # Modify the cam file intrinsics by factor of 4 to match the training image size
+            with open(os.path.join(args.input_folder, "Cameras_1/train", cam_file)) as f:
+                lines = [line.rstrip() for line in f.readlines()]
+            tmp = np.fromstring(lines[7], dtype=np.float32, sep=" ") * 4.0
+            lines[7] = "{} {} {}".format(tmp[0], tmp[1], tmp[2])
+            tmp = np.fromstring(lines[8], dtype=np.float32, sep=" ") * 4.0
+            lines[8] = "{} {} {}".format(tmp[0], tmp[1], tmp[2])
+            with open(os.path.join(cam_path, cam_file), "w") as f:
+                for line in lines:
+                    f.write(line + "\n")
+
+            # Copy GT depth map after resizing and cropping
+            depth_map = read_map(os.path.join(
+                args.input_folder, "Depths_raw", scan, "depth_map_{:0>4}.pfm".format(view_id)), 800)
+            depth_map = depth_map[44:556, 80:720]
+            save_map(os.path.join(depth_path, "{:0>8}.pfm".format(view_id)), depth_map)
+
+            # Copy mask after resizing and cropping
+            mask = read_image(os.path.join(
+                args.input_folder, "Depths_raw", scan, "depth_visual_{:0>4}.png".format(view_id)), 800)[0]
+            mask = mask[44:556, 80:720] > 0.04
+            save_image(os.path.join(mask_path, "{:0>8}.png".format(view_id)), mask)
+
+            for light_idx in range(7):
+                # Copy images for each light index into separate sub-folders
+                image_prefix_path = os.path.join(image_path, str(light_idx))
+                os.makedirs(image_prefix_path, exist_ok=True)
+                image = Image.open(os.path.join(
+                    args.input_folder, "Rectified/{}_train/rect_{:0>3}_{}_r5000.png".format(
+                        scan, view_id + 1, light_idx)))
+                image.save(os.path.join(image_prefix_path, "{:0>8}.jpg".format(view_id)))
diff --git a/convert_eth3d_dataset.py b/convert_eth3d_dataset.py
@@ -0,0 +1,69 @@
+import argparse
+import numpy as np
+import os
+import shutil
+from datasets.data_io import read_image_dictionary
+
+from datasets.data_io import read_map, save_image
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert ETH 3D training dataset to standard input format")
+    parser.add_argument("--input_folder", type=str, help="Input training data")
+    parser.add_argument("--output_folder", type=str, help="Output converted training data")
+    parser.add_argument("--scan_list", type=str, help="Input scan list for conversion")
+    args = parser.parse_args()
+
+    if args.input_folder is None or not os.path.isdir(args.input_folder):
+        raise Exception("Invalid input folder")
+
+    if args.output_folder is None or not os.path.isdir(args.output_folder):
+        raise Exception("Invalid output folder")
+
+    if args.scan_list is None or not os.path.isfile(args.scan_list):
+        raise Exception("Invalid input scan list")
+
+    # Read input scan list and create output scan list
+    with open(args.scan_list) as f:
+        scans = [line.rstrip() for line in f.readlines()]
+
+    # Process each input scan and copy files to output folders
+    for scan in scans:
+        # Create output folders
+        scan_path = os.path.join(args.output_folder, scan)
+        cam_path = os.path.join(scan_path, "cams")
+        depth_path = os.path.join(scan_path, "depth_gt")
+        image_path = os.path.join(scan_path, "images")
+        mask_path = os.path.join(scan_path, "masks")
+        os.makedirs(scan_path, exist_ok=True)
+        os.makedirs(cam_path, exist_ok=True)
+        os.makedirs(depth_path, exist_ok=True)
+        os.makedirs(image_path, exist_ok=True)
+        os.makedirs(mask_path, exist_ok=True)
+
+        input_cam_path = os.path.join(args.input_folder, scan, "cams")
+        # Read image index file
+        image_index = read_image_dictionary(os.path.join(input_cam_path, "index2prefix.txt"))
+
+        # Copy pair file
+        shutil.copy(os.path.join(input_cam_path, "pair.txt"), os.path.join(scan_path, "pair.txt"))
+
+        for cam_file in os.listdir(input_cam_path):
+            if cam_file == "index2prefix.txt" or cam_file == "pair.txt":
+                continue
+
+            # Extract view ID and write cam file
+            view_id = int(cam_file.split("_")[0])
+            shutil.copy(os.path.join(input_cam_path, cam_file), os.path.join(cam_path, cam_file))
+
+            # Copy image
+            image_filename = os.path.join(args.input_folder, scan, "images", image_index[view_id])
+            shutil.copy(image_filename, os.path.join(image_path, "{:0>8}.png".format(view_id)))
+
+            # Copy GT depth map
+            depth_gt_filename = os.path.join(args.input_folder, scan, "depths", image_index[view_id])
+            depth_gt_filename = os.path.splitext(depth_gt_filename.replace('_undistorted', ''))[0] + '.pfm'
+            shutil.copy(depth_gt_filename, os.path.join(depth_path, "{:0>8}.pfm".format(view_id)))
+
+            # Create mask from GT depth map and save in output
+            mask = (read_map(depth_gt_filename) > 0.0).squeeze(2).astype(bool)
+            save_image(os.path.join(mask_path, "{:0>8}.png".format(view_id)), mask)