예제 #1
0
def draw_result(fp_label,
                fp_res,
                fp_massmap,
                fp_posres,
                setups,
                masses_gt=[1, 1, 1, 1],
                iters=[
                    -1,
                    0,
                    1,
                    2,
                ]):

    cmap = matplotlib.cm.coolwarm
    #cnorm = matplotlib.colors.Normalize(vmin=0.1,vmax=4.9)

    obj = os.path.split(fp_label)[1].split('_')[0]

    sim = SimBullet(gui=False)
    #sim._p.setGravity(0,0,-9.8)
    sim.addObject('table_rect', 'table_rect', [1, 1, 1, 0])
    sim.setupRendering(0.75, -90, -60, 500, 500)
    img_label = cv2.imread(fp_label, cv2.IMREAD_UNCHANGED)
    labels = np.unique(img_label)
    cells = CellPhysicsBase.createCells(img_label, masses_gt, 0.53 / 640.0)
    sim.addObject(obj, cells, [0, 0, 1, 1])
    sim.setObjectPose(obj, 0, 0, 0)

    cellinfos = sim.getCellBodyProperties(obj)
    vmax = max([info['mass'] for info in cellinfos])
    vmin = min([info['mass'] for info in cellinfos])
    cnorm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax)
    for info in cellinfos:
        info['color'] = cmap(cnorm(info['mass']))
    sim.setCellBodyProperties(obj, cellinfos)
    img = sim.draw()
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    cv2.imwrite(fp_massmap % -3, img)

    f = open(fp_res, 'rb')
    res = pickle.load(f)
    f.close()

    # estimation example
    cellinfos = res['history'][-1]['cellinfos']
    #cnorm = matplotlib.colors.Normalize(vmin=0.1,vmax=5)
    for info in cellinfos[obj]:
        info['color'] = cmap(cnorm((vmax - vmin) * 0.5))
    sim.setCellBodyProperties(obj, cellinfos[obj])
    img = sim.draw()

    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    cv2.imwrite(fp_massmap % -2, img)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    for it in iters:

        if it >= len(res['history']):
            continue

        cellinfos = res['history'][it]['cellinfos']
        #vmax = max([info['mass'] for info in cellinfos[obj]])
        #vmin = min([info['mass'] for info in cellinfos[obj]])
        #cnorm = matplotlib.colors.Normalize(vmin=vmin,vmax=vmax)
        for info in cellinfos[obj]:
            info['color'] = cmap(cnorm(info['mass']))
        sim.setCellBodyProperties(obj, cellinfos[obj])
        img = sim.draw()
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        cv2.imwrite(fp_massmap % it, img)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    return

    # debug image
    sim.setupRendering(0.75, 90, -90, 500, 500)

    sim.addObject('finger1', 'finger', [1, 1, 0.9, 1])
    sim.addObject('finger2', 'finger', [1, 1, 0.6, 1])
    sim.addObject('finger3', 'finger', [1, 1, 0.3, 1])

    sim.addObject('init', cells, [0.7, 0.7, 0.7, 1])
    pos, rot = sim.getPosRot(sim.objects['init']['shape'], 0, 0, 0)
    pos[2] = pos[2] - 0.002
    sim._p.resetBasePositionAndOrientation(sim.objects['init']['bodyId'], pos,
                                           rot)

    sim.addObject('gt', cells, [0, 1, 0, 1])
    for i in range(len(res['history'][-1]['test_poses'])):

        if 'real' not in setups:
            pixel2meter(cells, setups['test'][i])
            sim.resetObject('finger1',setups['test'][i]['finger']['pos'][0] - 0.03,\
                                      setups['test'][i]['finger']['pos'][1],0)
            sim.resetObject('finger2',setups['test'][i]['finger']['pos'][0] + setups['test'][i]['finger']['velocity'][0]*0.40 - 0.03,\
                                      setups['test'][i]['finger']['pos'][1] + setups['test'][i]['finger']['velocity'][1]*0.40,0)
            sim.resetObject('finger3',setups['test'][i]['finger']['pos'][0] + setups['test'][i]['finger']['velocity'][0]*0.90 - 0.03,\
                                      setups['test'][i]['finger']['pos'][1] + setups['test'][i]['finger']['velocity'][1]*0.90,0)
        else:
            sim.resetObject('finger1',setups['real'][i]['finger']['pos'][0] - 0.03,\
                                      setups['real'][i]['finger']['pos'][1],0)
            sim.resetObject('finger2',setups['real'][i]['finger']['pos'][0] + setups['real'][i]['finger']['velocity'][0]*0.40 - 0.03,\
                                      setups['real'][i]['finger']['pos'][1] + setups['real'][i]['finger']['velocity'][1]*0.40,0)
            sim.resetObject('finger3',setups['real'][i]['finger']['pos'][0] + setups['real'][i]['finger']['velocity'][0]*0.90 - 0.03,\
                                      setups['real'][i]['finger']['pos'][1] + setups['real'][i]['finger']['velocity'][1]*0.90,0)

        pose_gt = res['history'][-1]['test_poses'][i]['gt']

        #sim.resetObject(obj,100,100,0)
        sim.resetObject('gt', pose_gt[obj]['x'], pose_gt[obj]['y'],
                        pose_gt[obj]['yaw'])
        #img = sim.draw()
        #img = cv2.cvtColor(img,cv2.COLOR_RGB2BGR)
        #cv2.imwrite(fp_posres % (i+20),img)
        #img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)

        #sim.resetObject('gt',100,100,0)
        pose_est = res['history'][-1]['test_poses'][i]['estimation']
        sim.resetObject(obj, pose_est[obj]['x'], pose_est[obj]['y'],
                        pose_est[obj]['yaw'])

        pos, rot = sim.getPosRot(sim.objects['gt']['shape'], pose_gt[obj]['x'],
                                 pose_gt[obj]['y'], pose_gt[obj]['yaw'])
        pos[0] = pos[0] - 0.001
        pos[2] = pos[2] - 0.001
        sim._p.resetBasePositionAndOrientation(sim.objects['gt']['bodyId'],
                                               pos, rot)

        img = sim.draw()
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        cv2.imwrite(fp_posres % i, img)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            if key == 'pos':
                pos_push = [0, 0]
                pos_push[
                    0] = np.cos(yaw0) * value[0] - np.sin(yaw0) * value[1] + x0
                pos_push[
                    1] = np.sin(yaw0) * value[0] + np.cos(yaw0) * value[1] + y0
            elif key == 'velocity':
                vec_push = [0, 0]
                vec_push[0] = np.cos(yaw0) * value[0] - np.sin(yaw0) * value[1]
                vec_push[1] = np.sin(yaw0) * value[0] + np.cos(yaw0) * value[1]
            elif key == 'duration':
                duration = value
        vec_push[0] = vec_push[0] * duration
        vec_push[1] = vec_push[1] * duration

        pos_begin = Pose2D(x0, y0, yaw0)
        pos_final = Pose2D(x1, y1, yaw1)
        pos_push = Pose2D(pos_push[0], pos_push[1], 0)
        vec_push = Pose2D(vec_push[0], vec_push[1], 0)
        planner.addObservation(pos_begin, pos_final, pos_push, vec_push,
                               duration)

        cellinfos = planner.getProperties()
        cellinfos_sim = sim.getCellBodyProperties(obj)
        cnorm = matplotlib.colors.Normalize(vmin=0.1, vmax=5)
        for info, info_sim in zip(cellinfos, cellinfos_sim):
            info_sim['color'] = matplotlib.cm.coolwarm(cnorm(info['mass']))
        sim.setCellBodyProperties(obj, cellinfos_sim)

        x0, y0, yaw0 = (x1, y1, yaw1)