frustum pointnets训练代码学习笔记——kitti_object.py
本文记录了博主学习frustum pointnets过程中遇到的2D和3D数据库显示程序。为了画出输出结果,博主希望在这个程序的基础上修改一个可以显示结果的程序。更新于2018.09.22。
本文首先给出代码原文的学习笔记,随后整理出修改后的结果显示程序,如果公开,会在这里放上链接,如果有帮助,请在代码页面点一下小星星哦。附可能有用的信息:各个集合所用的文件名在kitti/image_sets
文件夹下。
文章目录
总结
把总结写在前面,根据需要判断是否需要详细看源码分析。
- 这个文件的主要功能就是将KITTI库中的2d和3d结果画出来(至于是画training还是testing,文件的kitti_object函数的初始函数和get_label_objects分别有定义和判断)。
- 文件个数是人为在kitti_object函数中设定的,并非自动提取。
- 画图是机械第从第一个图片一直向后显示,且如果该图片中有多个目标,仅显示txt文件中排在第一的那个的数据。通过修改
objects[0].print_object()
中[]里面的标号可以指定画第几个目标,不过要注意的是,每个图片中含有的目标个数是不同的。
用到的语法规则
这一部分记录了代码原文中出现的语法规则,并不影响代码功能的理解,但是可能方便日后的使用,因此在这里记录下来。
from __future__ import print_function
加上这句话以后,即使在python2.X也要像python3.X一样的语法使用print函数(加括号)。类似地,如果有其他新的功能特性且该特性与当前版本中的使用不兼容,就可以从future模块导入。详细说明参考这里。
from PIL import Image
PIL已经是python平台事实上的图像处理标准库了,全称为Python Imaging Library。具体的使用方法说明可以参考这里。
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
其中,__file__
就是当前所执行的文件,也就是kitti_object.py
。os.path.abspath
命令获取的是当前文件的绝对路径,比如博主的运行结果:
>>> print os.path.abspath("/home/galaxy/Work/XXX/Pointnets/frustum-pointnets-master/kitti/kitti_object.py")
/home/galaxy/Work/XXX/Pointnets/frustum-pointnets-master/kitti/kitti_object.py
而前面的os.path.dirname
获取的则是当前路径所存在于的文件夹,因此,BASE_DIR指向的就是kitti_object.py
所处的文件夹了。运行结果为:
>>> print os.path.dirname(os.path.abspath("/home/galaxy/Work/XXX/Pointnets/frustum-pointnets-master/kitti/kitti_object.py"))
/home/galaxy/Work/XXX/Pointnets/frustum-pointnets-master/kitti
sys.path.append(os.path.join(ROOT_DIR, 'mayavi'))
其中,os.path.join
用于路径拼接。
sys.path.append
:在导入一个模块时,默认情况下python会搜索当前目录、已安装的内置模块和第三方模块,搜索路径存放在sys模块的path中。如果要用的模块和当前脚本不在一个目录下,就需要将其添加到path中。这种修改是临时的,脚本运行后失效。
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
opencv中提供了cvtColor函数用于实现图像格式类型的相互转换,具体说明可以参照这里。
代码原文分析
#代码作者信息
''' Helper class and functions for loading KITTI objects
Author: Charles R. Qi
Date: September 2017
'''
#加载必要的库
from __future__ import print_function
import os
import sys
import numpy as np
import cv2
from PIL import Image
#定义基础路径
BASE_DIR = os.path.dirname(os.path.abspath(__file__)) #指向当前文件所在文件夹(kitti)
ROOT_DIR = os.path.dirname(BASE_DIR) #指向frustum文件夹
sys.path.append(os.path.join(ROOT_DIR, 'mayavi'))
import kitti_util as utils #加载论文作者写的库
try:
raw_input # Python 2
except NameError:
raw_input = input # Python 3
#用于获取数据库各项路径路径(training或testing)
class kitti_object(object):
'''Load and parse object data into a usable format.'''
def __init__(self, root_dir, split='training'):
'''root_dir contains training and testing folders'''
self.root_dir = root_dir
self.split = split
self.split_dir = os.path.join(root_dir, split)
if split == 'training':
self.num_samples = 7481
elif split == 'testing':
self.num_samples = 7518
else:
print('Unknown split: %s' % (split))
exit(-1)
self.image_dir = os.path.join(self.split_dir, 'image_2')
self.calib_dir = os.path.join(self.split_dir, 'calib')
self.lidar_dir = os.path.join(self.split_dir, 'velodyne')
self.label_dir = os.path.join(self.split_dir, 'label_2')
# 用于后面获取样本库内的样本总数
def __len__(self):
return self.num_samples
def get_image(self, idx):
assert(idx<self.num_samples)
img_filename = os.path.join(self.image_dir, '%06d.png'%(idx))
return utils.load_image(img_filename)
def get_lidar(self, idx):
assert(idx<self.num_samples)
lidar_filename = os.path.join(self.lidar_dir, '%06d.bin'%(idx))
return utils.load_velo_scan(lidar_filename)
def get_calibration(self, idx):
assert(idx<self.num_samples)
calib_filename = os.path.join(self.calib_dir, '%06d.txt'%(idx))
return utils.Calibration(calib_filename)
# 获取idx指示的样本对应的label文件路径,并按行读取文件,返回文件内容
def get_label_objects(self, idx):
assert(idx<self.num_samples and self.split=='training')
label_filename = os.path.join(self.label_dir, '%06d.txt'%(idx))
return utils.read_label(label_filename)
def get_depth_map(self, idx):
pass
def get_top_down(self, idx):
pass
class kitti_object_video(object):
''' Load data for KITTI videos '''
def __init__(self, img_dir, lidar_dir, calib_dir):
self.calib = utils.Calibration(calib_dir, from_video=True)
self.img_dir = img_dir
self.lidar_dir = lidar_dir
self.img_filenames = sorted([os.path.join(img_dir, filename) \
for filename in os.listdir(img_dir)])
self.lidar_filenames = sorted([os.path.join(lidar_dir, filename) \
for filename in os.listdir(lidar_dir)])
print(len(self.img_filenames))
print(len(self.lidar_filenames))
#assert(len(self.img_filenames) == len(self.lidar_filenames))
self.num_samples = len(self.img_filenames)
def __len__(self):
return self.num_samples
def get_image(self, idx):
assert(idx<self.num_samples)
img_filename = self.img_filenames[idx]
return utils.load_image(img_filename)
def get_lidar(self, idx):
assert(idx<self.num_samples)
lidar_filename = self.lidar_filenames[idx]
return utils.load_velo_scan(lidar_filename)
def get_calibration(self, unused):
return self.calib
def viz_kitti_video():
video_path = os.path.join(ROOT_DIR, 'dataset/2011_09_26/')
dataset = kitti_object_video(\
os.path.join(video_path, '2011_09_26_drive_0023_sync/image_02/data'),
os.path.join(video_path, '2011_09_26_drive_0023_sync/velodyne_points/data'),
video_path)
print(len(dataset))
for i in range(len(dataset)):
img = dataset.get_image(0)
pc = dataset.get_lidar(0)
Image.fromarray(img).show()
draw_lidar(pc)
raw_input()
pc[:,0:3] = dataset.get_calibration().project_velo_to_rect(pc[:,0:3])
draw_lidar(pc)
raw_input()
return
def show_image_with_boxes(img, objects, calib, show3d=True):
''' Show image with 2D bounding boxes '''
img1 = np.copy(img) # for 2d bbox
img2 = np.copy(img) # for 3d bbox
for obj in objects:
if obj.type=='DontCare':continue
cv2.rectangle(img1, (int(obj.xmin),int(obj.ymin)),
(int(obj.xmax),int(obj.ymax)), (0,255,0), 2)
box3d_pts_2d, box3d_pts_3d = utils.compute_box_3d(obj, calib.P)
img2 = utils.draw_projected_box3d(img2, box3d_pts_2d)
Image.fromarray(img1).show()
if show3d:
Image.fromarray(img2).show()
def get_lidar_in_image_fov(pc_velo, calib, xmin, ymin, xmax, ymax,
return_more=False, clip_distance=2.0):
''' Filter lidar points, keep those in image FOV '''
pts_2d = calib.project_velo_to_image(pc_velo)
fov_inds = (pts_2d[:,0]<xmax) & (pts_2d[:,0]>=xmin) & \
(pts_2d[:,1]<ymax) & (pts_2d[:,1]>=ymin)
fov_inds = fov_inds & (pc_velo[:,0]>clip_distance)
imgfov_pc_velo = pc_velo[fov_inds,:]
if return_more:
return imgfov_pc_velo, pts_2d, fov_inds
else:
return imgfov_pc_velo
def show_lidar_with_boxes(pc_velo, objects, calib,
img_fov=False, img_width=None, img_height=None):
''' Show all LiDAR points.
Draw 3d box in LiDAR point cloud (in velo coord system) '''
if 'mlab' not in sys.modules: import mayavi.mlab as mlab
from viz_util import draw_lidar_simple, draw_lidar, draw_gt_boxes3d
print(('All point num: ', pc_velo.shape[0]))
fig = mlab.figure(figure=None, bgcolor=(0,0,0),
fgcolor=None, engine=None, size=(1000, 500))
if img_fov:
pc_velo = get_lidar_in_image_fov(pc_velo, calib, 0, 0,
img_width, img_height)
print(('FOV point num: ', pc_velo.shape[0]))
draw_lidar(pc_velo, fig=fig)
for obj in objects:
if obj.type=='DontCare':continue
# Draw 3d bounding box
box3d_pts_2d, box3d_pts_3d = utils.compute_box_3d(obj, calib.P)
box3d_pts_3d_velo = calib.project_rect_to_velo(box3d_pts_3d)
# Draw heading arrow
ori3d_pts_2d, ori3d_pts_3d = utils.compute_orientation_3d(obj, calib.P)
ori3d_pts_3d_velo = calib.project_rect_to_velo(ori3d_pts_3d)
x1,y1,z1 = ori3d_pts_3d_velo[0,:]
x2,y2,z2 = ori3d_pts_3d_velo[1,:]
draw_gt_boxes3d([box3d_pts_3d_velo], fig=fig)
mlab.plot3d([x1, x2], [y1, y2], [z1,z2], color=(0.5,0.5,0.5),
tube_radius=None, line_width=1, figure=fig)
mlab.show(1)
def show_lidar_on_image(pc_velo, img, calib, img_width, img_height):
''' Project LiDAR points to image '''
imgfov_pc_velo, pts_2d, fov_inds = get_lidar_in_image_fov(pc_velo,
calib, 0, 0, img_width, img_height, True)
imgfov_pts_2d = pts_2d[fov_inds,:]
imgfov_pc_rect = calib.project_velo_to_rect(imgfov_pc_velo)
import matplotlib.pyplot as plt
cmap = plt.cm.get_cmap('hsv', 256)
cmap = np.array([cmap(i) for i in range(256)])[:,:3]*255
for i in range(imgfov_pts_2d.shape[0]):
depth = imgfov_pc_rect[i,2]
color = cmap[int(640.0/depth),:]
cv2.circle(img, (int(np.round(imgfov_pts_2d[i,0])),
int(np.round(imgfov_pts_2d[i,1]))),
2, color=tuple(color), thickness=-1)
Image.fromarray(img).show()
return img
def dataset_viz():
dataset = kitti_object(os.path.join(ROOT_DIR, 'dataset/KITTI/object')) #获取数据库各项路径
for data_idx in range(len(dataset)): #从0开始到len获取的数据库样本总数
# 从数据库中加载数据
objects = dataset.get_label_objects(data_idx) #获取data_idx对应的结果
objects[0].print_object() #在屏幕上输出data_idx对应的第一个结果(如果有多个,修改[]内的值就可以变成对应的结果
img = dataset.get_image(data_idx) #获取data_idx对应的图片
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) #图像格式转换
img_height, img_width, img_channel = img.shape
print(('Image shape: ', img.shape))
pc_velo = dataset.get_lidar(data_idx)[:,0:3] #获取data_idx对应的3D点云
calib = dataset.get_calibration(data_idx)
# 在图像上画出2d和3dboxes
show_image_with_boxes(img, objects, calib, False)
raw_input()
# Show all LiDAR points. Draw 3d box in LiDAR point cloud
show_lidar_with_boxes(pc_velo, objects, calib, True, img_width, img_height)
raw_input()
if __name__=='__main__':
import mayavi.mlab as mlab
from viz_util import draw_lidar_simple, draw_lidar, draw_gt_boxes3d
dataset_viz() #显示数据