GAAS

#!/usr/bin/env python

import os.path
import argparse
from itertools import combinations

import matplotlib.pyplot as pl
import networkx as nx
import numpy as np

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


def plot_matches(im1, im2, p1, p2):
    h1, w1, c = im1.shape
    h2, w2, c = im2.shape
    image = np.zeros((max(h1, h2), w1 + w2, 3), dtype=im1.dtype)
    image[0:h1, 0:w1, :] = im1
    image[0:h2, w1:(w1 + w2), :] = im2

    p1 = features.denormalized_image_coordinates(p1, w1, h1)
    p2 = features.denormalized_image_coordinates(p2, w2, h2)
    pl.imshow(image)
    for a, b in zip(p1, p2):
        pl.plot([a[0], b[0] + w1], [a[1], b[1]], 'c')

    pl.plot(p1[:, 0], p1[:, 1], 'ob')
    pl.plot(p2[:, 0] + w1, p2[:, 1], 'ob')


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Plot matches between images')
    parser.add_argument('dataset',
                        help='path to the dataset to be processed')
    parser.add_argument('--image',
                        help='show tracks for a specific')
    parser.add_argument('--images',
                        help='show tracks between a subset of images (separated by commas)')
    parser.add_argument('--graph',
                        help='display image graph',
                        action='store_true')
    parser.add_argument('--save_figs',
                        help='save figures instead of showing them',
                        action='store_true')
    args = parser.parse_args()

    data = dataset.DataSet(args.dataset)
    images = data.images()

    if args.graph:
        # Plot graph
        graph = nx.Graph()

        for im1 in images:
            for im2 in images:
                if im1 != im2:
                    matches = data.find_matches(im1, im2)
                    graph.add_edge(im1, im2, weight=len(matches))

        weights = [i[2]['weight'] for i in graph.edges(data=True)]

        nx.draw_graphviz(
            graph, edge_color=weights, edge_cmap=pl.get_cmap('Blues'),
            edge_vmin=0, edge_vmax=200)
        pl.axis('off')
        pl.show()
    else:
        # Plot matches between images
        if args.image:
            pairs = [(args.image, o) for o in images if o != args.image]
        elif args.images:
            subset = args.images.split(',')
            pairs = combinations(subset, 2)
        else:
            pairs = combinations(images, 2)

        i = 0
        for im1, im2 in pairs:
            matches = data.find_matches(im1, im2)
            if len(matches) == 0:
                continue
            print('plotting {} matches between {} {}'.format(len(matches), im1, im2))

            p1, f1, c1 = data.load_features(im1)
            p2, f2, c2 = data.load_features(im2)
            p1 = p1[matches[:, 0]]
            p2 = p2[matches[:, 1]]

            pl.figure(figsize=(20, 10))
            pl.title('Images: ' + im1 + ' - ' + im2 + ', matches: ' +
                     str(matches.shape[0]))
            plot_matches(data.load_image(im1),
                         data.load_image(im2), p1, p2)
            i += 1
            if args.save_figs:
                p = os.path.join(args.dataset, 'plot_tracks')
                io.mkdir_p(p)
                pl.savefig(os.path.join(p, "{}_{}.jpg".format(im1, im2)), dpi=100)
                pl.close()
            else:
                if i >= 10:
                    i = 0
                    pl.show()

        if not args.save_figs and i > 0:
            pl.show()