export_pmvs

#!/usr/bin/env python
import os.path, sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
import argparse

import cv2
from networkx.algorithms import bipartite
import numpy as np

from opensfm import dataset
from opensfm import features
from opensfm import io
from opensfm import matching


# Prepare OpenSfM output for dense reconstruction with PMVS

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='Convert output from OpenSfM to PMVS')
    parser.add_argument('dataset', help='path to the dataset to be processed')
    parser.add_argument('--output', help='output pmvs directory')
    parser.add_argument('--undistorted',
                        action='store_true',
                        help='export the undistorted reconstruction')
    args = parser.parse_args()

    data = dataset.DataSet(args.dataset)
    if args.output:
        base_output_path = args.output
    else:
        base_output_path = os.path.join(data.data_path, 'pmvs')

    print "Converting dataset [%s] to PMVS dir [%s]" % (
        data.data_path, base_output_path)

    io.mkdir_p(base_output_path)

    # load tracks for vis.dat
    try:
        if args.undistorted:
            graph = data.load_undistorted_tracks_graph()
        else:
            graph = data.load_tracks_graph()
        tracks, images = matching.tracks_and_images(graph)
        image_graph = bipartite.weighted_projected_graph(graph, images)
        use_vis_data = True
    except IOError:
        use_vis_data = False

    if args.undistorted:
        reconstructions = data.load_undistorted_reconstruction()
    else:
        reconstructions = data.load_reconstruction()

    for h, reconstruction in enumerate(reconstructions):
        print "Reconstruction", h
        output_path = os.path.join(base_output_path, "recon%d" % h)
        io.mkdir_p(output_path)
        io.mkdir_p(os.path.join(output_path, "visualize"))
        io.mkdir_p(os.path.join(output_path, "txt"))
        io.mkdir_p(os.path.join(output_path, "models"))

        shot_index = {image: i for i, image in enumerate(reconstruction.shots)}

        fvis = open(os.path.join(output_path, "vis.dat"), "w")
        fvis.write("VISDATA\n")
        fvis.write("%d\n" % len(shot_index))

        for image, i in shot_index.items():
            shot = reconstruction.shots[image]
            base = "%08d" % i
            print "Image:", image, base

            # vis.dat for this image
            if use_vis_data:
                adj_indices = []
                for adj_image in image_graph[image]:
                    weight = image_graph[image][adj_image]["weight"]
                    if weight > 0 and adj_image in shot_index:
                        adj_indices.append(shot_index[adj_image])

                num_covisible = len(adj_indices)
                fvis.write("%d " % i)
                fvis.write("%d " % num_covisible)
                for ai in adj_indices:
                    fvis.write("%d " % ai)
                fvis.write("\n")

            # radially undistort the original image
            camera = shot.camera
            if args.undistorted:
                undistorted_image = data.load_undistorted_image(image)
            else:
                original_image = data.load_image(image)[:, :, ::-1]
                original_h, original_w = original_image.shape[:2]
                K = camera.get_K_in_pixel_coordinates(original_w, original_h)
                distortion = np.array([camera.k1, camera.k2, 0, 0])
                undistorted_image = cv2.undistort(original_image, K, distortion)

            # resize and save the undistorted to visualize/%08d.jpg
            resized_image = features.resized_image(undistorted_image, data.config)
            new_image_path = os.path.join(output_path, "visualize", base + ".jpg")
            cv2.imwrite(new_image_path, resized_image)

            # write camera projection matrix to txt/%08d.txt
            resized_h, resized_w = resized_image.shape[:2]
            resized_K = camera.get_K_in_pixel_coordinates(resized_w, resized_h)
            P = resized_K.dot(shot.pose.get_Rt())

            new_txt = os.path.join(output_path, "txt", base + ".txt")
            with open(new_txt, "w") as f:
                f.write("CONTOUR\n")
                np.savetxt(f, P, '%f')

        fvis.close()

        # options.txt
        with open(os.path.join(output_path, "pmvs_options.txt"), "w") as f:
            f.write("level 1\n")
            f.write("csize 2\n")
            f.write("threshold 0.7\n")
            f.write("wsize 7\n")
            f.write("minImageNum 3\n")
            f.write("CPU 8\n")
            f.write("setEdge 0\n")
            f.write("useBound 0\n")
            f.write("useVisData {}\n".format(int(use_vis_data)))
            f.write("sequence -1\n")
            f.write("timages -1 0 %d\n" % len(shot_index))
            f.write("oimages 0\n")