def extract_clouds(demofile):
	seg_num = 0
	leaf_size = 0.045
	keys = {}
	pclouds = []
	for demo_name in demofile:
		if demo_name != "ar_demo":
			for seg_name in demofile[demo_name]:
				if seg_name != 'done':
					keys[seg_num] = (demo_name, seg_name)
					pclouds.append(clouds.downsample(np.asarray(demofile[demo_name][seg_name]["cloud_xyz"]), leaf_size))
					print demo_name, seg_name
					seg_num += 1
	return keys, pclouds
示例#2
0
def track_video_frames_offline(video_dir, T_w_k, init_tfm, ref_image, rope_params = None):
    
    video_stamps = np.loadtxt(osp.join(video_dir,demo_names.stamps_name))    
    
    rgbs = []
    depths = []
    color_clouds = []
    
    for (index, stamp) in zip(range(len(video_stamps)), video_stamps):        
        #print index, stamp
        rgb = cv2.imread(osp.join(video_dir,demo_names.rgb_name%index))
        rgb = color_match.match(ref_image, rgb)
        assert rgb is not None
        depth = cv2.imread(osp.join(video_dir,demo_names.depth_name%index), 2)
        assert depth is not None
        color_cloud = extract_rope_tracking(rgb, depth, T_w_k)
        color_cloud = clouds.downsample_colored(color_cloud, .01)
        color_cloud[:,:3] = color_cloud[:,:3].dot(init_tfm[:3,:3].T) + init_tfm[:3,3][None,:]

        if index == 0:
            rope_xyz = extract_red(rgb, depth, T_w_k)
            rope_xyz = clouds.downsample(rope_xyz, .01)
            rope_xyz = rope_xyz.dot(init_tfm[:3, :3].T) + init_tfm[:3,3][None,:]
            rope_nodes = rope_initialization.find_path_through_point_cloud(rope_xyz)
            rope_radius = 0.005
                
        rgbs.append(rgb)
        depths.append(depth)
        color_clouds.append(color_cloud)
        
    rgbs = np.asarray(rgbs)
    depths = np.asarray(depths)    
    
    all_tracked_nodes = cbulletracpy2.py_tracking(rope_nodes, rope_radius, T_w_k, color_clouds, rgbs, depths, 5)
    
    for i in range(len(rgbs)):
        if i % 30 != 0:
            continue
        print i
        fig.clf()
        ax = fig.gca(projection='3d')
        ax.set_autoscale_on(False)
        tracked_nodes = all_tracked_nodes[i, :, :]
        ax.plot(tracked_nodes[:, 0], tracked_nodes[:,1], tracked_nodes[:,2], 'o')
        fig.show()
        raw_input()
        
    return all_tracked_nodes
def extract_segs(demofile):
	seg_num = 0
	leaf_size = 0.045
	keys = {}
	segs = []
	for demo_name in demofile:
		if demo_name != "ar_demo":
			for seg_name in demofile[demo_name]:
				if seg_name != 'done':
					keys[seg_num] = (demo_name, seg_name)
					seg = demofile[demo_name][seg_name]
					pc = clouds.downsample(np.asarray(seg['cloud_xyz']), leaf_size)
					segs.append((pc, np.asarray(seg['l']['tfms_s'])[:,0:3,3], np.asarray(seg['r']['tfms_s'])[:,0:3,3]))
					print demo_name, seg_name
					seg_num += 1
	return keys, segs
def grabcut(rgb, depth, T_w_k):
    xyz_k = clouds.depth_to_xyz(depth, asus_xtion_pro_f)
    xyz_w = xyz_k.dot(T_w_k[:3,:3].T) + T_w_k[:3,3][None,None,:]

    valid_mask = depth > 0

    from hd_rapprentice import interactive_roi as ir
    xys = ir.get_polyline(rgb, "rgb")
    xy_corner1 = np.clip(np.array(xys).min(axis=0), [0,0], [639,479])
    xy_corner2 = np.clip(np.array(xys).max(axis=0), [0,0], [639,479])
    polymask = ir.mask_from_poly(xys)
    #cv2.imshow("mask",mask)
        
    xy_tl = np.array([xy_corner1, xy_corner2]).min(axis=0)
    xy_br = np.array([xy_corner1, xy_corner2]).max(axis=0)

    xl, yl = xy_tl
    w, h = xy_br - xy_tl
    mask = np.zeros((h,w),dtype='uint8')    
    mask[polymask[yl:yl+h, xl:xl+w] > 0] = cv2.GC_PR_FGD
    print mask.shape
    #mask[h//4:3*h//4, w//4:3*w//4] = cv2.GC_PR_FGD

    tmp1 = np.zeros((1, 13 * 5))
    tmp2 = np.zeros((1, 13 * 5))    
    cv2.grabCut(rgb[yl:yl+h, xl:xl+w, :],mask,(0,0,0,0),tmp1, tmp2,10,mode=cv2.GC_INIT_WITH_MASK)

    mask = mask % 2
    #mask = ndi.binary_erosion(mask, utils_images.disk(args.erode)).astype('uint8')
    contours = cv2.findContours(mask,cv2.RETR_LIST,cv2.CHAIN_APPROX_NONE)[0]
    cv2.drawContours(rgb[yl:yl+h, xl:xl+w, :],contours,-1,(0,255,0),thickness=2)
    
    cv2.imshow('rgb', rgb)
    print "press enter to continue"
    cv2.waitKey()

    zsel = xyz_w[yl:yl+h, xl:xl+w, 2]
    mask = (mask%2==1) & np.isfinite(zsel)# & (zsel - table_height > -1)
    mask &= valid_mask[yl:yl+h, xl:xl+w]
    
    xyz_sel = xyz_w[yl:yl+h, xl:xl+w,:][mask.astype('bool')]
    return clouds.downsample(xyz_sel, .01)
def extract_segs(demofile, num_segs=None):
    seg_num = 0
    leaf_size = 0.045
    keys = {}
    segs = []
    for demo_name in demofile:
        if demo_name != "ar_demo":
            for seg_name in demofile[demo_name]:
                if seg_name != 'done':
                    keys[seg_num] = (demo_name, seg_name)
                    seg = demofile[demo_name][seg_name]
                    pc = clouds.downsample(np.asarray(seg['cloud_xyz']), leaf_size)
                    pc, params = registration.unit_boxify(pc)
                    ltraj = np.asarray(seg['l']['tfms_s'])[:,0:3,3]
                    ltraj = lerp(np.linspace(0,ltraj.shape[0],traj_n), range(ltraj.shape[0]), ltraj)
                    rtraj = np.asarray(seg['r']['tfms_s'])[:,0:3,3]
                    rtraj = lerp(np.linspace(0,rtraj.shape[0],traj_n), range(rtraj.shape[0]), rtraj)
                    
                    segs.append(((pc, params), ltraj, rtraj))
                    seg_num += 1
                    if num_segs is not None and seg_num >= num_segs: return keys, segs
    
    return keys, segs
def track_video_frames_online(video_dir, T_w_k, init_tfm, table_plane, ref_image = None, rope_params = None):    
    video_stamps = np.loadtxt(osp.join(video_dir,demo_names.stamps_name))
    
    env = openravepy.Environment()
    
    stdev = np.array([])
    
    
    
    rgb = cv2.imread(osp.join(video_dir,demo_names.rgb_name%0))
    rgb = color_match.match(ref_image, rgb)
    assert rgb is not None
    depth = cv2.imread(osp.join(video_dir,demo_names.depth_name%0), 2)
    assert depth is not None
    rope_xyz = extract_red(rgb, depth, T_w_k)

    table_dir = np.array([table_plane[0], table_plane[1], table_plane[2]])
    table_dir = table_dir / np.linalg.norm(table_dir)
    
    table_dir = init_tfm[:3,:3].dot(table_dir)    
    if np.dot([0,0,1], table_dir) < 0:
        table_dir = -table_dir
    table_axis = np.cross(table_dir, [0,0,1])
    table_angle = np.arccos(np.dot([0,0,1], table_dir))
    table_tfm = openravepy.matrixFromAxisAngle(table_axis, table_angle)
    table_center_xyz = np.mean(rope_xyz, axis=0)
        
    table_tfm[:3,3] = - table_tfm[:3,:3].dot(table_center_xyz) + table_center_xyz
    init_tfm = table_tfm.dot(init_tfm)
    
    tracked_nodes = None
    
    for (index, stamp) in zip(range(len(video_stamps)), video_stamps):
        print index, stamp
        rgb = cv2.imread(osp.join(video_dir,demo_names.rgb_name%index))
        rgb = color_match.match(ref_image, rgb)
        assert rgb is not None
        depth = cv2.imread(osp.join(video_dir,demo_names.depth_name%index), 2)
        assert depth is not None
        color_cloud = extract_rope_tracking(rgb, depth, T_w_k) # color_cloud is at first in camera frame
        
        
        #print "cloud in camera frame"
        #print "==================================="
        #print color_cloud[:, :3]

        
    
        color_cloud = clouds.downsample_colored(color_cloud, .01)
        color_cloud[:,:3] = color_cloud[:,:3].dot(init_tfm[:3,:3].T) + init_tfm[:3,3][None,:] # color_cloud now is in global frame
        
        raw_color_cloud = np.array(color_cloud)
        ##########################################
        ### remove the shadow points on the desk
        ##########################################
        color_cloud = remove_shadow(color_cloud)
        
        if tracked_nodes is not None:
            color_cloud = rope_shape_filter(tracked_nodes, color_cloud)
             
        if index == 0:
            rope_xyz = extract_red(rgb, depth, T_w_k)
            rope_xyz = clouds.downsample(rope_xyz, .01)
            rope_xyz = rope_xyz.dot(init_tfm[:3, :3].T) + init_tfm[:3,3][None,:]
            rope_nodes = rope_initialization.find_path_through_point_cloud(rope_xyz) # rope_nodes and rope_xyz are in global frame
                        
            # print rope_nodes
            
            table_height = rope_xyz[:,2].min() - .02
            table_xml = make_table_xml(translation=[1, 0, table_height], extents=[.85, .55, .01])
            env.LoadData(table_xml)            
            
            bulletsimpy.sim_params.scale = 10
            bulletsimpy.sim_params.maxSubSteps = 200
            if rope_params is None:
                rope_params = bulletsimpy.CapsuleRopeParams()
                rope_params.radius = 0.005
                #angStiffness: a rope with a higher angular stiffness seems to have more resistance to bending.
                #orig self.rope_params.angStiffness = .1
                rope_params.angStiffness = .1
                #A higher angular damping causes the ropes joints to change angle slower.
                #This can cause the rope to be dragged at an angle by the arm in the air, instead of falling straight.
                #orig self.rope_params.angDamping = 1
                rope_params.angDamping = 1
                #orig self.rope_params.linDamping = .75
                #Not sure what linear damping is, but it seems to limit the linear accelertion of centers of masses.
                rope_params.linDamping = .75
                #Angular limit seems to be the minimum angle at which the rope joints can bend.
                #A higher angular limit increases the minimum radius of curvature of the rope.
                rope_params.angLimit = .4
                #TODO--Find out what the linStopErp is
                #This could be the tolerance for error when the joint is at or near the joint limit
                rope_params.linStopErp = .2  
            
            bt_env = bulletsimpy.BulletEnvironment(env, [])
            bt_env.SetGravity([0,0,-0.1])
            rope = bulletsimpy.CapsuleRope(bt_env, 'rope', rope_nodes, rope_params)
            
            continue
        
        
        
 #==============================================================================
 #       rope_nodes = rope.GetNodes()
 #       #print "rope nodes in camera frame"
 #       R = init_tfm[:3,:3].T
 #       t = - R.dot(init_tfm[:3,3])
 #       rope_in_camera_frame = rope_nodes.dot(R.T) + t[None,:]
 #       #print rope_in_camera_frame
 #       uvs = XYZ_to_xy(rope_in_camera_frame[:,0], rope_in_camera_frame[:,1], rope_in_camera_frame[:,2], asus_xtion_pro_f)
 #       uvs = np.vstack(uvs)
 #       #print uvs
 #       #print "uvs"
 #       uvs = uvs.astype(int)
 #       
 #       n_rope_nodes = len(rope_nodes)
 #       
 #       DEPTH_OCCLUSION_DIST = .03
 #       occ_dist = DEPTH_OCCLUSION_DIST
 #       
 #       vis = np.ones(n_rope_nodes)
 #       
 #       rope_depth_in_camera = np.array(rope_in_camera_frame)
 #       depth_xyz = depth_to_xyz(depth, asus_xtion_pro_f)
 # 
 #       for i in range(n_rope_nodes):
 #           u = uvs[0, i]
 #           v = uvs[1, i]
 #           
 #           neighbor_radius = 10;
 #           v_range = [max(0, v-neighbor_radius), v+neighbor_radius+1]
 #           u_range = [max(0, u-neighbor_radius), u+neighbor_radius+1]
 #           
 #           xyzs = depth_xyz[v_range[0]:v_range[1], u_range[0]:u_range[1]]
 #                       
 #           xyzs = np.reshape(xyzs, (xyzs.shape[0]*xyzs.shape[1], xyzs.shape[2]))
 #           dists_to_origin = np.linalg.norm(xyzs, axis=1)
 #           
 #           dists_to_origin = dists_to_origin[np.isfinite(dists_to_origin)]
 #           
 #           #print dists_to_origin
 #           min_dist_to_origin = np.min(dists_to_origin)
 #                       
 #           print v, u, min_dist_to_origin, np.linalg.norm(rope_in_camera_frame[i])
 #                                   
 #           if min_dist_to_origin + occ_dist > np.linalg.norm(rope_in_camera_frame[i]):
 #               vis[i] = 1
 #           else:
 #               vis[i] = 0
 #               
 #       #print "vis result"
 #       #print vis
 #       
 #       rope_depth_in_global = rope_depth_in_camera.dot(init_tfm[:3,:3].T) + init_tfm[:3,3][None,:]
 #==============================================================================
        
        
        depth_cloud = extract_cloud(depth, T_w_k)
        depth_cloud[:,:3] = depth_cloud[:,:3].dot(init_tfm[:3,:3].T) + init_tfm[:3,3][None,:] # depth_cloud now is in global frame


        [tracked_nodes, new_stdev] = bulletsimpy.py_tracking(rope, bt_env, init_tfm, color_cloud, rgb, depth, 5, stdev)
        stdev = new_stdev
        
        #print tracked_nodes

        #if index % 10 != 0:
        #    continue
        
        print index
        
        
        xx, yy = np.mgrid[-1:3, -1:3]
        zz = np.ones(xx.shape) * table_height
        table_cloud = [xx, yy, zz]
        
        fig.clf()
        ax = fig.gca(projection='3d')
        ax.set_autoscale_on(False)     
        
        print init_tfm[:,3]
        ax.plot(depth_cloud[:,0], depth_cloud[:,1], depth_cloud[:,2], 'go', alpha=0.1)
        ax.plot(color_cloud[:,0], color_cloud[:,1], color_cloud[:,2]-0.1, 'go')
        ax.plot(raw_color_cloud[:,0], raw_color_cloud[:,1], raw_color_cloud[:,2] -0.2, 'ro')
        #ax.plot(rope_depth_in_global[:,0], rope_depth_in_global[:,1], rope_depth_in_global[:,2], 'ro')
        ax.plot_surface(table_cloud[0], table_cloud[1], table_cloud[2], color = (0,1,0,0.5))
        ax.plot(tracked_nodes[:,0], tracked_nodes[:,1], tracked_nodes[:,2], 'bo')
        ax.plot([init_tfm[0,3]], [init_tfm[1,3]], [init_tfm[2,3]], 'ro')   


        fig.show()
        raw_input()
    annotation_file = osp.join(demo_dir,"ann.yaml")
    with open(annotation_file, "r") as fh: annotations = yaml.load(fh)
    
    look_stamps = [seg_info['look'] for seg_info in annotations]
    rgb_imgs, depth_imgs= get_video_frames(rgbd_dir, look_stamps)
    
    if demo_name in perturb_demofile.keys():
        demo_group = perturb_demofile[demo_name]
    else:
        demo_group = perturb_demofile.create_group(demo_name)
        
    
    n_perturb_existed = len(demo_group.keys()) # number of perturbations
    
    object_xyz = cloud_proc_func(rgb_imgs[0], depth_imgs[0], np.eye(4)) 
    object_xyz = clouds.downsample(object_xyz, .01)

    hitch_xyz = None
    hitch_pos = None
    if args.has_hitch:
        hitch_normal = clouds.clouds_plane(object_xyz)
        hitch_xyz, hitch_pos = hitch_proc_func(rgb_imgs[0], depth_imgs[0], np.eye(4), hitch_normal)
        hitch_xyz = clouds.downsample(hitch_xyz, .01)
        xyz = np.r_[object_xyz, hitch_xyz]
    else:
        xyz = object_xyz
    
    mlab.figure(0)
    mlab.clf()
    mlab.points3d(xyz[:,0], xyz[:,1], xyz[:,2], color=(1,0,0), scale_factor=.005)
    
                try:
                    print "1",  seg_info.keys()
                    print "2", del_fields
                    if field in seg_info: del seg_info[field]
                except Exception as e:
                    import IPython
                    IPython.embed()

            if args.has_hitch:
                if not hitch_found:
                    hitch_normal = clouds.clouds_plane(cloud)
                    hitch, hitch_pos = hitch_proc_func(rgb_image, np.asarray(seg_info["depth"]), np.eye(4), hitch_normal)
                    hitch_found = True
                seg_info["full_hitch"] = hitch
                seg_info["full_object"] = cloud
                seg_info["hitch"] = clouds.downsample(hitch, .01)
                seg_info["object"] = clouds.downsample(cloud, .01)
                seg_info["hitch_pos"] = hitch_pos
                seg_info["cloud_xyz"] = np.r_[seg_info["hitch"], seg_info["object"]]
            else:
                seg_info["full_cloud_xyz"] = cloud
                seg_info["cloud_xyz"] = clouds.downsample(cloud, .01)

            seg_info["cloud_proc_func"] = args.cloud_proc_func
            seg_info["cloud_proc_mod"] = args.cloud_proc_mod
            seg_info["cloud_proc_code"] = inspect.getsource(cloud_proc_func)
            
            if table_cloud is None:
                table_cloud, table_plane = cloud_proc_mod.extract_table_ransac(rgb_image, np.asarray(seg_info["depth"]), np.eye(4))