~mrol-dev/mrol/trunk

« back to all changes in this revision

Viewing changes to test/benchmarktest.py

  • Committer: Vikas Dhiman
  • Date: 2012-05-14 16:28:10 UTC
  • Revision ID: wecacuee@gmail.com-20120514162810-gqqevqvisr4qyi5r
Backtracing must be working fine now.

Show diffs side-by-side

added added

removed removed

Lines of Context:
1
 
import argparse
2
 
import tempfile
3
 
import rgbd_benchmark_tools.associate as bmtools_assc
4
 
import os
5
 
import re
6
 
import mrol_mapping.mapper as mapper
7
 
import itertools as it
8
 
import mrol_mapping.depth_image as depth_image
9
 
import scipy.misc as scimisc
10
 
 
11
 
curdir = os.path.dirname(__file__)
12
 
 
13
 
def path(relpath):
14
 
    return os.path.join(curdir, relpath)
15
 
 
16
 
def slicegen(gen, n):
17
 
    """ Generator equivalent of gen[:n] """
18
 
    return it.imap(lambda x:x[1],
19
 
                   it.takewhile(lambda x:x[0] < n, enumerate(gen)))
20
 
 
21
 
def get_pointclouds(dir, depth, rgb, timestamp):
22
 
    for d, r, ts in it.izip(depth, rgb, timestamp):
23
 
        di = depth_image.load_image(os.path.join(dir, d[0]))
24
 
        ri = scimisc.imread(os.path.join(dir, r[0]))
25
 
        yield depth_image.image_to_points(di, ri,
26
 
                                          max_range=di.max(),
27
 
                                          timestamp=ts)
28
 
_backup_pattern = re.compile("\.(\d+)$")
29
 
def isbackupfile(fname):
30
 
    match = _backup_pattern.search(fname)
31
 
    return match
32
 
def backup_file(fname):
33
 
    isbak = isbackupfile(fname)
34
 
    next = 1
35
 
    if isbak:
36
 
        # assume that fname exists
37
 
        next = int(isbak.group(1)) + 1
38
 
        nextfname = _backup_pattern.subs(".%d"%next, fname)
39
 
    else:
40
 
        nextfname = fname + ".%d"%next
41
 
    if os.path.exists(nextfname):
42
 
        backup_file(nextfname)
43
 
    os.rename(fname, nextfname)
44
 
    return True
45
 
 
46
 
class TestMapper:
47
 
    def __init__(self, dir):
48
 
        depth  = bmtools_assc.read_file_list(os.path.join(dir, "depth.txt"))
49
 
        rgb    = bmtools_assc.read_file_list(os.path.join(dir, "rgb.txt"  ))
50
 
        mapped = bmtools_assc.associate(depth, rgb, 0, 0.2)
51
 
        self.pointclouds = get_pointclouds(dir,
52
 
                                           (depth[m[0]] for m in mapped),
53
 
                                           (rgb[m[1]] for m in mapped),
54
 
                                           (m[0] for m in mapped)
55
 
                                          )
56
 
 
57
 
    def _map_build(self, dir, scans):
58
 
        # TODO use a proper iterator to save memory rather than list of pointclouds?
59
 
 
60
 
        #pcs[0].pose = pcs[0].true_pose 
61
 
        # make first guess pose correct to fix trajectory in same coordinate 
62
 
        # system as true poses
63
 
        builder = mapper.MapBuilder(0.03, visualise=True, odometry_map=False)
64
 
        est_traj = os.path.join(dir, 'estimated_trajectory.txt')
65
 
        if os.path.exists(est_traj) and not backup_file(est_traj):
66
 
            raise RuntimeError("Unable to backup file:{0}".format(est_traj))
67
 
        tf = open(est_traj, 'w')
68
 
        pointclouds = (self.pointclouds if scans == -1 else
69
 
                       slicegen(self.pointclouds, scans))
70
 
        for pc in pointclouds:
71
 
            ((vdata, vres), bestpose) = builder.add_pointcloud(pc)
72
 
            tf.write(' '.join([str(pc.timestamp),]
73
 
                              + [str(p) for p in bestpose.getPosition()]
74
 
                              + [str(q) for q in bestpose.getQuaternion()]
75
 
                             ) + "\n")
76
 
        tf.close()
77
 
 
78
 
def benchmark_test(dir, scans):
79
 
    tm = TestMapper(dir)
80
 
    vmap = tm._map_build(dir, scans)
81
 
    return vmap
82
 
 
83
 
if __name__ == '__main__':
84
 
    parser = argparse.ArgumentParser(description='''
85
 
                                     test with benchmark data
86
 
    ''')
87
 
    parser.add_argument('--dir', 
88
 
                        default=path('data/rgbd_dataset_freiburg1_360/'),
89
 
                        help='''directory with data set. It should have
90
 
                        the following files: depth.txt, rgb.txt,
91
 
                        groundtruth.txt, depth/, rgb/''')
92
 
    parser.add_argument('--scans', 
93
 
                        default=-1,
94
 
                        help='''Number of scans to read for testing''')
95
 
    args = parser.parse_args()
96
 
    benchmark_test(args.dir, args.scans)