in tools/data_prepare/patch_data_prepare_val.py [0:0]
def extract_patch_data_val_detection(split, output_filename,
whitelist=['Car'],
img_height_threshold=25):
''' Extract point clouds in frustums extruded from 2D detection boxes.
Update: Lidar points and 3d boxes are in *rect camera* coord system
(as that in 3d box label files)
Input:
det_filename: string, each line is
img_path typeid confidence xmin ymin xmax ymax
split: string, either trianing or testing
output_filename: string, the name for output .pickle file
whitelist: a list of strings, object types we are interested in.
img_height_threshold: int, neglect image with height lower than that.
lidar_point_threshold: int, neglect frustum with too few points.
Output:
None (will write a .pickle file to the disk)
'''
data_dir = os.path.join(ROOT_DIR, 'data')
dataset = KittiDataset(root_dir=data_dir, split=split)
# det_id_list, det_type_list, det_box2d_list, det_prob_list = read_det_file(det_filename)
# cache_id = -1
# cache = None
feat_dir = '../../data/KITTI/pickle_files/org'
id_list = []
type_list = []
box2d_list = []
prob_list = []
patch_xyz_list = []
patch_rgb_list = []
frustum_angle_list = []
feats_list = []
progress_bar = tqdm.tqdm(total=len(dataset.idx_list), leave=True, desc='%s split patch data gen (from 2d detections)' % split)
for data_idx in dataset.idx_list: # image idx
data_idx = int(data_idx)
calib = dataset.get_calib(data_idx)
objects = dataset.get_label(data_idx)
feature_path = os.path.join(feat_dir, '%06d.pickle' % data_idx)
df = open(feature_path, 'rb')
datas = pickle.load(df)
df.close()
features = datas['features'].data.cpu()
trans_out = datas['trans']
c = trans_out[0]
s = trans_out[1]
out_h = trans_out[2]
out_w = trans_out[3]
trans_output = get_affine_transform(c, s, 0, [out_w, out_h])
# compute x,y,z for each pixel in depth map
depth = dataset.get_depth(data_idx)
image = dataset.get_image(data_idx)
assert depth.size == image.size
width, height = depth.size
depth = np.array(depth).astype(np.float32) / 256
uvdepth = np.zeros((height, width, 3), dtype=np.float32)
for v in range(height):
for u in range(width):
uvdepth[v, u, 0] = u
uvdepth[v, u, 1] = v
uvdepth[:, :, 2] = depth
uvdepth = uvdepth.reshape(-1, 3)
xyz = calib.img_to_rect(uvdepth[:, 0], uvdepth[:, 1], uvdepth[:, 2]) # rect coord sys
xyz = xyz.reshape(height, width, 3) # record xyz, data type: float32
uvdepth = uvdepth.reshape(height, width, 3)
rgb = np.array(image)
for object in objects:
if object.cls_type not in whitelist:
continue
# get 2d box from ground truth
box2d = object.box2d
xmin, ymin, xmax, ymax = box2d
if ymax - ymin < img_height_threshold:
progress_bar.update()
continue
# Get frustum angle (according to center pixel in 2D BOX)
box2d_center = np.array([(xmin + xmax) / 2.0, (ymin + ymax) / 2.0])
uvdepth = np.zeros((1, 3))
uvdepth[0, 0:2] = box2d_center
uvdepth[0, 2] = 20 # some random depth
box2d_center_rect = calib.img_to_rect(uvdepth[:, 0], uvdepth[:, 1], uvdepth[:, 2])
frustum_angle = -1 * np.arctan2(box2d_center_rect[0, 2], box2d_center_rect[0, 0])
# filter
# if object.level > 3: continue
xmin, ymin = max(xmin, 0), max(ymin, 0) # check range
xmax, ymax = min(xmax, width), min(ymax, height) # check range
patch_xyz = xyz[int(ymin):int(ymax), int(xmin):int(xmax), :]
patch_rgb = rgb[int(ymin):int(ymax), int(xmin):int(xmax), :]
if int(ymin)>=int(ymax) or int(xmin)>=int(xmax):#patch_xyz.shape[0] == 0 or patch_xyz.shape[0] == 0:
continue
x1_out, y1_out = affine_transform((xmin, ymin), trans_output)
x2_out, y2_out = affine_transform((xmax, ymax), trans_output)
roi_feat = torchvision.ops.roi_align(features,
[torch.FloatTensor(np.array([[x1_out, y1_out, x2_out, y2_out]]))],
(16, 16)).squeeze(0).numpy()
feats_list.append(roi_feat)
id_list.append(data_idx)
box2d_list.append(box2d)
patch_xyz_list.append(patch_xyz)
patch_rgb_list.append(patch_rgb)
type_list.append(object.cls_type)
frustum_angle_list.append(frustum_angle)
prob_list.append(object.score)
progress_bar.update()
progress_bar.close()
with open(output_filename, 'wb') as fp:
pickle.dump(id_list, fp)
pickle.dump(box2d_list, fp)
pickle.dump(patch_xyz_list, fp)
pickle.dump(patch_rgb_list, fp)
pickle.dump(type_list, fp)
pickle.dump(frustum_angle_list, fp)
pickle.dump(prob_list, fp)
with open(output_filename.replace('.pickle','_feat.pickle'),'wb') as fp:
pickle.dump(feats_list, fp)