def _plot_sample_light(fig, ax, ib, count, data_dict): plt.cla() ax.set_xlim(_X_LIM[0], _X_LIM[1]) ax.set_ylim(_Y_LIM[0], _Y_LIM[1]) ax.set_xlabel("x [m]") ax.set_ylabel("y [m]") ax.set_aspect("equal") # ax.set_title(f"Frame {data_dict['idx'][ib]}. Press any key to exit.") # scan and cls label scan_r = data_dict["scans"][ib][-1] scan_phi = data_dict["scan_phi"][ib] scan_x, scan_y = u.rphi_to_xy(scan_r, scan_phi) ax.scatter(scan_x, scan_y, s=0.5, c="blue") # annotation ann = data_dict["dets_wp"][ib] ann_valid_mask = data_dict["anns_valid_mask"][ib] if len(ann) > 0: ann = np.array(ann) det_x, det_y = u.rphi_to_xy(ann[:, 0], ann[:, 1]) for x, y, valid in zip(det_x, det_y, ann_valid_mask): if valid: # c = plt.Circle((x, y), radius=0.1, color="red", fill=True) c = plt.Circle((x, y), radius=0.4, color="red", fill=False) ax.add_artist(c)
def _plot_sample(fig, ax, ib, count, data_dict): plt.cla() ax.set_xlim(_X_LIM[0], _X_LIM[1]) ax.set_ylim(_Y_LIM[0], _Y_LIM[1]) ax.set_xlabel("x [m]") ax.set_ylabel("y [m]") ax.set_aspect("equal") ax.set_title(f"Frame {data_dict['idx'][ib]}. Press any key to exit.") # scan and cls label scan_r = data_dict["scans"][ib][-1] scan_phi = data_dict["scan_phi"][ib] scan_x, scan_y = u.rphi_to_xy(scan_r, scan_phi) target_cls = data_dict["target_cls"][ib] ax.scatter(scan_x[target_cls == -2], scan_y[target_cls == -2], s=1, c="yellow") ax.scatter(scan_x[target_cls == -1], scan_y[target_cls == -1], s=1, c="orange") ax.scatter(scan_x[target_cls == 0], scan_y[target_cls == 0], s=1, c="black") ax.scatter(scan_x[target_cls > 0], scan_y[target_cls > 0], s=1, c="green") # annotation ann = data_dict["dets_wp"][ib] ann_valid_mask = data_dict["anns_valid_mask"][ib] if len(ann) > 0: ann = np.array(ann) det_x, det_y = u.rphi_to_xy(ann[:, 0], ann[:, 1]) for x, y, valid in zip(det_x, det_y, ann_valid_mask): c = "blue" if valid else "orange" c = plt.Circle((x, y), radius=0.4, color=c, fill=False) ax.add_artist(c) # reg label target_reg = data_dict["target_reg"][ib] dets_r, dets_phi = u.canonical_to_global(scan_r, scan_phi, target_reg[:, 0], target_reg[:, 1]) dets_r = dets_r[target_cls > 0] dets_phi = dets_phi[target_cls > 0] dets_x, dets_y = u.rphi_to_xy(dets_r, dets_phi) ax.scatter(dets_x, dets_y, s=10, c="red")
def _plot_annotation(ann, ax, color, radius): if len(ann) > 0: ann = np.array(ann) det_x, det_y = u.rphi_to_xy(ann[:, 0], ann[:, 1]) for x, y in zip(det_x, det_y): c = plt.Circle((x, y), radius=radius, color=color, fill=False) ax.add_artist(c)
def plot_one_frame( batch_dict, frame_idx, pred_cls=None, pred_reg=None, dets_cls=None, dets_xy=None, xlim=_X_LIM, ylim=_Y_LIM, ): """Plot one frame from a batch, specified by frame_idx. Returns: fig: figure handle ax: axis handle """ fig, ax = _create_figure("", xlim, ylim) # scan and cls label scan_r = batch_dict["scans"][frame_idx][-1] scan_phi = batch_dict["scan_phi"][frame_idx] target_cls = batch_dict["target_cls"][frame_idx] _plot_scan(ax, scan_r, scan_phi, target_cls, s=1) # annotation ann = batch_dict["dets_wp"][frame_idx] ann_valid_mask = batch_dict["anns_valid_mask"][frame_idx] if len(ann) > 0: ann = np.array(ann) det_x, det_y = u.rphi_to_xy(ann[:, 0], ann[:, 1]) for x, y, valid in zip(det_x, det_y, ann_valid_mask): c = "blue" if valid else "orange" c = plt.Circle((x, y), radius=0.4, color=c, fill=False) ax.add_artist(c) # regression target target_reg = batch_dict["target_reg"][frame_idx] _plot_target(ax, target_reg, target_cls > 0, scan_r, scan_phi, s=10, c="blue") # regression result if dets_xy is not None and dets_cls is not None: _plot_detection(ax, dets_cls, dets_xy, s=40, color_dim=1) if pred_cls is not None and pred_reg is not None: _plot_prediction(ax, pred_cls, pred_reg, scan_r, scan_phi, s=2, color_dim=1) return fig, ax
def _plot_scan(ax, scan_r, scan_phi, target_cls, s): scan_x, scan_y = u.rphi_to_xy(scan_r, scan_phi) ax.scatter(scan_x[target_cls < 0], scan_y[target_cls < 0], s=s, c="orange") ax.scatter(scan_x[target_cls == 0], scan_y[target_cls == 0], s=s, c="black") ax.scatter(scan_x[target_cls > 0], scan_y[target_cls > 0], s=s, c="green")
def _plot_target(ax, target_reg, target_flag, scan_r, scan_phi, s, c): dets_r, dets_phi = u.canonical_to_global(scan_r, scan_phi, target_reg[:, 0], target_reg[:, 1]) dets_r = dets_r[target_flag] dets_phi = dets_phi[target_flag] dets_x, dets_y = u.rphi_to_xy(dets_r, dets_phi) ax.scatter(dets_x, dets_y, s=s, c=c)
def _plot_sequence(): drow_handle = DROWHandle( split="train", cfg={ "num_scans": 1, "scan_stride": 1, "data_dir": "./data/DROWv2-data" }, ) fig = plt.figure(figsize=(10, 10)) ax = fig.add_subplot(111) _break = False if _INTERACTIVE: def p(event): nonlocal _break _break = True fig.canvas.mpl_connect("key_press_event", p) else: if os.path.exists(_SAVE_DIR): shutil.rmtree(_SAVE_DIR) os.makedirs(_SAVE_DIR) for i, data_dict in enumerate(drow_handle): if _break: break scan_x, scan_y = u.rphi_to_xy(data_dict["scans"][-1], data_dict["scan_phi"]) plt.cla() ax.set_aspect("equal") ax.set_xlim(_X_LIM[0], _X_LIM[1]) ax.set_ylim(_Y_LIM[0], _Y_LIM[1]) ax.set_xlabel("x [m]") ax.set_ylabel("y [m]") ax.set_title(f"Frame {data_dict['idx']}. Press any key to exit.") # ax.axis("off") ax.scatter(scan_x, scan_y, s=1, c="black") _plot_annotation(data_dict["dets_wc"], ax, "red", 0.6) _plot_annotation(data_dict["dets_wa"], ax, "green", 0.4) _plot_annotation(data_dict["dets_wp"], ax, "blue", 0.35) if _INTERACTIVE: plt.pause(0.1) else: plt.savefig( os.path.join(_SAVE_DIR, f"frame_{data_dict['idx']:04}.png")) if _INTERACTIVE: plt.show()
def _plot_annotation_detr(ax, anns, radius, color, linestyle="-"): if len(anns) == 0: return det_x, det_y = u.rphi_to_xy(anns[0], anns[1]) for x, y in zip(det_x, det_y): c = plt.Circle((x, y), radius=radius, color=color, fill=False, linestyle=linestyle) ax.add_artist(c)
def _model_eval_fn(model, batch_dict): _, tb_dict, rtn_dict = _model_fn(model, batch_dict) pred_cls = torch.sigmoid(rtn_dict["pred_cls"]).data.cpu().numpy() pred_reg = rtn_dict["pred_reg"].data.cpu().numpy() # # DEBUG use perfect predictions # pred_cls = batch_dict["target_cls"] # pred_cls[pred_cls < 0] = 1 # pred_reg = batch_dict["target_reg"] fig_dict = {} file_dict = {} # postprocess network prediction to get detection scans = batch_dict["scans"] scan_phi = batch_dict["scan_phi"] for ib in range(len(scans)): # store detection, which will be used by _model_eval_collate_fn to compute AP dets_xy, dets_cls, _ = u.nms_predicted_center(scans[ib][-1], scan_phi[ib], pred_cls[ib], pred_reg[ib]) frame_id = "%06d" % batch_dict['frame_id'][ib] sequence = batch_dict["sequence"][ib] # save detection results for evaluation det_str = pru.drow_detection_to_kitti_string(dets_xy, dets_cls, None) file_dict["detections/" + str(sequence) + "/" + str(frame_id)] = det_str # save corresponding groundtruth for evaluation anns_rphi = batch_dict["dets_wp"][ib] if len(anns_rphi) > 0: anns_rphi = np.array(anns_rphi, dtype=np.float32) gts_xy = np.stack(u.rphi_to_xy(anns_rphi[:, 0], anns_rphi[:, 1]), axis=1) gts_occluded = np.logical_not( batch_dict["anns_valid_mask"][ib]).astype(np.int) gts_str = pru.drow_detection_to_kitti_string( gts_xy, None, gts_occluded) file_dict["groundtruth/{}/{}".format(sequence, frame_id)] = gts_str else: file_dict["groundtruth/{}/{}".format(sequence, frame_id)] = "" # TODO When to plot if _PLOTTING: fig, ax = plot_one_frame(batch_dict, ib, pred_cls[ib], pred_reg[ib], dets_cls, dets_xy) fig_dict["figs/{}/{}".format(sequence, frame_id)] = (fig, ax) return tb_dict, file_dict, fig_dict
def _plot_sequence(): jrdb_handle = JRDBHandle( split="train", cfg={"data_dir": "./data/JRDB", "num_scans": 10, "scan_stride": 1}, ) fig = plt.figure(figsize=(20, 10)) gs = GridSpec(3, 2, figure=fig) ax_im = fig.add_subplot(gs[0, :]) ax_bev = fig.add_subplot(gs[1:, 1]) ax_fpv_xz = fig.add_subplot(gs[1, 0]) ax_fpv_yz = fig.add_subplot(gs[2, 0]) color_pool = np.random.uniform(size=(100, 3)) _break = False if _INTERACTIVE: def p(event): nonlocal _break _break = True fig.canvas.mpl_connect("key_press_event", p) else: if os.path.exists(_SAVE_DIR): shutil.rmtree(_SAVE_DIR) os.makedirs(_SAVE_DIR) for i, data_dict in enumerate(jrdb_handle): if _break: break # lidar pc_xyz_upper = jt.transform_pts_upper_velodyne_to_base( data_dict["pc_data"]["upper_velodyne"] ) pc_xyz_lower = jt.transform_pts_lower_velodyne_to_base( data_dict["pc_data"]["lower_velodyne"] ) # labels boxes, label_ids = [], [] for ann in data_dict["pc_anns"]: jrdb_handle.box_is_on_ground(ann) boxes.append(jt.box_from_jrdb(ann["box"])) label_ids.append(int(ann["label_id"].split(":")[-1]) % len(color_pool)) # laser laser_r = data_dict["laser_data"][-1] laser_phi = data_dict["laser_grid"] laser_z = data_dict["laser_z"] laser_x, laser_y = u.rphi_to_xy(laser_r, laser_phi) pc_xyz_laser = jt.transform_pts_laser_to_base( np.stack((laser_x, laser_y, laser_z), axis=0) ) # BEV ax_bev.cla() ax_bev.set_aspect("equal") ax_bev.set_xlim(_XY_LIM[0], _XY_LIM[1]) ax_bev.set_ylim(_XY_LIM[0], _XY_LIM[1]) ax_bev.set_title(f"Frame {data_dict['idx']}. Press any key to exit.") ax_bev.set_xlabel("x [m]") ax_bev.set_ylabel("y [m]") # ax_bev.axis("off") for rgb_dim, pc_xyz in zip( (2, 1, 0), (pc_xyz_upper, pc_xyz_lower, pc_xyz_laser) ): ax_bev.scatter(pc_xyz[0], pc_xyz[1], s=1, c=_get_pts_color(pc_xyz, rgb_dim)) for label_id, box in zip(label_ids, boxes): box.draw_bev(ax_bev, c=color_pool[label_id]) # side view for dim, ax_fpv in zip((0, 1), (ax_fpv_xz, ax_fpv_yz)): ax_fpv.cla() ax_fpv.set_aspect("equal") ax_fpv.set_xlim(_XY_LIM[0], _XY_LIM[1]) ax_fpv.set_ylim(_Z_LIM[0], _Z_LIM[1]) ax_fpv.set_title(f"Frame {data_dict['idx']}. Press any key to exit.") ax_fpv.set_xlabel("x [m]" if dim == 0 else "y [m]") ax_fpv.set_ylabel("z [m]") # ax_fpv.axis("off") for rgb_dim, pc_xyz in zip( (2, 1, 0), (pc_xyz_upper, pc_xyz_lower, pc_xyz_laser) ): ax_fpv.scatter( pc_xyz[dim], pc_xyz[2], s=1, c=_get_pts_color(pc_xyz, rgb_dim) ) for label_id, box in zip(label_ids, boxes): box.draw_fpv(ax_fpv, dim=dim, c=color_pool[label_id]) # image ax_im.cla() ax_im.axis("off") ax_im.imshow(data_dict["im_data"]["stitched_image0"]) # detection bounding box for box_dict in data_dict["im_dets"]: x0, y0, w, h = box_dict["box"] verts = np.array( [(x0, y0), (x0, y0 + h), (x0 + w, y0 + h), (x0 + w, y0), (x0, y0)] ) c = max(float(box_dict["score"]) - 0.5, 0) * 2.0 ax_im.plot(verts[:, 0], verts[:, 1], c=(1.0 - c, 1.0 - c, 1.0)) # laser points on image p_xy, ib_mask = jt.transform_pts_base_to_stitched_im(pc_xyz_laser) ax_im.scatter( p_xy[0, ib_mask], p_xy[1, ib_mask], s=1, c=_get_pts_color(pc_xyz_laser[:, ib_mask], dim=0), ) if _INTERACTIVE: plt.pause(0.1) else: plt.savefig(os.path.join(_SAVE_DIR, f"frame_{data_dict['idx']:04}.png")) if _INTERACTIVE: plt.show()
def generate_pseudo_labels(): with open("./base_dr_spaam_jrdb_cfg.yaml", "r") as f: cfg = yaml.safe_load(f) cfg["dataset"]["pseudo_label"] = True cfg["dataset"]["pl_correction_level"] = 0 test_loader = get_dataloader( split=_SPLIT, batch_size=1, num_workers=1, shuffle=False, dataset_cfg=cfg["dataset"], ) if os.path.exists(_SAVE_DIR): shutil.rmtree(_SAVE_DIR) time.sleep(1.0) os.makedirs(_SAVE_DIR) sequences_tp_fp_tn_fn = {} # for computing true positive and negative rate # generate pseudo labels for all sample for count, batch_dict in enumerate(tqdm(test_loader)): if count >= _MAX_COUNT: break for ib in range(len(batch_dict["input"])): frame_id = f"{batch_dict['frame_id'][ib]:06d}" sequence = batch_dict["sequence"][ib] if count % _PLOTTING_INTERVAL == 0: # visualize the whole frame _plot_frame(batch_dict, ib) # visualize each pseudo labels _plot_pseudo_labels(batch_dict, ib) # save pseudo labels as detection results for evaluation pl_xy = batch_dict["pseudo_label_loc_xy"][ib] pl_str = ( pru.drow_detection_to_kitti_string(pl_xy, None, None) if len(pl_xy) > 0 else "" ) pl_file = os.path.join(_SAVE_DIR, f"detections/{sequence}/{frame_id}.txt") _write_file_make_dir(pl_file, pl_str) # save groundtruth anns_rphi = batch_dict["dets_wp"][ib] if len(anns_rphi) > 0: anns_rphi = np.array(anns_rphi, dtype=np.float32) gts_xy = np.stack( u.rphi_to_xy(anns_rphi[:, 0], anns_rphi[:, 1]), axis=1 ) gts_occluded = np.logical_not(batch_dict["anns_valid_mask"][ib]).astype( np.int ) gts_str = pru.drow_detection_to_kitti_string(gts_xy, None, gts_occluded) else: gts_str = "" gts_file = os.path.join(_SAVE_DIR, f"groundtruth/{sequence}/{frame_id}.txt") _write_file_make_dir(gts_file, gts_str) # compute true positive and negative rate target_cls = batch_dict["target_cls"][ib] target_cls_gt = batch_dict["target_cls_real"][ib] tn = np.sum( np.logical_or( np.logical_and(target_cls == 0, target_cls_gt == 0), np.logical_and(target_cls == 0, target_cls_gt == -1), ) ) fn = np.sum(target_cls == 0) - tn tp = np.sum( np.logical_or( np.logical_and(target_cls == 1, target_cls_gt == 1), np.logical_and(target_cls == 1, target_cls_gt == -1), ) ) fp = np.sum(target_cls == 1) - tp if sequence in sequences_tp_fp_tn_fn.keys(): tp0, fp0, tn0, fn0 = sequences_tp_fp_tn_fn[sequence] sequences_tp_fp_tn_fn[sequence] = ( tp + tp0, fp + fp0, tn + tn0, fn + fn0, ) else: sequences_tp_fp_tn_fn[sequence] = (tp, fp, tn, fn) # write sequence statistics to file for sequence, (tp, fp, tn, fn) in sequences_tp_fp_tn_fn.items(): st_file = os.path.join(_SAVE_DIR, f"evaluation/{sequence}/tp_fp_tn_fn.txt") _write_file_make_dir(st_file, f"{tp},{fp},{tn},{fn}")
def _plot_prediction(ax, pred_cls, pred_reg, scan_r, scan_phi, s, color_dim): pred_r, pred_phi = u.canonical_to_global(scan_r, scan_phi, pred_reg[:, 0], pred_reg[:, 1]) pred_x, pred_y = u.rphi_to_xy(pred_r, pred_phi) pred_color = _cls_to_color(pred_cls, color_dim=color_dim) ax.scatter(pred_x, pred_y, s=s, c=pred_color)
def _get_regression_target(scan_rphi, dets_rphi, person_radius_small, person_radius_large, min_close_points): """Generate classification and regression label. Args: scan_rphi (np.ndarray[2, N]): Scan points in polar coordinate dets_rphi (np.ndarray[2, M]): Annotated person centers in polar coordinate person_radius_small (float): Points less than this distance away from an annotation is assigned to that annotation and marked as fg. person_radius_large (float): Points with no annotation smaller than this distance is marked as bg. min_close_points (int): Annotations with supportive points fewer than this value is marked as invalid. Supportive points are those within the small radius. Returns: target_cls (np.ndarray[N]): Classification label, 1=fg, 0=bg, -1=ignore target_reg (np.ndarray[N, 2]): Regression label anns_valid_mask (np.ndarray[M]) """ N = scan_rphi.shape[1] # no annotation in this frame if len(dets_rphi) == 0: return np.zeros(N, dtype=np.int64), np.zeros((N, 2), dtype=np.float32), [] scan_xy = np.stack(u.rphi_to_xy(scan_rphi[0], scan_rphi[1]), axis=0) dets_xy = np.stack(u.rphi_to_xy(dets_rphi[0], dets_rphi[1]), axis=0) dist_scan_dets = np.hypot( scan_xy[0].reshape(1, -1) - dets_xy[0].reshape(-1, 1), scan_xy[1].reshape(1, -1) - dets_xy[1].reshape(-1, 1), ) # (M, N) pairwise distance between scan and detections # mark out annotations that has too few scan points anns_valid_mask = ( np.sum(dist_scan_dets < person_radius_small, axis=1) > min_close_points ) # (M, ) # for each point, find the distance to its closest annotation argmin_dist_scan_dets = np.argmin(dist_scan_dets, axis=0) # (N, ) min_dist_scan_dets = dist_scan_dets[argmin_dist_scan_dets, np.arange(N)] # points within small radius, whose corresponding annotation is valid, is marked # as foreground target_cls = -1 * np.ones(N, dtype=np.int64) fg_mask = np.logical_and(anns_valid_mask[argmin_dist_scan_dets], min_dist_scan_dets < person_radius_small) target_cls[fg_mask] = 1 target_cls[min_dist_scan_dets > person_radius_large] = 0 # regression target dets_matched_rphi = dets_rphi[:, argmin_dist_scan_dets] target_reg = np.stack( u.global_to_canonical(scan_rphi[0], scan_rphi[1], dets_matched_rphi[0], dets_matched_rphi[1]), axis=1, ) return target_cls, target_reg, anns_valid_mask
def play_sequence(seq_name, tracking=False): # scans scans_data = np.genfromtxt(seq_name, delimiter=',') scans_t = scans_data[:, 1] scans = scans_data[:, 2:] scan_phi = u.get_laser_phi() # odometry, used only for plotting odo_name = seq_name[:-3] + 'odom2' odos = np.genfromtxt(odo_name, delimiter=',') odos_t = odos[:, 1] odos_phi = odos[:, 4] # detector detector = DrSpaamDetector(num_pts=450, ang_inc_degree=0.5, tracking=tracking, gpu=True, ckpt=default_ckpts) # detector.set_laser_spec(angle_inc=np.radians(0.5), num_pts=450) # scanner location rad_tmp = 0.5 * np.ones(len(scan_phi), dtype=np.float) xy_scanner = u.rphi_to_xy(rad_tmp, scan_phi) xy_scanner = np.stack(xy_scanner[::-1], axis=1) # plot fig = plt.figure(figsize=(10, 10)) ax = fig.add_subplot(111) _break = False _pause = False def p(event): nonlocal _break, _pause if event.key == 'escape': _break = True if event.key == ' ': _pause = not _pause fig.canvas.mpl_connect('key_press_event', p) # video sequence odo_idx = 0 for i in range(len(scans)): # for i in range(0, len(scans), 20): plt.cla() ax.set_aspect('equal') ax.set_xlim(-15, 15) ax.set_ylim(-15, 15) # ax.set_title('Frame: %s' % i) ax.set_title('Press escape key to exit.') ax.axis("off") # find matching odometry while odo_idx < len(odos_t) - 1 and odos_t[odo_idx] < scans_t[i]: odo_idx += 1 odo_phi = odos_phi[odo_idx] odo_rot = np.array([[np.cos(odo_phi), -np.sin(odo_phi)], [np.sin(odo_phi), np.cos(odo_phi)]], dtype=np.float32) # plot scanner location xy_scanner_rot = np.matmul(xy_scanner, odo_rot.T) ax.plot(xy_scanner_rot[:, 0], xy_scanner_rot[:, 1], c='black') ax.plot((0, xy_scanner_rot[0, 0] * 1.0), (0, xy_scanner_rot[0, 1] * 1.0), c='black') ax.plot((0, xy_scanner_rot[-1, 0] * 1.0), (0, xy_scanner_rot[-1, 1] * 1.0), c='black') # plot points scan = scans[i] scan_y, scan_x = u.rphi_to_xy(scan, scan_phi + odo_phi) ax.scatter(scan_x, scan_y, s=1, c='blue') # inference dets_xy, dets_cls = detector.detect(scan) # plot detection dets_xy_rot = np.matmul(dets_xy, odo_rot.T) cls_thresh = 0.3 for j in range(len(dets_xy)): if dets_cls[j] < cls_thresh: continue # c = plt.Circle(dets_xy_rot[j], radius=0.5, color='r', fill=False) c = plt.Circle(dets_xy_rot[j], radius=0.5, color='r', fill=False, linewidth=2) ax.add_artist(c) # plot track if tracking: cls_thresh = 0.2 tracks, tracks_cls = detector.get_tracklets() for t, tc in zip(tracks, tracks_cls): if tc >= cls_thresh and len(t) > 1: t_rot = np.matmul(t, odo_rot.T) ax.plot(t_rot[:, 0], t_rot[:, 1], color='g', linewidth=2) # plt.savefig('/home/dan/tmp/det_img/frame_%04d.png' % i) plt.pause(0.001) if _break: break if _pause: plt.pause(1)
def _get_sample(self, idx): data_dict = self.__handle[idx] # DROW defines laser frame as x-forward, y-right, z-downward # JRDB defines laser frame as x-forward, y-left, z-upward # Use DROW frame for DR-SPAAM or DROW3 # equivalent of flipping y axis (inversing laser phi angle) data_dict["laser_data"] = data_dict["laser_data"][:, ::-1] scan_rphi = np.stack( (data_dict["laser_data"][-1], data_dict["laser_grid"]), axis=0) # get annotation in laser frame ann_xyz = [(ann["box"]["cx"], ann["box"]["cy"], ann["box"]["cz"]) for ann in data_dict["pc_anns"]] if len(ann_xyz) > 0: ann_xyz = np.array(ann_xyz, dtype=np.float32).T ann_xyz = jt.transform_pts_base_to_laser(ann_xyz) ann_xyz[1] = -ann_xyz[1] # to DROW frame dets_rphi = np.stack(u.xy_to_rphi(ann_xyz[0], ann_xyz[1]), axis=0) else: dets_rphi = [] # regression target target_cls, target_reg, anns_valid_mask = _get_regression_target( scan_rphi, dets_rphi, person_radius_small=0.4, person_radius_large=0.8, min_close_points=5, ) data_dict["target_cls"] = target_cls data_dict["target_reg"] = target_reg data_dict["anns_valid_mask"] = anns_valid_mask # regression target from pseudo labels if self._pseudo_label: # get pixels of laser points projected on image scan_x, scan_y = u.rphi_to_xy(scan_rphi[0], scan_rphi[1]) scan_y = -scan_y # convert DROW frame to JRDB laser frame scan_xyz_laser = np.stack((scan_x, scan_y, data_dict["laser_z"]), axis=0) scan_pixel_xy, _ = jt.transform_pts_laser_to_stitched_im( scan_xyz_laser) # get detection boxes boxes = [] boxes_confs = [] for box_dict in data_dict["im_dets"]: x0, y0, w, h = box_dict["box"] boxes.append((x0, y0, x0 + w, y0 + h)) boxes_confs.append(box_dict["score"]) ######## # NOTE for ablation, using 2D annotation to generate pseudo labels # for box_dict in data_dict["im_anns"]: # x0, y0, w, h = box_dict["box"] # boxes.append((x0, y0, x0 + w, y0 + h)) # boxes_confs.append(1.0) ######## boxes = np.array(boxes, dtype=np.float32) boxes_confs = np.array(boxes_confs, dtype=np.float32) # pseudo label pl_xy, pl_boxes, pl_neg_mask = u.generate_pseudo_labels( scan_rphi[0], scan_rphi[1], scan_pixel_xy, boxes, boxes_confs) ( target_cls_pseudo, target_reg_pseudo, ) = _get_regression_target_from_pseudo_labels( scan_rphi, pl_xy, pl_neg_mask, person_radius_small=0.4, person_radius_large=0.8, min_close_points=5, pl_correction_level=self._pl_correction_level, target_cls_annotated=data_dict["target_cls"], target_reg_annotated=data_dict["target_reg"], ) data_dict["pseudo_label_loc_xy"] = pl_xy data_dict["pseudo_label_boxes"] = pl_boxes # still keep the original target for debugging purpose data_dict["target_cls_real"] = data_dict["target_cls"] data_dict["target_reg_real"] = data_dict["target_reg"] data_dict["target_cls"] = target_cls_pseudo data_dict["target_reg"] = target_reg_pseudo # to be consistent with DROWDataset in order to use the same evaluation function dets_wp = [] for i in range(dets_rphi.shape[1]): dets_wp.append((dets_rphi[0, i], dets_rphi[1, i])) data_dict["dets_wp"] = dets_wp data_dict["scans"] = data_dict["laser_data"] data_dict["scan_phi"] = data_dict["laser_grid"] if self._augment_data: data_dict = u.data_augmentation(data_dict) data_dict["input"] = u.scans_to_cutout( data_dict["laser_data"], data_dict["laser_grid"], stride=1, **self._cutout_kwargs, ) return data_dict
def _test_detr_dataloader(): with open("./tests/test.yaml", "r") as f: cfg = yaml.safe_load(f) cfg["dataset"]["DataHandle"]["tracking"] = True cfg["dataset"]["DataHandle"]["num_scans"] = 1 test_loader = get_dataloader( split="train", batch_size=8, num_workers=1, shuffle=True, dataset_cfg=cfg["dataset"], ) fig = plt.figure(figsize=(10, 10)) ax = fig.add_subplot(111) _break = False if _INTERACTIVE: def p(event): nonlocal _break _break = True fig.canvas.mpl_connect("key_press_event", p) else: if os.path.exists(_SAVE_DIR): shutil.rmtree(_SAVE_DIR) os.makedirs(_SAVE_DIR) for count, data_dict in enumerate(test_loader): for ib in range(len(data_dict["input"])): fr_idx = data_dict["frame_dict_curr"][ib]["idx"] plt.cla() ax.set_xlim(_X_LIM[0], _X_LIM[1]) ax.set_ylim(_Y_LIM[0], _Y_LIM[1]) ax.set_xlabel("x [m]") ax.set_ylabel("y [m]") ax.set_aspect("equal") ax.set_title(f"Frame {fr_idx}. Press any key to exit.") # scan and cls label scan_r = data_dict["frame_dict_curr"][ib]["laser_data"][-1] scan_phi = data_dict["frame_dict_curr"][ib]["laser_grid"] scan_x, scan_y = u.rphi_to_xy(scan_r, scan_phi) target_cls = data_dict["target_cls"][ib] ax.scatter(scan_x[target_cls < 0], scan_y[target_cls < 0], s=1, c="orange") ax.scatter(scan_x[target_cls == 0], scan_y[target_cls == 0], s=1, c="black") ax.scatter(scan_x[target_cls > 0], scan_y[target_cls > 0], s=1, c="green") # annotation for tracking anns_tracking = data_dict["frame_dict_curr"][ib]["dets_rphi_prev"] anns_tracking_mask = data_dict["anns_tracking_mask"][ib] anns_tracking = anns_tracking[:, anns_tracking_mask] if len(anns_tracking) > 0: det_x, det_y = u.rphi_to_xy(anns_tracking[0], anns_tracking[1]) for x, y in zip(det_x, det_y): c = plt.Circle((x, y), radius=0.5, color="gray", fill=False, linestyle="--") ax.add_artist(c) # annotation anns = data_dict["frame_dict_curr"][ib]["dets_rphi"] anns_valid_mask = data_dict["anns_valid_mask"][ib] if len(anns) > 0: det_x, det_y = u.rphi_to_xy(anns[0], anns[1]) for x, y, valid in zip(det_x, det_y, anns_valid_mask): c = "blue" if valid else "orange" c = plt.Circle((x, y), radius=0.4, color=c, fill=False) ax.add_artist(c) # reg label for previous frame target_reg_prev = data_dict["target_reg_prev"][ib] target_tracking_flag = data_dict["target_tracking_flag"][ib] dets_r_prev, dets_phi_prev = u.canonical_to_global( scan_r, scan_phi, target_reg_prev[:, 0], target_reg_prev[:, 1]) dets_r_prev = dets_r_prev[target_tracking_flag] dets_phi_prev = dets_phi_prev[target_tracking_flag] dets_x_prev, dets_y_prev = u.rphi_to_xy(dets_r_prev, dets_phi_prev) ax.scatter(dets_x_prev, dets_y_prev, s=25, c="gray") # reg label for current frame target_reg = data_dict["target_reg"][ib] dets_r, dets_phi = u.canonical_to_global(scan_r, scan_phi, target_reg[:, 0], target_reg[:, 1]) dets_r = dets_r[target_cls > 0] dets_phi = dets_phi[target_cls > 0] dets_x, dets_y = u.rphi_to_xy(dets_r, dets_phi) ax.scatter(dets_x, dets_y, s=10, c="red") if _INTERACTIVE: plt.pause(0.1) else: plt.savefig( os.path.join( _SAVE_DIR, f"b{count:03}s{ib:02}f{fr_idx:04}.png", )) if _INTERACTIVE: plt.show()
def test_detector(): data_handle = JRDBHandle( split="train", cfg={ "data_dir": "./data/JRDB", "num_scans": 10, "scan_stride": 1 }, ) # ckpt_file = "/home/jia/ckpts/ckpt_jrdb_ann_drow3_e40.pth" # d = Detector( # ckpt_file, model="DROW3", gpu=True, stride=1, panoramic_scan=True # ) ckpt_file = "/home/jia/ckpts/ckpt_jrdb_ann_dr_spaam_e20.pth" d = Detector(ckpt_file, model="DR-SPAAM", gpu=True, stride=1, panoramic_scan=True) d.set_laser_fov(360) fig = plt.figure(figsize=(10, 10)) ax = fig.add_subplot(111) _break = False if _INTERACTIVE: def p(event): nonlocal _break _break = True fig.canvas.mpl_connect("key_press_event", p) else: if os.path.exists(_SAVE_DIR): shutil.rmtree(_SAVE_DIR) os.makedirs(_SAVE_DIR) for i, data_dict in enumerate(data_handle): if _break: break # plot scans scan_r = data_dict["laser_data"][-1, ::-1] # to DROW frame scan_x, scan_y = u.rphi_to_xy(scan_r, data_dict["laser_grid"]) plt.cla() ax.set_aspect("equal") ax.set_xlim(_X_LIM[0], _X_LIM[1]) ax.set_ylim(_Y_LIM[0], _Y_LIM[1]) ax.set_xlabel("x [m]") ax.set_ylabel("y [m]") ax.set_title(f"Frame {data_dict['idx']}. Press any key to exit.") # ax.axis("off") ax.scatter(scan_x, scan_y, s=1, c="black") # plot annotation ann_xyz = [(ann["box"]["cx"], ann["box"]["cy"], ann["box"]["cz"]) for ann in data_dict["pc_anns"]] if len(ann_xyz) > 0: ann_xyz = np.array(ann_xyz, dtype=np.float32).T ann_xyz = jt.transform_pts_base_to_laser(ann_xyz) ann_xyz[1] = -ann_xyz[1] # to DROW frame for xyz in ann_xyz.T: c = plt.Circle( (xyz[0], xyz[1]), radius=0.4, color="red", fill=False, linestyle="--", ) ax.add_artist(c) # plot detection dets_xy, dets_cls, _ = d(scan_r) dets_cls_norm = np.clip(dets_cls, 0, 0.3) / 0.3 for xy, cls_norm in zip(dets_xy, dets_cls_norm): color = (1.0 - cls_norm, 1.0, 1.0 - cls_norm) c = plt.Circle((xy[0], xy[1]), radius=0.4, color=color, fill=False, linestyle="-") ax.add_artist(c) if _INTERACTIVE: plt.pause(0.1) else: plt.savefig( os.path.join(_SAVE_DIR, f"frame_{data_dict['idx']:04}.png")) if _INTERACTIVE: plt.show()
def _plot_frame_pts(batch_dict, ib, pred_cls, pred_reg, pred_cls_p, pred_reg_p): frame_id = f"{batch_dict['frame_id'][ib]:06d}" sequence = batch_dict["sequence"][ib] fig = plt.figure(figsize=(10, 10)) ax = fig.add_subplot(111) ax.set_xlim(_X_LIM[0], _X_LIM[1]) ax.set_ylim(_Y_LIM[0], _Y_LIM[1]) ax.set_xlabel("x [m]") ax.set_ylabel("y [m]") ax.set_aspect("equal") # ax.set_title(f"Frame {data_dict['idx'][ib]}. Press any key to exit.") # scan and cls label scan_r = batch_dict["scans"][ib][-1] scan_phi = batch_dict["scan_phi"][ib] scan_x, scan_y = u.rphi_to_xy(scan_r, scan_phi) ax.scatter(scan_x, scan_y, s=0.5, c="blue") # annotation ann = batch_dict["dets_wp"][ib] ann_valid_mask = batch_dict["anns_valid_mask"][ib] if len(ann) > 0: ann = np.array(ann) det_x, det_y = u.rphi_to_xy(ann[:, 0], ann[:, 1]) for x, y, valid in zip(det_x, det_y, ann_valid_mask): if valid: # c = plt.Circle((x, y), radius=0.1, color="red", fill=True) c = plt.Circle((x, y), radius=0.4, color="red", fill=False, linestyle="--") ax.add_artist(c) # plot detections if pred_cls is not None and pred_reg is not None: dets_xy, dets_cls, _ = u.nms_predicted_center(scan_r, scan_phi, pred_cls[ib].reshape(-1), pred_reg[ib]) dets_xy = dets_xy[dets_cls >= 0.9438938] # at EER if len(dets_xy) > 0: for x, y in dets_xy: c = plt.Circle((x, y), radius=0.4, color="green", fill=False) ax.add_artist(c) fig_file = os.path.join(_SAVE_DIR, f"figs/{sequence}/scan_det_{frame_id}.png") # plot in addition detections from a pre-trained if pred_cls_p is not None and pred_reg_p is not None: dets_xy, dets_cls, _ = u.nms_predicted_center( scan_r, scan_phi, pred_cls_p[ib].reshape(-1), pred_reg_p[ib]) dets_xy = dets_xy[dets_cls > 0.29919282] # at EER if len(dets_xy) > 0: for x, y in dets_xy: c = plt.Circle((x, y), radius=0.4, color="green", fill=False) ax.add_artist(c) # plot pre-trained detections only elif pred_cls_p is not None and pred_reg_p is not None: dets_xy, dets_cls, _ = u.nms_predicted_center( scan_r, scan_phi, pred_cls_p[ib].reshape(-1), pred_reg_p[ib]) dets_xy = dets_xy[dets_cls > 0.29919282] # at EER if len(dets_xy) > 0: for x, y in dets_xy: c = plt.Circle((x, y), radius=0.4, color="green", fill=False) ax.add_artist(c) fig_file = os.path.join( _SAVE_DIR, f"figs/{sequence}/scan_pretrain_{frame_id}.png") # plot pseudo-labels only else: pl_neg_mask = batch_dict["target_cls"][ib] == 0 ax.scatter(scan_x[pl_neg_mask], scan_y[pl_neg_mask], s=0.5, c="orange") pl_xy = batch_dict["pseudo_label_loc_xy"][ib] if len(pl_xy) > 0: for x, y in pl_xy: c = plt.Circle((x, y), radius=0.4, color="green", fill=False) ax.add_artist(c) fig_file = os.path.join(_SAVE_DIR, f"figs/{sequence}/scan_pl_{frame_id}.png") # save fig os.makedirs(os.path.dirname(fig_file), exist_ok=True) fig.savefig(fig_file, dpi=200) plt.close(fig)
def play_sequence_with_tracking(): # scans seq_name = './data/DROWv2-data/train/lunch_2015-11-26-12-04-23.bag.csv' seq0, seq1 = 109170, 109360 scans, scans_t = [], [] with open(seq_name) as f: for line in f: scan_seq, scan_t, scan = line.split(",", 2) scan_seq = int(scan_seq) if scan_seq < seq0: continue scans.append(np.fromstring(scan, sep=',')) scans_t.append(float(scan_t)) if scan_seq > seq1: break scans = np.stack(scans, axis=0) scans_t = np.array(scans_t) scan_phi = u.get_laser_phi() # odometry, used only for plotting odo_name = seq_name[:-3] + 'odom2' odos = np.genfromtxt(odo_name, delimiter=',') odos_t = odos[:, 1] odos_phi = odos[:, 4] # detector ckpt = './ckpts/dr_spaam_e40.pth' detector = Detector(model_name="DR-SPAAM", ckpt_file=ckpt, gpu=True, stride=1, tracking=True) detector.set_laser_spec(angle_inc=np.radians(0.5), num_pts=450) # scanner location rad_tmp = 0.5 * np.ones(len(scan_phi), dtype=np.float) xy_scanner = u.rphi_to_xy(rad_tmp, scan_phi) xy_scanner = np.stack(xy_scanner, axis=1) # plot fig = plt.figure(figsize=(6, 8)) ax = fig.add_subplot(111) _break = False def p(event): nonlocal _break _break = True fig.canvas.mpl_connect('key_press_event', p) # video sequence odo_idx = 0 for i in range(len(scans)): plt.cla() ax.set_aspect('equal') ax.set_xlim(-10, 5) ax.set_ylim(-5, 15) # ax.set_title('Frame: %s' % i) ax.set_title('Press any key to exit.') ax.axis("off") # find matching odometry while odo_idx < len(odos_t) - 1 and odos_t[odo_idx] < scans_t[i]: odo_idx += 1 odo_phi = odos_phi[odo_idx] odo_rot = np.array( [[np.cos(odo_phi), np.sin(odo_phi)], [-np.sin(odo_phi), np.cos(odo_phi)]], dtype=np.float32) # plot scanner location xy_scanner_rot = np.matmul(xy_scanner, odo_rot.T) ax.plot(xy_scanner_rot[:, 0], xy_scanner_rot[:, 1], c='black') ax.plot((0, xy_scanner_rot[0, 0] * 1.0), (0, xy_scanner_rot[0, 1] * 1.0), c='black') ax.plot((0, xy_scanner_rot[-1, 0] * 1.0), (0, xy_scanner_rot[-1, 1] * 1.0), c='black') # plot points scan = scans[i] scan_x, scan_y = u.rphi_to_xy(scan, scan_phi + odo_phi) ax.scatter(scan_x, scan_y, s=1, c='blue') # inference dets_xy, dets_cls, instance_mask = detector(scan) # plot detection dets_xy_rot = np.matmul(dets_xy, odo_rot.T) cls_thresh = 0.3 for j in range(len(dets_xy)): if dets_cls[j] < cls_thresh: continue c = plt.Circle(dets_xy_rot[j], radius=0.5, color='r', fill=False, linewidth=2) ax.add_artist(c) # plot track cls_thresh = 0.2 tracks, tracks_cls = detector.get_tracklets() for t, tc in zip(tracks, tracks_cls): if tc >= cls_thresh and len(t) > 1: t_rot = np.matmul(t, odo_rot.T) ax.plot(t_rot[:, 0], t_rot[:, 1], color='g', linewidth=2) # plt.savefig('/home/dan/tmp/track3_img/frame_%04d.png' % i) plt.pause(0.001) if _break: break
def _plot_pseudo_labels(batch_dict, ib): # pseudo labels pl_xy = batch_dict["pseudo_label_loc_xy"][ib] pl_boxes = batch_dict["pseudo_label_boxes"][ib] if len(pl_xy) == 0: return # groundtruth anns_rphi = np.array(batch_dict["dets_wp"][ib], dtype=np.float32)[ batch_dict["anns_valid_mask"][ib] ] # match pseudo labels with groundtruth if len(anns_rphi) > 0: gts_x, gts_y = u.rphi_to_xy(anns_rphi[:, 0], anns_rphi[:, 1]) x_diff = pl_xy[:, 0].reshape(-1, 1) - gts_x.reshape(1, -1) y_diff = pl_xy[:, 1].reshape(-1, 1) - gts_y.reshape(1, -1) d_diff = np.sqrt(x_diff * x_diff + y_diff * y_diff) match_found = d_diff < 0.3 # (pl, gt) match_found = match_found.max(axis=1) else: match_found = np.zeros(len(pl_xy), dtype=np.bool) # overlay image with laser im = batch_dict["im_data"][ib]["stitched_image0"] scan_r = batch_dict["scans"][ib][-1] scan_phi = batch_dict["scan_phi"][ib] scan_x, scan_y = u.rphi_to_xy(scan_r, scan_phi) scan_z = batch_dict["laser_z"][ib] scan_xyz_laser = np.stack((scan_x, -scan_y, scan_z), axis=0) # in JRDB laser frame p_xy, ib_mask = jt.transform_pts_laser_to_stitched_im(scan_xyz_laser) p_xy = p_xy[:, ib_mask] far_v = 1.0 far_s = 0 close_v = 0.75 close_s = 0.59 dist_normalized = np.clip(scan_r[ib_mask], 0.0, 20.0) / 20.0 c_hsv = np.empty((1, p_xy.shape[1], 3), dtype=np.float32) c_hsv[0, :, 0] = 0.0 # c_hsv[0, :, 1] = 1.0 - np.clip(scan_r[ib_mask], 0.0, 20.0) / 20.0 c_hsv[0, :, 1] = close_s * (1.0 - dist_normalized) + far_s * dist_normalized c_hsv[0, :, 2] = close_v * (1.0 - dist_normalized) + far_v * dist_normalized c_bgr = cv2.cvtColor(c_hsv, cv2.COLOR_HSV2RGB)[0] # plot frame_id = f"{batch_dict['frame_id'][ib]:06d}" sequence = batch_dict["sequence"][ib] for count, (xy, box, is_pos) in enumerate(zip(pl_xy, pl_boxes, match_found)): # image x0, y0, x1, y1 = box x0 = int(x0) x1 = int(x1) y0 = int(y0) y1 = int(y1) im_box = im[y0 : y1 + 1, x0 : x1 + 1] height = y1 - y0 width = x1 - x0 fig_w_inch = 0.314961 * 2.0 fig_h_inch = 0.708661 * 2.0 fig_im = plt.figure() fig_im.set_size_inches(fig_w_inch, fig_h_inch, forward=False) ax_im = plt.Axes(fig_im, [0.0, 0.0, 1.0, 1.0]) ax_im.imshow(im_box) ax_im.set_axis_off() ax_im.axis(([0, width, height, 0])) ax_im.set_aspect((fig_h_inch / fig_w_inch) / (height / width)) fig_im.add_axes(ax_im) in_box_mask = np.logical_and( np.logical_and(p_xy[0] >= x0, p_xy[0] <= x1), np.logical_and(p_xy[1] >= y0, p_xy[1] <= y1), ) plt.scatter( p_xy[0, in_box_mask] - x0, p_xy[1, in_box_mask] - y0, s=3, c=c_bgr[in_box_mask], ) pos_neg_dir = "true" if is_pos else "false" fig_file = os.path.join( _SAVE_DIR, f"samples/{sequence}/{pos_neg_dir}/{frame_id}_{count}_im.pdf" ) os.makedirs(os.path.dirname(fig_file), exist_ok=True) plt.savefig(fig_file, dpi=height / fig_h_inch) plt.close(fig_im) # lidar plot_range = 0.5 close_mask = np.hypot(scan_x - xy[0], scan_y - xy[1]) < plot_range fig = plt.figure(figsize=(5, 5)) ax = fig.add_subplot() ax.set_aspect("equal") ax.axis("off") # ax.set_xlim(-plot_range, plot_range) # ax.set_ylim(-plot_range, plot_range) # ax.set_xlabel("x [m]") # ax.set_ylabel("y [m]") # ax.set_aspect("equal") # ax.set_title(f"Frame {batch_dict['idx'][ib]}") # plot points in local frame (so it looks aligned with image) ang = np.mean(scan_phi[close_mask]) - 0.5 * np.pi ca, sa = np.cos(ang), np.sin(ang) xy_plotting = np.array([[ca, sa], [-sa, ca]]) @ np.stack( (scan_x[close_mask] - xy[0], scan_y[close_mask] - xy[1]), axis=0 ) ax.scatter( -xy_plotting[0], xy_plotting[1], s=80, color=(191 / 255, 83 / 255, 79 / 255) ) ax.scatter( 0, 0, s=500, color=(18 / 255, 105 / 255, 176 / 255), marker="+", linewidth=5 ) fig_file = os.path.join( _SAVE_DIR, f"samples/{sequence}/{pos_neg_dir}/{frame_id}_{count}_pt.pdf" ) fig.savefig(fig_file) plt.close(fig)
def play_sequence(): # scans seq_name = './data/DROWv2-data/test/run_t_2015-11-26-11-22-03.bag.csv' # seq_name = './data/DROWv2-data/val/run_2015-11-26-15-52-55-k.bag.csv' scans_data = np.genfromtxt(seq_name, delimiter=',') scans_t = scans_data[:, 1] scans = scans_data[:, 2:] scan_phi = u.get_laser_phi() # odometry, used only for plotting odo_name = seq_name[:-3] + 'odom2' odos = np.genfromtxt(odo_name, delimiter=',') odos_t = odos[:, 1] odos_phi = odos[:, 4] # detector ckpt = './ckpts/dr_spaam_e40.pth' detector = Detector(model_name="DR-SPAAM", ckpt_file=ckpt, gpu=True, stride=1) detector.set_laser_spec(angle_inc=np.radians(0.5), num_pts=450) # scanner location rad_tmp = 0.5 * np.ones(len(scan_phi), dtype=np.float) xy_scanner = u.rphi_to_xy(rad_tmp, scan_phi) xy_scanner = np.stack(xy_scanner, axis=1) # plot fig = plt.figure(figsize=(10, 10)) ax = fig.add_subplot(111) _break = False def p(event): nonlocal _break _break = True fig.canvas.mpl_connect('key_press_event', p) # video sequence odo_idx = 0 for i in range(len(scans)): # for i in range(0, len(scans), 20): plt.cla() ax.set_aspect('equal') ax.set_xlim(-15, 15) ax.set_ylim(-15, 15) # ax.set_title('Frame: %s' % i) ax.set_title('Press any key to exit.') ax.axis("off") # find matching odometry while odo_idx < len(odos_t) - 1 and odos_t[odo_idx] < scans_t[i]: odo_idx += 1 odo_phi = odos_phi[odo_idx] odo_rot = np.array( [[np.cos(odo_phi), np.sin(odo_phi)], [-np.sin(odo_phi), np.cos(odo_phi)]], dtype=np.float32) # plot scanner location xy_scanner_rot = np.matmul(xy_scanner, odo_rot.T) ax.plot(xy_scanner_rot[:, 0], xy_scanner_rot[:, 1], c='black') ax.plot((0, xy_scanner_rot[0, 0] * 1.0), (0, xy_scanner_rot[0, 1] * 1.0), c='black') ax.plot((0, xy_scanner_rot[-1, 0] * 1.0), (0, xy_scanner_rot[-1, 1] * 1.0), c='black') # plot points scan = scans[i] scan_x, scan_y = u.rphi_to_xy(scan, scan_phi + odo_phi) ax.scatter(scan_x, scan_y, s=1, c='blue') # inference dets_xy, dets_cls, instance_mask = detector(scan) # plot detection dets_xy_rot = np.matmul(dets_xy, odo_rot.T) cls_thresh = 0.5 for j in range(len(dets_xy)): if dets_cls[j] < cls_thresh: continue # c = plt.Circle(dets_xy_rot[j], radius=0.5, color='r', fill=False) c = plt.Circle(dets_xy_rot[j], radius=0.5, color='r', fill=False, linewidth=2) ax.add_artist(c) # plt.savefig('/home/dan/tmp/det_img/frame_%04d.png' % i) plt.pause(0.001) if _break: break
def _plot_frame(batch_dict, ib): frame_id = f"{batch_dict['frame_id'][ib]:06d}" sequence = batch_dict["sequence"][ib] fig = plt.figure(figsize=(10, 10)) gs = GridSpec(3, 1, figure=fig) ax_im = fig.add_subplot(gs[0, 0]) ax_bev = fig.add_subplot(gs[1:, 0]) ax_bev.set_xlim(_X_LIM[0], _X_LIM[1]) ax_bev.set_ylim(_Y_LIM[0], _Y_LIM[1]) ax_bev.set_xlabel("x [m]") ax_bev.set_ylabel("y [m]") ax_bev.set_aspect("equal") ax_bev.set_title(f"Frame {batch_dict['idx'][ib]}") # scan and cls label scan_r = batch_dict["scans"][ib][-1] scan_phi = batch_dict["scan_phi"][ib] scan_x, scan_y = u.rphi_to_xy(scan_r, scan_phi) target_cls = batch_dict["target_cls"][ib] ax_bev.scatter(scan_x[target_cls == -2], scan_y[target_cls == -2], s=1, c="yellow") ax_bev.scatter(scan_x[target_cls == -1], scan_y[target_cls == -1], s=1, c="orange") ax_bev.scatter(scan_x[target_cls == 0], scan_y[target_cls == 0], s=1, c="black") ax_bev.scatter(scan_x[target_cls > 0], scan_y[target_cls > 0], s=1, c="green") # annotation ann = batch_dict["dets_wp"][ib] ann_valid_mask = batch_dict["anns_valid_mask"][ib] if len(ann) > 0: ann = np.array(ann) det_x, det_y = u.rphi_to_xy(ann[:, 0], ann[:, 1]) for x, y, valid in zip(det_x, det_y, ann_valid_mask): c = "blue" if valid else "orange" c = plt.Circle((x, y), radius=0.4, color=c, fill=False) ax_bev.add_artist(c) # reg label target_reg = batch_dict["target_reg"][ib] dets_r, dets_phi = u.canonical_to_global( scan_r, scan_phi, target_reg[:, 0], target_reg[:, 1] ) dets_r = dets_r[target_cls > 0] dets_phi = dets_phi[target_cls > 0] dets_x, dets_y = u.rphi_to_xy(dets_r, dets_phi) ax_bev.scatter(dets_x, dets_y, s=10, c="red") # image ax_im.axis("off") ax_im.imshow(batch_dict["im_data"][ib]["stitched_image0"]) # detection bounding box for box_dict in batch_dict["im_dets"][ib]: x0, y0, w, h = box_dict["box"] x1 = x0 + w y1 = y0 + h verts = _get_bounding_box_plotting_vertices(x0, y0, x1, y1) c = max(float(box_dict["score"]) - 0.5, 0) * 2.0 ax_im.plot(verts[:, 0], verts[:, 1], c=(1.0 - c, 1.0 - c, 1.0)) for box in batch_dict["pseudo_label_boxes"][ib]: x0, y0, x1, y1 = box verts = _get_bounding_box_plotting_vertices(x0, y0, x1, y1) ax_im.plot(verts[:, 0], verts[:, 1], c="green") # laser points on image scan_z = batch_dict["laser_z"][ib] scan_xyz_laser = np.stack((scan_x, -scan_y, scan_z), axis=0) # in JRDB laser frame p_xy, ib_mask = jt.transform_pts_laser_to_stitched_im(scan_xyz_laser) c = np.clip(scan_r, 0.0, 20.0) / 20.0 c = c.reshape(-1, 1).repeat(3, axis=1) c[:, 0] = 1.0 ax_im.scatter(p_xy[0, ib_mask], p_xy[1, ib_mask], s=1, c=c[ib_mask]) # save fig fig_file = os.path.join(_SAVE_DIR, f"figs/{sequence}/{frame_id}.png") os.makedirs(os.path.dirname(fig_file), exist_ok=True) fig.savefig(fig_file) plt.close(fig)
def _plot_frame_im(batch_dict, ib, show_pseudo_labels=False): frame_id = f"{batch_dict['frame_id'][ib]:06d}" sequence = batch_dict["sequence"][ib] im = batch_dict["im_data"][ib]["stitched_image0"] crop_min_x = 0 im = im[:, crop_min_x:] height = im.shape[0] width = im.shape[1] dpi = height / 1.0 fig = plt.figure() fig.set_size_inches(1.0 * width / height, 1, forward=False) ax_im = plt.Axes(fig, [0.0, 0.0, 1.0, 1.0]) fig.add_axes(ax_im) # image ax_im.axis("off") ax_im.imshow(im) plt.xlim(0, width) plt.ylim(height, 0) # laser points on image scan_r = batch_dict["scans"][ib][-1] scan_phi = batch_dict["scan_phi"][ib] scan_x, scan_y = u.rphi_to_xy(scan_r, scan_phi) scan_z = batch_dict["laser_z"][ib] scan_xyz_laser = np.stack((scan_x, -scan_y, scan_z), axis=0) # in JRDB laser frame p_xy, ib_mask = jt.transform_pts_laser_to_stitched_im(scan_xyz_laser) if show_pseudo_labels: # detection bounding box for box_dict in batch_dict["im_dets"][ib]: x0, y0, w, h = box_dict["box"] x1 = x0 + w y1 = y0 + h verts = _get_bounding_box_plotting_vertices(x0, y0, x1, y1) ax_im.plot(verts[:, 0] - crop_min_x, verts[:, 1], c=(0.0, 0.0, 1.0), alpha=0.3) # c = max(float(box_dict["score"]) - 0.5, 0) * 2.0 # ax_im.plot(verts[:, 0] - crop_min_x, verts[:, 1], c=(1.0 - c, 1.0 - c, 1.0)) # ax_im.plot(verts[:, 0] - crop_min_x, verts[:, 1], # c=(0.0, 0.0, 1.0), alpha=1.0) # x1_large = x1 + 0.05 * w # x0_large = x0 - 0.05 * w # y1_large = y1 + 0.05 * w # y0_large = y0 - 0.05 * w # in_box_mask = np.logical_and( # np.logical_and(p_xy[0] > x0_large, p_xy[0] < x1_large), # np.logical_and(p_xy[1] > y0_large, p_xy[1] < y1_large) # ) # neg_mask[in_box_mask] = False for box in batch_dict["pseudo_label_boxes"][ib]: x0, y0, x1, y1 = box verts = _get_bounding_box_plotting_vertices(x0, y0, x1, y1) ax_im.plot(verts[:, 0] - crop_min_x, verts[:, 1], c="green") # overlay only pseudo-label laser points on image pl_pos_mask = np.logical_and(batch_dict["target_cls"][ib] == 1, ib_mask) pl_neg_mask = np.logical_and(batch_dict["target_cls"][ib] == 0, ib_mask) ax_im.scatter( p_xy[0, pl_pos_mask] - crop_min_x, p_xy[1, pl_pos_mask], s=1, color="green", ) ax_im.scatter( p_xy[0, pl_neg_mask] - crop_min_x, p_xy[1, pl_neg_mask], s=1, color="orange", ) fig_file = os.path.join(_SAVE_DIR, f"figs/{sequence}/im_pl_{frame_id}.png") else: # overlay all laser points on image c_bgr = _distance_to_bgr_color(scan_r) ax_im.scatter(p_xy[0, ib_mask] - crop_min_x, p_xy[1, ib_mask], s=1, color=c_bgr[ib_mask]) fig_file = os.path.join(_SAVE_DIR, f"figs/{sequence}/im_raw_{frame_id}.png") # save fig os.makedirs(os.path.dirname(fig_file), exist_ok=True) fig.savefig(fig_file, dpi=dpi) plt.close(fig)
def _plot_sequence(): jrdb_handle = JRDBHandle( split="train", cfg={ "data_dir": "./data/JRDB", "num_scans": 10, "scan_stride": 1 }, ) color_pool = np.random.uniform(size=(100, 3)) for i, data_dict in enumerate(jrdb_handle): # lidar pc_xyz_upper = jt.transform_pts_upper_velodyne_to_base( data_dict["pc_data"]["upper_velodyne"]) pc_xyz_lower = jt.transform_pts_lower_velodyne_to_base( data_dict["pc_data"]["lower_velodyne"]) # laser laser_r = data_dict["laser_data"][-1] laser_phi = data_dict["laser_grid"] laser_z = data_dict["laser_z"] laser_x, laser_y = u.rphi_to_xy(laser_r, laser_phi) pc_xyz_laser = jt.transform_pts_laser_to_base( np.stack((laser_x, laser_y, laser_z), axis=0)) if _COLOR_INSTANCE: # labels boxes, label_ids = [], [] for ann in data_dict["pc_anns"]: # jrdb_handle.box_is_on_ground(ann) box, b_id = ub3d.box_from_jrdb(ann) boxes.append(box) label_ids.append(b_id) boxes = np.array(boxes) # (B, 7) pc = np.concatenate([pc_xyz_laser, pc_xyz_upper, pc_xyz_lower], axis=1) in_box_mask, closest_box_inds = ub3d.associate_points_and_boxes( pc, boxes, resize_factor=1.0) # plot bg points bg_pc = pc[:, np.logical_not(in_box_mask)] mlab.points3d( bg_pc[0], bg_pc[1], bg_pc[2], scale_factor=0.05, color=(1.0, 0.0, 0.0), ) # plot box and fg points fg_pc = pc[:, in_box_mask] fg_box_inds = closest_box_inds[in_box_mask] corners_xyz, connect_inds = ub3d.boxes_to_corners( boxes, rtn_connect_inds=True) for box_idx, (p_id, corner_xyz) in enumerate(zip(label_ids, corners_xyz)): color = tuple(color_pool[p_id % 100]) # box for inds in connect_inds: mlab.plot3d( corner_xyz[0, inds], corner_xyz[1, inds], corner_xyz[2, inds], tube_radius=None, line_width=5, color=color, ) # point in_box_pc = fg_pc[:, fg_box_inds == box_idx] mlab.points3d( in_box_pc[0], in_box_pc[1], in_box_pc[2], scale_factor=0.05, color=color, ) else: # plot points mlab.points3d( pc_xyz_lower[0], pc_xyz_lower[1], pc_xyz_lower[2], scale_factor=0.05, color=(0.0, 1.0, 0.0), ) mlab.points3d( pc_xyz_upper[0], pc_xyz_upper[1], pc_xyz_upper[2], scale_factor=0.05, color=(0.0, 0.0, 1.0), ) mlab.points3d( pc_xyz_laser[0], pc_xyz_laser[1], pc_xyz_laser[2], scale_factor=0.05, color=(1.0, 0.0, 0.0), ) # plot box boxes = [] for ann in data_dict["pc_anns"]: # jrdb_handle.box_is_on_ground(ann) box = ub3d.box_from_jrdb(ann, fast_mode=False) corners_xyz, connect_inds = box.to_corners( resize_factor=1.0, rtn_connect_inds=True) for inds in connect_inds: mlab.plot3d( corners_xyz[0, inds], corners_xyz[1, inds], corners_xyz[2, inds], tube_radius=None, line_width=5, color=tuple(color_pool[box.get_id() % 100]), ) mlab.show()