Basic loading/saving of unionfind, but not full solution

Author: Brook Roberts

opensfm/commands/create_tracks.py

@@ -19,6 +19,16 @@ def run(self, args):
         data = dataset.DataSet(args.dataset)
         images = data.images()
 
+        try:
+            graph = data.load_tracks_graph()
+            tracks, processed_images = matching.tracks_and_images(graph)
+        except IOError:
+            graph = None
+            tracks = None
+            processed_images = []
+
+        remaining_images = set(images) - set(processed_images)
+
         # Read local features
         logging.info('reading features')
         features = {}
@@ -30,7 +40,7 @@ def run(self, args):
 
         # Read matches
         matches = {}
-        for im1 in images:
+        for im1 in remaining_images:
             try:
                 im1_matches = data.load_matches(im1)
             except IOError:
@@ -39,7 +49,7 @@ def run(self, args):
                 matches[im1, im2] = im1_matches[im2]
 
         tracks_graph = matching.create_tracks_graph(features, colors, matches,
-                                                    data.config)
+                                                    data.config, data)
         data.save_tracks_graph(tracks_graph)
 
         end = time.time()

opensfm/dataset.py

@@ -3,7 +3,7 @@
 import os
 import json
 import errno
-import pickle
+import cPickle as pickle
 import gzip
 import numpy as np
 import networkx as nx
@@ -329,6 +329,19 @@ def find_matches(self, im1, im2):
                     return im2_matches[im1][:, [1, 0]]
         return []
 
+    def __unionfind_file(self, filename=None):
+        """Return path of unionfind file"""
+        return os.path.join(self.data_path, filename or 'unionfind.pkl')
+
+    def load_unionfind_file(self, filename=None):
+        """Return unionfind of tracks"""
+        with open(self.__unionfind_file(filename)) as fin:
+            return load_unionfind_file(fin)
+
+    def save_unionfind_file(self, unionfind, filename=None):
+        with open(self.__unionfind_file(filename), 'w') as fout:
+            save_unionfind_file(fout, unionfind)
+
     def __tracks_graph_file(self, filename=None):
         """Return path of tracks file"""
         return os.path.join(self.data_path, filename or 'tracks.csv')
@@ -501,3 +514,9 @@ def save_tracks_graph(fileobj, graph):
                 r, g, b = data['feature_color']
                 fileobj.write('%s\t%s\t%d\t%g\t%g\t%g\t%g\t%g\n' % (
                     str(image), str(track), fid, x, y, r, g, b))
+
+def load_unionfind_file(fileobj):
+    return pickle.load(fileobj)
+
+def save_unionfind_file(fileobj, unionfind):
+    pickle.dump(unionfind, fileobj, protocol=pickle.HIGHEST_PROTOCOL)

opensfm/matching.py

@@ -142,10 +142,14 @@ def good_track(track, min_length):
         return False
     return True
 
-
-def create_tracks_graph(features, colors, matches, config):
+def create_tracks_graph(features, colors, matches, config, data):
     logger.debug('Merging features onto tracks')
-    uf = UnionFind()
+
+    try:
+        uf = data.load_unionfind_file()
+    except IOError:
+        uf = UnionFind()
+
     for im1, im2 in matches:
         for f1, f2 in matches[im1, im2]:
             uf.union((im1, f1), (im2, f2))
@@ -158,6 +162,8 @@ def create_tracks_graph(features, colors, matches, config):
         else:
             sets[p] = [i]
 
+    data.save_unionfind_file(uf)
+
     tracks = [t for t in sets.values() if good_track(t, config.get('min_track_length', 2))]
     logger.debug('Good tracks: {}'.format(len(tracks)))