masses = [] frics = [] forces = [] for info in cellinfo: masses.append(info['mass']) frics.append(info['fric']) forces.append(info['mass'] * info['fric']) cnormF = matplotlib.colors.Normalize(vmin=min(forces), vmax=max(forces)) #cnormF = matplotlib.colors.Normalize(vmin=0.1,vmax=5) for info in cellinfo: info['color'] = matplotlib.cm.coolwarm( cnormF(info['mass'] * info['fric'])) sim.setCellBodyProperties(obj, cellinfo) img = sim.draw() img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) cv2.imwrite( os.path.join(dir_sav + '/fig/' + filename + '_force.png'), img) #cnormM = matplotlib.colors.Normalize(vmin=min(masses),vmax=max(masses)) cnormM = matplotlib.colors.Normalize(vmin=1, vmax=4) for info in cellinfo: info['color'] = matplotlib.cm.coolwarm(cnormM(info['mass'])) sim.setCellBodyProperties(obj, cellinfo) img = sim.draw() img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) cv2.imwrite( os.path.join(dir_sav + '/fig/' + filename + '_mass.png'), img) #cnormf = matplotlib.colors.Normalize(vmin=min(frics),vmax=max(frics))
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)
sim.addObject('finger0', 'finger', [1, 1, 0, 1]) sim.addObject('finger1', 'finger', [1, 0, 1, 1]) sim.addObject('table_rect', 'table_rect', [1, 1, 1, 1]) img_label = cv2.imread(fp_label, cv2.IMREAD_UNCHANGED) cells = CellPhysicsBase.createCells( img_label, [0, 0, 0, 0], meter_per_pixel=0.53 / 640.0) sim.addObject(obj + '0', cells, [0, 0, 1, 1]) sim.addObject(obj + '1', cells, [1, 0, 0, 1]) sim.setObjectPose('finger0', push_x, push_y, 0) sim.setObjectPose('finger1', push_x + float(node['push']['vec']['x']), push_y, 0) sim.setObjectPose(obj + '0', 0, 0, 0) sim.setObjectPose(obj + '1', x_res, y_res, yaw_res) draw = sim.draw() plt.imshow(draw) plt.show() if True: engine = get_engine(args.method) engine.setup(setup, n_group=-1) n_group = len(engine.meta_targets[obj]['cells']['cells']) history = engine.infer(infer_type='cell_mass', nsimslimit=100) print( '[Test Result] cell err: %.3f (m) pose err: %.3f (m) %.3f (deg)' % (history[-1]['test_error_cell'], history[-1]['test_error_xy'], history[-1]['test_error_yaw']))