def __getitem__(self, index):
sample_id = int(self.sample_id_list[index])
if self.mode in ['train', 'val']:
lidarData = self.get_lidar(sample_id)
objects = self.get_label(sample_id)
calib = self.get_calib(sample_id)
labels, noObjectLabels = kitti_bev_utils.read_labels_for_bevbox(
objects)
if not noObjectLabels:
labels[:, 1:] = transformation.camera_to_lidar_box(
labels[:, 1:], calib.V2C, calib.R0,
calib.P) # convert rect cam to velo cord
if self.data_aug and self.mode == 'train':
lidarData, labels[:,
1:] = transformation.complex_yolo_pc_augmentation(
lidarData, labels[:, 1:], True)
b = kitti_bev_utils.removePoints(lidarData, cnf.boundary)
rgb_map = kitti_bev_utils.makeBVFeature(b, cnf.DISCRETIZATION,
cnf.boundary)
target = kitti_bev_utils.build_yolo_target(labels)
img_file = os.path.join(self.image_dir,
'{:06d}.png'.format(sample_id))
ntargets = 0
for i, t in enumerate(target):
if t.sum(0):
ntargets += 1
# (box_idx, class, y, x, w, l, sin(yaw), cos(yaw))
targets = torch.zeros((ntargets, 8))
for i, t in enumerate(target):
if t.sum(0):
targets[i, 1:] = torch.from_numpy(t)
img = torch.from_numpy(rgb_map).float()
if self.data_aug:
if np.random.random() < 0.5:
img, targets = self.horizontal_flip(img, targets)
return img_file, img, targets
else:
lidarData = self.get_lidar(sample_id)
b = kitti_bev_utils.removePoints(lidarData, cnf.boundary)
rgb_map = kitti_bev_utils.makeBVFeature(b, cnf.DISCRETIZATION,
cnf.boundary)
img_file = os.path.join(self.image_dir,
'{:06d}.png'.format(sample_id))
return img_file, rgb_map
def load_img_with_targets(self, index):
"""Load images and targets for the training and validation phase"""
sample_id = int(self.sample_id_list[index])
lidarData = self.get_lidar(sample_id)
objects = self.get_label(sample_id)
calib = self.get_calib(sample_id)
labels, noObjectLabels = kitti_bev_utils.read_labels_for_bevbox(objects)
if not noObjectLabels:
labels[:, 1:] = transformation.camera_to_lidar_box(labels[:, 1:], calib.V2C, calib.R0,
calib.P) # convert rect cam to velo cord
if self.aug_transforms is not None:
lidarData, labels[:, 1:] = self.aug_transforms(lidarData, labels[:, 1:])
b = kitti_bev_utils.removePoints(lidarData, cnf.boundary)
rgb_map = kitti_bev_utils.makeBVFeature(b, cnf.DISCRETIZATION, cnf.boundary)
target = kitti_bev_utils.build_yolo_target(labels)
img_file = os.path.join(self.image_dir, '{:06d}.png'.format(sample_id))
# on image space: targets are formatted as (box_idx, class, x, y, w, l, sin(yaw), cos(yaw))
n_target = len(target)
targets = torch.zeros((n_target, 8))
if n_target > 0:
targets[:, 1:] = torch.from_numpy(target)
rgb_map = torch.from_numpy(rgb_map).float()
if self.hflip_transform is not None:
rgb_map, targets = self.hflip_transform(rgb_map, targets)
return img_file, rgb_map, targets
def load_img_only(self, index):
"""Load only image for the testing phase"""
sample_id = int(self.sample_id_list[index])
lidarData = self.get_lidar(sample_id)
b = kitti_bev_utils.removePoints(lidarData, cnf.boundary)
rgb_map = kitti_bev_utils.makeBVFeature(b, cnf.DISCRETIZATION, cnf.boundary)
img_file = os.path.join(self.image_dir, '{:06d}.png'.format(sample_id))
return img_file, rgb_map
def callback(self,data):
rospy.loginfo("detection")
with torch.no_grad():
gen = point_cloud2.read_points(data)
for idx, p in enumerate(gen):
print(idx)
print(p)
b = kitti_bev_utils.removePoints(gen, cnf.boundary)
imgs_bev = kitti_bev_utils.makeBVFeature(b, cnf.DISCRETIZATION, cnf.boundary)
input_imgs = imgs_bev.to(device=configs.device).float()
t1 = time_synchronized()
outputs = self.model(input_imgs)
t2 = time_synchronized()
detections = post_processing_v2(outputs, conf_thresh=configs.conf_thresh, nms_thresh=configs.nms_thresh)
img_detections = [] # Stores detections for each image index
img_detections.extend(detections)
img_bev = imgs_bev.squeeze() * 255
img_bev = img_bev.permute(1, 2, 0).numpy().astype(np.uint8)
img_bev = cv2.resize(img_bev, (configs.img_size, configs.img_size))
for detections in img_detections:
if detections is None:
continue
# Rescale boxes to original image
detections = rescale_boxes(detections, configs.img_size, img_bev.shape[:2])
for x, y, w, l, im, re, *_, cls_pred in detections:
yaw = np.arctan2(im, re)
# Draw rotated box
kitti_bev_utils.drawRotatedBox(img_bev, x, y, w, l, yaw, cnf.colors[int(cls_pred)])
img_bev = cv2.flip(cv2.flip(img_bev, 0), 1)
out_img = img_bev
cv2.imshow('test-img', out_img)
cv2.waitKey(1)
