File size: 5,071 Bytes
62a2f1c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 |
import json
import os.path as osp
from collections import defaultdict
import cv2
import numpy as np
class Octopus(object):
"""
dataset structure:
- data_root
- train_split.txt
- val_split.txt
- test_split.txt
-
"""
camera_names = ['cam01', 'cam03', 'cam05', 'cam06', 'cam07', 'cam08', 'cam09']
camera_tags = ['top', 'top2', 'left_back', 'left_front', 'right_front', 'right_back', 'back']
def __init__(self, dataset_root):
self.dataset_root = dataset_root
self.data_root = osp.join(self.dataset_root, 'data')
self._collect_basic_infos()
@property
def train_split_list(self):
if not osp.isfile(osp.join(self.dataset_root, 'ImageSets', 'train_set.txt')):
train_split_list = None
else:
train_split_list = set(map(lambda x: x.strip(),
open(osp.join(self.data_root, 'train_set.txt')).readlines()))
return train_split_list
@property
def val_split_list(self):
if not osp.isfile(osp.join(self.dataset_root, 'ImageSets', 'val_set.txt')):
val_split_list = None
else:
val_split_list = set(map(lambda x: x.strip(),
open(osp.join(self.data_root, 'val_set.txt')).readlines()))
return val_split_list
@property
def test_split_list(self):
if not osp.isfile(osp.join(self.dataset_root, 'ImageSets', 'test_set.txt')):
test_split_list = None
else:
test_split_list = set(map(lambda x: x.strip(),
open(osp.join(self.data_root, 'test_set.txt')).readlines()))
return test_split_list
@property
def raw_split_list(self):
if not osp.isfile(osp.join(self.dataset_root, 'ImageSets', 'raw_set.txt')):
raw_split_list = None
else:
raw_split_list = set(map(lambda x: x.strip(),
open(osp.join(self.data_root, 'raw_set.txt')).readlines()))
return raw_split_list
def _find_split_name(self, seq_id):
if seq_id in self.raw_split_list:
return 'raw'
if seq_id in self.train_split_list:
return 'train'
if seq_id in self.test_split_list:
return 'test'
if seq_id in self.val_split_list:
return 'val'
print("sequence id {} corresponding to no split".format(seq_id))
raise NotImplementedError
def _collect_basic_infos(self):
self.train_info = defaultdict(dict)
if self.train_split_list is not None:
for train_seq in self.train_split_list:
anno_file_path = osp.join(self.data_root, train_seq, '{}.json'.format(train_seq))
if not osp.isfile(anno_file_path):
print("no annotation file for sequence {}".format(train_seq))
raise FileNotFoundError
anno_file = json.load(open(anno_file_path, 'r'))
for frame_anno in anno_file['frames']:
self.train_info[train_seq][frame_anno['frame_id']] = {
'pose': frame_anno['pose'],
'calib': anno_file['calib'],
}
def get_frame_anno(self, seq_id, frame_id):
split_name = self._find_split_name(seq_id)
frame_info = getattr(self, '{}_info'.format(split_name))[seq_id][frame_id]
if 'anno' in frame_info:
return frame_info['anno']
return None
def load_point_cloud(self, seq_id, frame_id):
bin_path = osp.join(self.data_root, seq_id, 'lidar_roof', '{}.bin'.format(frame_id))
points = np.fromfile(bin_path, dtype=np.float32).reshape(-1, 4)
return points
def load_image(self, seq_id, frame_id, cam_name):
cam_path = osp.join(self.data_root, seq_id, cam_name, '{}.jpg'.format(frame_id))
img_buf = cv2.cvtColor(cv2.imread(cam_path), cv2.COLOR_BGR2RGB)
return img_buf
def project_lidar_to_image(self, seq_id, frame_id):
points = self.load_point_cloud(seq_id, frame_id)
split_name = self._find_split_name(seq_id)
frame_info = getattr(self, '{}_info'.format(split_name))[seq_id][frame_id]
points_img_dict = dict()
for cam_name in self.__class__.camera_names:
calib_info = frame_info['calib'][cam_name]
cam_2_velo = calib_info['cam_to_velo']
cam_intri = calib_info['cam_intrinsic']
point_xyz = points[:, :3]
points_homo = np.hstack(
[point_xyz, np.ones(point_xyz.shape[0], dtype=np.float32).reshape((-1, 1))])
points_lidar = np.dot(points_homo, np.linalg.inv(cam_2_velo).T)
mask = points_lidar[:, 2] > 0
points_lidar = points_lidar[mask]
points_img = np.dot(points_lidar, cam_intri.T)
points_img_dict[cam_name] = points_img
return points_img_dict
def undistort_image(self, seq_id, frame_id):
pass |