~wecacuee/mrol/mrol-dev

« back to all changes in this revision

Viewing changes to test/benchmarktest.py

  • Committer: Vikas Dhiman
  • Date: 2012-05-22 15:26:38 UTC
  • Revision ID: wecacuee@gmail.com-20120522152638-j6srlkgrpdq4kdfe
Added support for userdata in terms of colored voxels. Added benchmark test and corresponding data

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 mrol_mapping.mapper as mapper
 
6
import itertools as it
 
7
import mrol_mapping.depth_image as depth_image
 
8
import scipy.misc as scimisc
 
9
 
 
10
curdir = os.path.dirname(__file__)
 
11
 
 
12
def path(relpath):
 
13
    return os.path.join(curdir, relpath)
 
14
 
 
15
def slicegen(gen, n):
 
16
    """ Generator equivalent of gen[:n] """
 
17
    return it.imap(lambda x:x[1],
 
18
                   it.takewhile(lambda x:x[0] < n, enumerate(gen)))
 
19
 
 
20
def get_pointclouds(dir, depth, rgb, timestamp):
 
21
    for d, r, ts in it.izip(depth, rgb, timestamp):
 
22
        di = depth_image.load_image(os.path.join(dir, d[0]))
 
23
        ri = scimisc.imread(os.path.join(dir, r[0]))
 
24
        yield depth_image.image_to_points(di, ri,
 
25
                                          max_range=di.max(),
 
26
                                          timestamp=ts)
 
27
 
 
28
class TestMapper:
 
29
    def __init__(self, dir):
 
30
        depth  = bmtools_assc.read_file_list(os.path.join(dir, "depth.txt"))
 
31
        rgb    = bmtools_assc.read_file_list(os.path.join(dir, "rgb.txt"  ))
 
32
        mapped = bmtools_assc.associate(depth, rgb, 0, 0.2)
 
33
        self.pointclouds = get_pointclouds(dir,
 
34
                                           (depth[m[0]] for m in mapped),
 
35
                                           (rgb[m[1]] for m in mapped),
 
36
                                           (m[0] for m in mapped)
 
37
                                          )
 
38
 
 
39
    def _map_build(self, dir, scans):
 
40
        # TODO use a proper iterator to save memory rather than list of pointclouds?
 
41
 
 
42
        #pcs[0].pose = pcs[0].true_pose 
 
43
        # make first guess pose correct to fix trajectory in same coordinate 
 
44
        # system as true poses
 
45
        builder = mapper.MapBuilder(0.03, visualise=True, odometry_map=False)
 
46
        tf = open(os.path.join(dir, 'estimated_trajectory.txt'), 'w')
 
47
        pointclouds = (self.pointclouds if scans == -1 else
 
48
                       slicegen(self.pointclouds, scans))
 
49
        for pc in pointclouds:
 
50
            ((vdata, vres), bestpose) = builder.add_pointcloud(pc)
 
51
            tf.write(' '.join([str(pc.timestamp),]
 
52
                              + [str(p) for p in bestpose.getPosition()]
 
53
                              + [str(q) for q in bestpose.getQuaternion()]
 
54
                             ) + "\n")
 
55
        tf.close()
 
56
 
 
57
def benchmark_test(dir, scans):
 
58
    tm = TestMapper(dir)
 
59
    vmap = tm._map_build(dir, scans)
 
60
    return vmap
 
61
 
 
62
if __name__ == '__main__':
 
63
    parser = argparse.ArgumentParser(description='''
 
64
                                     test with benchmark data
 
65
    ''')
 
66
    parser.add_argument('--dir', 
 
67
                        default=path('data/rgbd_dataset_freiburg1_360/'),
 
68
                        help='''directory with data set. It should have
 
69
                        the following files: depth.txt, rgb.txt,
 
70
                        groundtruth.txt, depth/, rgb/''')
 
71
    parser.add_argument('--scans', 
 
72
                        default=-1,
 
73
                        help='''Number of scans to read for testing''')
 
74
    args = parser.parse_args()
 
75
    benchmark_test(args.dir, args.scans)