def shuffle_log(subset, log: ArgoverseTrackingLoader):
index = np.arange(log.num_lidar_frame)
random.shuffle(index)
for idx in index:
lidar = log.get_lidar(idx)
label = log.get_label_object(idx)
yield idx, subset, lidar, label, log
def test_loading_image_lidar(data_loader: ArgoverseTrackingLoader) -> None:
camera = CAMERA_LIST[0]
log = "1"
image_1 = data_loader.get_image_at_timestamp(0, camera, log)
image_2 = data_loader.get_image_list_sync(camera, log, load=True)[0]
image_3 = data_loader.get_image(0, camera, log)
image_4 = data_loader.get_image_sync(0, camera, log)
assert np.array_equal(image_1, image_2) and np.array_equal(
image_1, image_3) and np.array_equal(image_1, image_4)
lidar_0 = data_loader.get_lidar(0, log)
lidar_gt = np.array([
[0.0, 0.0, 5.0],
[1.0, 0.0, 5.0],
[2.0, 0.0, 5.0],
[3.0, 0.0, 5.0],
[4.0, 0.0, 5.0],
[5.0, 0.0, 5.0],
[6.0, 0.0, 5.0],
[7.0, 0.0, 5.0],
[8.0, 0.0, 5.0],
[9.0, 0.0, 5.0],
])
assert np.array_equal(lidar_0, lidar_gt)
def test_calib(data_loader: ArgoverseTrackingLoader) -> None:
assert data_loader.calib
camera = "ring_front_center"
calib = data_loader.get_calibration(camera)
pc = data_loader.get_lidar(0)
uv = calib.project_ego_to_image(pc)
uv_cam = calib.project_ego_to_cam(pc)
assert (uv == calib.project_cam_to_image(uv_cam)).all
assert (uv_cam == calib.project_image_to_cam(uv)).all
assert (pc == calib.project_image_to_ego(uv)).all
def __init__(self,
data_path,
filenames,
height,
width,
frame_idxs,
num_scales,
is_train=False,
img_ext='.jpg'):
super(Argo_MonoDataset, self).__init__()
self.data_path = data_path
self.filenames = filenames
self.height = height #2056
self.width = width #2464
self.num_scales = num_scales
self.interp = Image.ANTIALIAS
self.frame_idxs = frame_idxs
self.is_train = is_train
self.img_ext = img_ext
self.loader = argoverse_load_image
self.argoverse_loader = ArgoverseTrackingLoader(self.data_path)
self.camera = self.argoverse_loader.CAMERA_LIST[0]
self.argoverse_data = 0
self.to_tensor = transforms.ToTensor()
# We need to specify augmentations differently in newer versions of torchvision.
# We first try the newer tuple version; if this fails we fall back to scalars
try:
self.brightness = (0.8, 1.2)
self.contrast = (0.8, 1.2)
self.saturation = (0.8, 1.2)
self.hue = (-0.1, 0.1)
transforms.ColorJitter.get_params(self.brightness, self.contrast,
self.saturation, self.hue)
except TypeError:
self.brightness = 0.2
self.contrast = 0.2
self.saturation = 0.2
self.hue = 0.1
self.resize = {}
for i in range(self.num_scales):
s = 2**i
self.resize[i] = transforms.Resize(
(self.height // s, self.width // s), interpolation=self.interp)
self.load_depth = self.check_depth()
def extract_datapoints(root_dir, test_set=False):
argoverse_loader = ArgoverseTrackingLoader(root_dir=root_dir)
data = []
for log_id in argoverse_loader.log_list:
# print(f"Processing log: {os.path.join(root_dir, log_id)}", end="\r")
argoverse_data = argoverse_loader.get(log_id=log_id)
# calibL = argoverse_data.get_calibration(camera=camL, log_id=log_id)
# calibR = argoverse_data.get_calibration(camera=camR, log_id=log_id)
calibs = {}
for cam in STEREO_CAMERA_LIST + RING_CAMERA_LIST:
calibs[cam] = argoverse_data.get_calibration(camera=cam,
log_id=log_id)
count = 0
for lidar_timestamp in argoverse_data.lidar_timestamp_list:
data_point = dict()
data_point["log_id"] = log_id
data_point["frame_id"] = count
count += 1
for cam in STEREO_CAMERA_LIST + RING_CAMERA_LIST:
cam_timestamp = argoverse_loader.sync.get_closest_cam_channel_timestamp(
lidar_timestamp=lidar_timestamp,
camera_name=cam,
log_id=log_id)
if cam_timestamp is not None:
data_point[cam] = argoverse_loader.get_image_at_timestamp(
timestamp=cam_timestamp,
camera=cam,
log_id=log_id,
load=False)
else:
data_point[cam] = None
data_point["timestamp"] = lidar_timestamp
data_point["lidar"] = argoverse_loader.timestamp_lidar_dict[
lidar_timestamp]
data_point["calibs"] = calibs
d = argoverse_data.get_pose(
argoverse_data.get_idx_from_timestamp(
lidar_timestamp)).translation
r = Rotation.from_dcm(
argoverse_data.get_pose(
argoverse_data.get_idx_from_timestamp(
lidar_timestamp)).rotation)
data_point["pose"] = (d, r.as_euler('xyz'))
if not test_set:
data_point["labels"] = argoverse_loader.get_label_object(
idx=argoverse_loader.lidar_timestamp_list.index(
lidar_timestamp),
log_id=log_id)
data.append(data_point)
return data
def __init__(self, args, transform=None):
self.args = args
self.transform = transform
self.argoverse_loader = ArgoverseTrackingLoader(args.root_dir)
self.image_list = []
for log_id in self.argoverse_loader.log_list:
argoverse_data = self.argoverse_loader.get(log_id)
for idx in range(
len(argoverse_data.image_list_sync['ring_front_center'])):
self.image_list.append(str(log_id) + '@' + str(idx))
self.camera_name = [
'ring_front_center', 'ring_side_left', 'ring_side_right',
'ring_front_left', 'ring_front_right', 'ring_rear_right',
'ring_rear_left'
]
def draw_point_cloud(
ax: plt.Axes,
title: str,
argoverse_data: ArgoverseTrackingLoader,
idx: int,
axes: Optional[Any] = None,
xlim3d: Any = None,
ylim3d: Any = None,
zlim3d: Any = None,
) -> None:
axes = _get_axes_or_default(axes)
pc = argoverse_data.get_lidar(idx)
assert isinstance(pc, np.ndarray)
objects = argoverse_data.get_label_object(idx)
ax.scatter(*np.transpose(pc[:, axes]),
s=point_size,
c=pc[:, 2],
cmap="gray")
ax.set_title(title)
ax.set_xlabel("{} axis".format(axes_str[axes[0]]))
ax.set_ylabel("{} axis".format(axes_str[axes[1]]))
if len(axes) > 2:
ax.set_xlim3d(*axes_limits[axes[0]])
ax.set_ylim3d(*axes_limits[axes[1]])
ax.set_zlim3d(*axes_limits[axes[2]])
ax.set_zlabel("{} axis".format(axes_str[axes[2]]))
else:
ax.set_xlim(*axes_limits[axes[0]])
ax.set_ylim(*axes_limits[axes[1]])
# User specified limits
if xlim3d != None:
ax.set_xlim3d(xlim3d)
if ylim3d != None:
ax.set_ylim3d(ylim3d)
if zlim3d != None:
ax.set_zlim3d(zlim3d)
for obj in objects:
if obj.occlusion == 100:
continue
if obj.label_class is None:
logger.warning("No label class, just picking the default color.")
color = _COLOR_MAP[-1]
else:
color = _COLOR_MAP[OBJ_CLASS_MAPPING_DICT[obj.label_class]]
draw_box(ax, obj.as_3d_bbox().T, axes=axes, color=color)
def __init__(self, info_path, root_path, class_names, num_point_features,
target_assigner, feature_map_size, prep_func):
#with open(info_path, 'rb') as f:
# infos = pickle.load(f)
self._root_path = root_path
self.argoverse_loader = ArgoverseTrackingLoader(root_path)
self.am = ArgoverseMap()
self.inform = info_path
self._num_point_features = num_point_features
self.class_names = class_names
#print("remain number of infos:", len(self._kitti_infos))
# generate anchors cache
# [352, 400]
ret = target_assigner.generate_anchors(feature_map_size)
anchors = ret["anchors"]
anchors = anchors.reshape([-1, 7])
matched_thresholds = ret["matched_thresholds"]
unmatched_thresholds = ret["unmatched_thresholds"]
anchors_bv = box_np_ops.rbbox2d_to_near_bbox(anchors[:,
[0, 1, 3, 4, 6]])
anchor_cache = {
"anchors": anchors,
"anchors_bv": anchors_bv,
"matched_thresholds": matched_thresholds,
"unmatched_thresholds": unmatched_thresholds,
}
self._prep_func = partial(prep_func, anchor_cache=anchor_cache)
def main(data_path):
# avm = ArgoverseMap()
# argoverse_tracker_loader = ArgoverseTrackingLoader('argoverse-tracking/') #simply change to your local path of the data
# argoverse_forecasting_loader = ArgoverseForecastingLoader('argoverse-forecasting/')
# argoVerseDataset = ArgoVerseDataset()
# argoVerseDataset.readCameraParameters()
argoverse_tracker_loader = ArgoverseTrackingLoader(data_path)
camera1_name = 'ring_front_center'
camera2_name = 'ring_front_right'
camera3_name = 'ring_front_left'
camera1_images = argoverse_tracker_loader.image_list[camera1_name]
camera2_images = argoverse_tracker_loader.image_list[camera2_name]
camera3_images = argoverse_tracker_loader.image_list[camera3_name]
camera1_calib = argoverse_tracker_loader.calib[camera1_name]
camera2_calib = argoverse_tracker_loader.calib[camera2_name]
camera3_calib = argoverse_tracker_loader.calib[camera3_name]
cv2.namedWindow('img1', cv2.WINDOW_NORMAL)
cv2.resizeWindow('img1', 600,400)
cv2.namedWindow('img2', cv2.WINDOW_NORMAL)
cv2.resizeWindow('img2', 600,400)
cv2.namedWindow('img3', cv2.WINDOW_NORMAL)
cv2.resizeWindow('img3', 600,400)
cv2.namedWindow('Stiched Img', cv2.WINDOW_NORMAL)
cv2.resizeWindow('Stiched Img', 1200,800)
for i in range(len(camera1_images)):
im1 = cv2.imread(camera1_images[i])
im2 = cv2.imread(camera2_images[i])
im3 = cv2.imread(camera3_images[i])
#Merge Imgs into single
sticker = cv2.Stitcher_create()
status, stichedImg = sticker.stitch([im1, im2, im3])
# cv2.imshow('img2', im2)
# cv2.waitKey()
if status == 0:
cv2.imshow('Stiched Img', stichedImg)
cv2.imshow('img1', im1)
cv2.imshow('img2', im2)
cv2.imshow('img3', im3)
# cv2.waitKey()
if cv2.waitKey(1) & 0xFF == ord('q'):
break
def __init__(self, root_dir, split='train'):
self.split = split
is_test = (self.split == 'test')
self.lidar_pathlist = []
self.label_pathlist = []
self.lidar_dir = os.path.join(root_dir)
data_loader = ArgoverseTrackingLoader(os.path.join(root_dir))
self.lidar_pathlist.extend(data_loader.lidar_list)
self.label_pathlist.extend(data_loader.label_list)
#self.calib_file = data_loader.calib_filename
self.lidar_filename = [x.split('.')[0].rsplit('/',1)[1] for x in self.lidar_pathlist]
assert len(self.lidar_pathlist) == len(self.label_pathlist)
self.num_sample = len(self.lidar_pathlist)
self.lidar_idx_list = ['%06d'%l for l in range(self.num_sample)]
self.lidar_idx_table = dict(zip(self.lidar_idx_list, self.lidar_filename))
self.argo_to_kitti = np.array([[0, -1, 0],
[0, 0, -1],
[1, 0, 0 ]])
self.ground_removal = False
self.lidar_dir = os.path.join('/data/')
self.label_dir = os.path.join('/data/')
def test_pose(data_loader: ArgoverseTrackingLoader) -> None:
for idx in range(len(data_loader.lidar_list)):
pose = data_loader.get_pose(idx)
assert np.array_equal(
pose.inverse().transform_point_cloud(np.array([[0, 0, 0]])),
pose.inverse_transform_point_cloud(np.array([[0, 0, 0]])),
)
def build_argoverse_datasets(config):
print('==> Loading Argoverse dataset...')
dataroot = os.path.expandvars(config.dataroot)
# Load native argoverse splits
loaders = {
'train': ArgoverseTrackingLoader(os.path.join(dataroot, 'train')),
'val': ArgoverseTrackingLoader(os.path.join(dataroot, 'val'))
}
# Create datasets using new argoverse splits
train_data = ArgoverseMapDataset(loaders, config.label_root,
config.img_size, TRAIN_LOGS)
val_data = ArgoverseMapDataset(loaders, config.label_root, config.img_size,
VAL_LOGS)
return train_data, val_data
def test_label_object(data_loader: ArgoverseTrackingLoader) -> None:
label_at_frame_0 = data_loader.get_label_object(0)
for label in label_at_frame_0:
assert label.label_class == "VEHICLE"
assert label.width == 2
assert label.height == 2
assert label.length == 2
def __init__(self, split, cfg,
phase,
transform=None,
random_rotation=True,
random_scale=True,
manual_seed=False,
config=None):
#super().__init__(split, cfg, )
self.phase = phase
self.files = []
self.data_objects = []
self.transform = transform
self.voxel_size = config.voxel_size
self.matching_search_voxel_size = \
config.voxel_size * config.positive_pair_search_voxel_size_multiplier
self.random_scale = random_scale
self.min_scale = config.min_scale
self.max_scale = config.max_scale
self.random_rotation = random_rotation
self.rotation_range = config.rotation_range
self.randg = np.random.RandomState()
if manual_seed:
self.reset_seed()
self.root = root = os.path.join(cfg.path_dataset, 'dataset')
self.split = phase
self.cfg = cfg
self.path_map_dict = os.path.join(
root, "argo_map_files_d3feat_%s_fcgf.pkl" % self.split)
# self.read_map_data()
# self.prepare_kitti_ply()#split=split)
from argoverse.data_loading.argoverse_tracking_loader import ArgoverseTrackingLoader
self.name = "argoverse_tracking"
self.cfg = cfg
self.split = split
self.transform_map = True
if split == "train":
self.list_id_log = np.load(os.path.join(
self.cfg.path_dataset, 'argoverse-tracking', 'train_ids.npy'))
list_folders = ["train1", "train2", "train3", "train4"]
elif split == "test":
self.list_id_log = np.load(os.path.join(
self.cfg.path_dataset, 'argoverse-tracking', 'test_ids.npy'))
list_folders = ["test"]
self.dict_tracking_data_loader = {}
for name_folder in list_folders:
self.dict_tracking_data_loader[name_folder] = ArgoverseTrackingLoader(
os.path.join(self.cfg.path_dataset, 'argoverse-tracking', name_folder))
self.transform_map = False
self.read_data()
def __init__(self, args):
# Setup Argoverse loader
self.dataset_dir = args.dataset_dir
self.log_id = args.log_id
self.argoverse_loader = ArgoverseTrackingLoader(self.dataset_dir)
self.argoverse_data = self.argoverse_loader.get(self.log_id)
print(self.argoverse_data)
# List of cameras to publish
self.cameras_list = [camera for camera in self.argoverse_loader.CAMERA_LIST
if camera in args.cameras]
self.load_camera_info()
# Images timestamps
self.image_files_dict = {
camera: self.argoverse_data.timestamp_image_dict[camera]
for camera in self.cameras_list}
self.image_timestamps = {
camera: list(self.image_files_dict[camera].keys())
for camera in self.cameras_list}
# Save image files dict to be added into the bag file later using Python2
with open('/tmp/image_data_%s.json' % self.log_id, 'w') as f:
data = {'image_files_dict': self.image_files_dict,
'image_timestamps': self.image_timestamps,
'cameras_list': self.cameras_list}
json.dump(data, f, indent=4)
# LiDAR timestamps
self.lidar_files_dict = self.argoverse_data.timestamp_lidar_dict
self.lidar_timestamps = list(self.lidar_files_dict.keys())
# ROSBAG output path
if not os.path.isdir(args.output_dir): os.makedirs(args.output_dir)
self.output_filename = os.path.join(args.output_dir, '%s_no_images.bag' % self.log_id)
self.bag = rosbag.Bag(self.output_filename, 'w')
# Topic names
self.lidar_topic = '/argoverse/lidar/pointcloud'
self.camera_info_topic_template = '/argoverse/%s/camera_info'
# TF frames
self.map_frame = 'city'
self.vehicle_frame = 'egovehicle'
self.load_camera_static_tf()
def process_split(split, map_data, config):
# Create an Argoverse loader instance
path = os.path.join(os.path.expandvars(config.argoverse.root), split)
print("Loading Argoverse tracking data at " + path)
loader = ArgoverseTrackingLoader(path)
for scene in loader:
process_scene(split, scene, map_data, config)
def __init__(self, root_path, subsets: list):
super().__init__()
self.root_path = root_path
self.atls = {
subset: ArgoverseTrackingLoader(Path(self.root_path) / subset)
for subset in subsets
}
self._len = 0
pass
def __init__(self, question_h5, image_feature_h5_path, lidar_feature_h5_path,vocab,load_lidar=True,npoint=1024,normal_channel=True,uniform=False,cache_size=15000,drivable_area=False,
mode='prefix', image_h5=None, lidar_h5=None, max_samples=None, question_families=None,
image_idx_start_from=None):
#############read whole question_h5 file in memory#############################
self.all_questions = question_h5['questions'][:]
self.all_answers = get_answer_classes(question_h5['answers'][:], vocab)
self.all_image_idxs = question_h5['image_idxs'][:]
self.all_video_names = (question_h5['video_names'][:]).astype(str)
self.questions_length = question_h5['question_length'][:]
self.image_feature_h5 = image_feature_h5_path
self.load_lidar=load_lidar
############for lidar##########################################################
if self.load_lidar:
self.argoverse_loader = ArgoverseTrackingLoader(lidar_feature_h5_path)
self.am = ArgoverseMap()
self.drivable_area=drivable_area
def __init__(self, root_dir, split='train'):
self.split = split
is_test = (self.split == 'test')
self.lidar_pathlist = []
self.label_pathlist = []
if split == 'train':
for i in np.arange(1, 5):
self.imageset_dir = os.path.join(root_dir, split + str(i))
data_loader = ArgoverseTrackingLoader(
os.path.join(root_dir, split + str(i)))
self.log_list = data_loader.log_list
for log in self.log_list:
self.lidar_pathlist.extend(data_loader.get(log).lidar_list)
self.label_pathlist.extend(data_loader.get(log).label_list)
print(len(self.lidar_pathlist))
else:
self.imageset_dir = os.path.join(root_dir, split)
data_loader = ArgoverseTrackingLoader(os.path.join(
root_dir, split))
self.lidar_pathlist.extend(data_loader.lidar_list)
self.label_pathlist.extend(data_loader.label_list)
self.calib_file = data_loader.calib_filename
assert len(self.lidar_pathlist) == len(self.label_pathlist)
#z = list(zip(self.lidar_pathlist, self.label_pathlist))
#random.shuffle(z)
#self.lidar_pathlist[:], self.label_pathlist[:] = zip(*z)
self.num_sample = len(self.lidar_pathlist)
self.image_idx_list = np.arange(self.num_sample)
self.argo_to_kitti = np.array(
[[6.927964e-03, -9.999722e-01, -2.757829e-03],
[-1.162982e-03, 2.749836e-03, -9.999955e-01],
[9.999753e-01, 6.931141e-03, -1.143899e-03]])
self.ground_removal = True
self.image_dir = os.path.join('/data/')
self.lidar_dir = os.path.join('/data/')
self.calib_dir = os.path.join('/data/')
self.label_dir = os.path.join('/data/')
def __init__(self, split, cfg, config_d3feat, root,
transform=None,
random_rotation=True,
random_scale=True,
manual_seed=False,
config=None, input_threads=8,
first_subsampling_dl=0.30):
#super().__init__(split, cfg, )
self.network_model = 'descriptor'
self.num_threads = input_threads
self.root = root
self.icp_path = "data" # 'data/kitti/icp'
self.voxel_size = first_subsampling_dl
self.matching_search_voxel_size = first_subsampling_dl * 0.5 # 1.5
self.split = split
# Initiate containers
#self.anc_points = {'train': [], 'val': [], 'test': []}
#self.files = {'train': [], 'val': [], 'test': []}
self.config = config_d3feat
self.path_map_dict = os.path.join(
root, "argo_map_files_d3feat_%s.pkl" % self.split)
# to match the dataset file used in the benchmark
from argoverse.data_loading.argoverse_tracking_loader import ArgoverseTrackingLoader
self.name = "argoverse_tracking"
self.cfg = cfg
#self.split = split
self.transform_map = True
if split == "train":
self.list_id_log = np.load(os.path.join(
self.cfg.path_dataset, 'argoverse-tracking', 'train_ids.npy'))
list_folders = ["train1", "train2", "train3", "train4"]
elif split == "test":
self.list_id_log = np.load(os.path.join(
self.cfg.path_dataset, 'argoverse-tracking', 'test_ids.npy'))
list_folders = ["test"]
self.dict_tracking_data_loader = {}
for name_folder in list_folders:
self.dict_tracking_data_loader[name_folder] = ArgoverseTrackingLoader(
os.path.join(self.cfg.path_dataset, 'argoverse-tracking', name_folder))
#self.root = root = os.path.join(cfg.path_dataset, 'dataset')
# self.read_map_data()
# self.prepare_kitti_ply()#split=split)
self.transform_map=False
self.read_data()
def __init__(self, input_log_dir: str, output_save_path: str) -> None:
self.input_log_dir = input_log_dir
self.output_save_path = output_save_path
self.log_id = os.path.basename(input_log_dir)
print("Log ID ", self.log_id)
# Load Argo data
dataset = os.path.dirname(self.input_log_dir)
self.argoverse_loader = ArgoverseTrackingLoader(dataset)
self.argoverse_data = self.argoverse_loader.get(self.log_id)
# Count the number of LiDAR ply files in the log dir
self.lidar_frame_counter = len(
glob.glob1(os.path.join(self.input_log_dir, "lidar"), "*.ply"))
# Setup depth dataset dir
self.depth_data_dir_setup()
# Extract depth data and ring camera frames
self.depth_extraction()
def __init__(self, dataset_path, out_path, version='sample'):
assert version in ['train', 'val', 'test', 'sample']
self.is_test = 'test' in version
self.out_path = out_path
self.dataset_path = join(dataset_path, version)
self.argo = ArgoverseTrackingLoader(self.dataset_path)
print(f"Total number of logs : {len(self.argo)}, version : {version}")
self.version = version
class ArgoDataset(Dataset):
def __init__(self, args, transform=None):
self.args = args
self.transform = transform
self.argoverse_loader = ArgoverseTrackingLoader(args.root_dir)
self.image_list = []
for log_id in self.argoverse_loader.log_list:
argoverse_data = self.argoverse_loader.get(log_id)
for idx in range(
len(argoverse_data.image_list_sync['ring_front_center'])):
self.image_list.append(str(log_id) + '@' + str(idx))
self.camera_name = [
'ring_front_center', 'ring_side_left', 'ring_side_right',
'ring_front_left', 'ring_front_right', 'ring_rear_right',
'ring_rear_left'
]
def __len__(self):
return len(self.image_list)
def __getitem__(self, index):
lidar_index = int(str(self.image_list[index]).split('@')[1])
log_id = str(self.image_list[index]).split('@')[0]
argoverse_data = self.argoverse_loader.get(log_id)
image_7 = []
for camera in self.camera_name:
img = argoverse_data.get_image_sync(lidar_index, camera=camera)
img = Image.fromarray(img)
img = self.transform(img) if self.transform is not None else img
image_7.append(img)
sample = {
'image_7': image_7,
'video_idxs': lidar_index,
'image_names': log_id
}
return sample
def __init__(self,
tracking_dataset_dir,
dataset_name=None,
argoverse_map=None,
argoverse_loader=None,
save_imgs=False):
logger = logging.getLogger()
logger.setLevel(logging.CRITICAL)
self.dataset_dir = tracking_dataset_dir
self.am = ArgoverseMap() if argoverse_map is None else argoverse_map
self.argoverse_loader = ArgoverseTrackingLoader(
tracking_dataset_dir
) if argoverse_loader is None else argoverse_loader
self.dataset_prefix_name = dataset_name
self.objects_from_to = self._get_objects_from_to()
self.valid_target_objects = self._get_valid_target_objects(
save_imgs=save_imgs)
def save_all_to_pickle():
datasets = ["train1", "train2", "train3", "train4"]
final_dict = {}
for dataset in datasets:
tracking_dataset_dir = '/media/bartosz/hdd1TB/workspace_hdd/datasets/argodataset/argoverse-tracking/' + dataset
###################
am = ArgoverseMap()
argoverse_loader = ArgoverseTrackingLoader(tracking_dataset_dir)
###################
argoverse = Argoverse(tracking_dataset_dir=tracking_dataset_dir,
dataset_name=dataset,
argoverse_map=am,
argoverse_loader=argoverse_loader)
final_dict.update(argoverse.valid_target_objects)
print("Processed {}".format(dataset))
f = "/media/bartosz/hdd1TB/workspace_hdd/SS-LSTM/data/argoverse/train1234_48x48.pickle"
pickle_out = open(f, "wb")
pickle.dump(final_dict, pickle_out, protocol=2)
pickle_out.close()
print("Saved to pickle {}".format(f))
def __init__(self,
root,
seed=None,
train=True,
sequence_length=3,
transform=None,
skip_frames=1,
dataset='kitti'):
np.random.seed(seed)
random.seed(seed)
self.root = Path(root)
scene_list_path = self.root / 'train.txt' if train else self.root / 'val.txt'
self.scenes = [
self.root / folder[:-1] for folder in open(scene_list_path)
]
self.transform = transform
self.dataset = dataset
self.k = skip_frames
self.root_dir = '/ssd_scratch/cvit/raghava.modhugu/'
self.argoverse_loader = ArgoverseTrackingLoader(self.root_dir)
self.camera = self.argoverse_loader.CAMERA_LIST[0]
self.crawl_folders(sequence_length, self.camera, self.argoverse_loader)
def draw_point_cloud_trajectory(
ax: plt.Axes,
title: str,
argoverse_data: ArgoverseTrackingLoader,
idx: int,
axes: Optional[Any] = None,
xlim3d: Any = None,
ylim3d: Any = None,
zlim3d: Any = None,
) -> None:
axes = _get_axes_or_default(axes)
unique_id_list = set()
for i in range(len(argoverse_data.lidar_list)):
for label in argoverse_data.get_label_object(i):
unique_id_list.add(label.track_id)
color_map = {
track_id: (
float(np.random.rand()),
float(np.random.rand()),
float(np.random.rand()),
)
for track_id in unique_id_list
}
pc = argoverse_data.get_lidar(idx)
assert isinstance(pc, np.ndarray)
objects = argoverse_data.get_label_object(idx)
ax.scatter(*np.transpose(pc[:, axes]),
s=point_size,
c=pc[:, 2],
cmap="gray")
ax.set_title(title)
ax.set_xlabel("{} axis".format(axes_str[axes[0]]))
ax.set_ylabel("{} axis".format(axes_str[axes[1]]))
if len(axes) > 2:
ax.set_xlim3d(*axes_limits[axes[0]])
ax.set_ylim3d(*axes_limits[axes[1]])
ax.set_zlim3d(*axes_limits[axes[2]])
ax.set_zlabel("{} axis".format(axes_str[axes[2]]))
else:
ax.set_xlim(*axes_limits[axes[0]])
ax.set_ylim(*axes_limits[axes[1]])
# User specified limits
if xlim3d != None:
ax.set_xlim3d(xlim3d)
if ylim3d != None:
ax.set_ylim3d(ylim3d)
if zlim3d != None:
ax.set_zlim3d(zlim3d)
visible_track_id = set()
for obj in objects:
if obj.occlusion == 100:
continue
draw_box(ax,
obj.as_3d_bbox().T,
axes=axes,
color=color_map[obj.track_id])
visible_track_id.add(obj.track_id)
current_pose = argoverse_data.get_pose(idx)
traj_by_id: Dict[Optional[str], List[Any]] = defaultdict(list)
for i in range(0, idx, 1):
if current_pose is None:
logger.warning("`current_pose` is missing at index %d", idx)
break
pose = argoverse_data.get_pose(i)
if pose is None:
logger.warning("`pose` is missing at index %d", i)
continue
objects = argoverse_data.get_label_object(i)
for obj in objects:
if obj.occlusion == 100:
continue
if obj.track_id is None or obj.track_id not in visible_track_id:
continue
x, y, z = pose.transform_point_cloud(
np.array([np.array(obj.translation)]))[0]
x, y, _ = current_pose.inverse_transform_point_cloud(
np.array([np.array([x, y, z])]))[0]
# ax.scatter(x,y, s=point_size, c=color_map[obj.track_id])
if obj.track_id is None:
logger.warning(
"Label has no track_id. Collisions with other tracks that are missing IDs could happen"
)
traj_by_id[obj.track_id].append([x, y])
for track_id in traj_by_id.keys():
traj = np.array(traj_by_id[track_id])
ax.plot(
traj[:, 0],
traj[:, 1],
color=color_map[track_id],
linestyle="--",
linewidth=1,
)
def make_grid_ring_camera(argoverse_data: ArgoverseTrackingLoader,
idx: int) -> Tuple[plt.Figure, plt.Axes]:
f, ax = plt.subplots(3, 3, figsize=(20, 15))
camera = "ring_front_left"
img = argoverse_data.get_image_sync(idx, camera=camera)
objects = argoverse_data.get_label_object(idx)
calib = argoverse_data.get_calibration(camera)
img_vis = Image.fromarray(show_image_with_boxes(img, objects, calib))
ax[0, 0].imshow(img_vis)
ax[0, 0].set_title("Ring Front Left")
ax[0, 0].axis("off")
camera = "ring_front_center"
img = argoverse_data.get_image_sync(idx, camera=camera)
objects = argoverse_data.get_label_object(idx)
calib = argoverse_data.get_calibration(camera)
img_vis = Image.fromarray(show_image_with_boxes(img, objects, calib))
ax[0, 1].imshow(img_vis)
ax[0, 1].set_title("Right Front Center")
ax[0, 1].axis("off")
camera = "ring_front_right"
img = argoverse_data.get_image_sync(idx, camera=camera)
objects = argoverse_data.get_label_object(idx)
calib = argoverse_data.get_calibration(camera)
img_vis = Image.fromarray(show_image_with_boxes(img, objects, calib))
ax[0, 2].imshow(img_vis)
ax[0, 2].set_title("Ring Front Right")
ax[0, 2].axis("off")
camera = "ring_side_left"
img = argoverse_data.get_image_sync(idx, camera=camera)
objects = argoverse_data.get_label_object(idx)
calib = argoverse_data.get_calibration(camera)
img_vis = Image.fromarray(show_image_with_boxes(img, objects, calib))
ax[1, 0].imshow(img_vis)
ax[1, 0].set_title("Ring Side Left")
ax[1, 0].axis("off")
ax[1, 1].axis("off")
camera = "ring_side_right"
img = argoverse_data.get_image_sync(idx, camera=camera)
objects = argoverse_data.get_label_object(idx)
calib = argoverse_data.get_calibration(camera)
img_vis = Image.fromarray(show_image_with_boxes(img, objects, calib))
ax[1, 2].imshow(img_vis)
ax[1, 2].set_title("Ring Side Right")
ax[1, 2].axis("off")
camera = "ring_rear_left"
img = argoverse_data.get_image_sync(idx, camera=camera)
objects = argoverse_data.get_label_object(idx)
calib = argoverse_data.get_calibration(camera)
img_vis = Image.fromarray(show_image_with_boxes(img, objects, calib))
ax[2, 0].imshow(img_vis)
ax[2, 0].set_title("Ring Rear Left")
ax[2, 0].axis("off")
ax[2, 1].axis("off")
camera = "ring_rear_right"
img = argoverse_data.get_image_sync(idx, camera=camera)
objects = argoverse_data.get_label_object(idx)
calib = argoverse_data.get_calibration(camera)
img_vis = Image.fromarray(show_image_with_boxes(img, objects, calib))
ax[2, 2].imshow(img_vis)
ax[2, 2].set_title("Ring Rear Right")
ax[2, 2].axis("off")
return f, ax
def generate_vehicle_bev(dataset_dir="", log_id="", output_dir=""):
argoverse_loader = ArgoverseTrackingLoader(dataset_dir)
argoverse_data = argoverse_loader.get(log_id)
camera = argoverse_loader.CAMERA_LIST[7]
calib = argoverse_data.get_calibration(camera)
sdb = SynchronizationDB(dataset_dir, collect_single_log_id=log_id)
calib_path = f"{dataset_dir}/{log_id}/vehicle_calibration_info.json"
calib_data = read_json_file(calib_path)
ind = 0
for i in range(9):
if calib_data['camera_data_'][i]['key'] == \
'image_raw_stereo_front_left':
ind = i
break
rotation = np.array(calib_data['camera_data_'][ind]['value']
['vehicle_SE3_camera_']['rotation']['coefficients'])
translation = np.array(calib_data['camera_data_'][ind]['value']
['vehicle_SE3_camera_']['translation'])
egovehicle_SE3_cam = SE3(rotation=quat2rotmat(rotation),
translation=translation)
if not os.path.exists(os.path.join(output_dir, log_id, "car_bev_gt")):
os.makedirs(os.path.join(output_dir, log_id, "car_bev_gt"))
lidar_dir = os.path.join(dataset_dir, log_id, "lidar")
ply_list = os.listdir(lidar_dir)
pfa = PerFrameLabelAccumulator(dataset_dir,
dataset_dir,
"argoverse_bev_viz",
save=False,
bboxes_3d=True)
pfa.accumulate_per_log_data(log_id)
log_timestamp_dict = pfa.log_timestamp_dict
for i, ply_name in enumerate(ply_list):
lidar_timestamp = ply_name.split('.')[0].split('_')[1]
lidar_timestamp = int(lidar_timestamp)
cam_timestamp = sdb.get_closest_cam_channel_timestamp(
lidar_timestamp, "stereo_front_left", str(log_id))
image_path = os.path.join(
output_dir, str(log_id), "car_bev_gt",
"stereo_front_left_" + str(cam_timestamp) + ".jpg")
objects = log_timestamp_dict[log_id][lidar_timestamp]
top_view = np.zeros((256, 256))
all_occluded = True
for frame_rec in objects:
if frame_rec.occlusion_val != IS_OCCLUDED_FLAG:
all_occluded = False
if not all_occluded:
for i, frame_rec in enumerate(objects):
bbox_ego_frame = frame_rec.bbox_ego_frame
uv = calib.project_ego_to_image(bbox_ego_frame).T
idx_ = np.all(
np.logical_and(
np.logical_and(
np.logical_and(uv[0, :] >= 0.0,
uv[0, :] < size[1] - 1.0),
np.logical_and(uv[1, :] >= 0.0,
uv[1, :] < size[0] - 1.0)),
uv[2, :] > 0))
if not idx_:
continue
bbox_cam_fr = egovehicle_SE3_cam.inverse().\
transform_point_cloud(bbox_ego_frame)
X = bbox_cam_fr[:, 0]
Z = bbox_cam_fr[:, 2]
if (frame_rec.occlusion_val != IS_OCCLUDED_FLAG
and frame_rec.obj_class_str == "VEHICLE"):
y_img = (-Z / res).astype(np.int32)
x_img = (X / res).astype(np.int32)
x_img -= int(np.floor(-20 / res))
y_img += int(np.floor(40 / res))
box = np.array([x_img[2], y_img[2]])
box = np.vstack((box, [x_img[6], y_img[6]]))
box = np.vstack((box, [x_img[7], y_img[7]]))
box = np.vstack((box, [x_img[3], y_img[3]]))
cv2.drawContours(top_view, [box], 0, 255, -1)
cv2.imwrite(image_path, top_view)
def generate_road_bev(dataset_dir="", log_id="", output_dir=""):
argoverse_loader = ArgoverseTrackingLoader(dataset_dir)
argoverse_data = argoverse_loader.get(log_id)
city_name = argoverse_data.city_name
avm = ArgoverseMap()
sdb = SynchronizationDB(dataset_dir, collect_single_log_id=log_id)
try:
path = os.path.join(output_dir, log_id, 'road_gt')
command = "rm -r " + path
os.system(command)
except BaseException:
pass
try:
os.makedirs(os.path.join(output_dir, log_id, 'road_gt'))
except BaseException:
pass
ply_fpath = os.path.join(dataset_dir, log_id, 'lidar')
ply_locs = []
for idx, ply_name in enumerate(os.listdir(ply_fpath)):
ply_loc = np.array([idx, int(ply_name.split('.')[0].split('_')[-1])])
ply_locs.append(ply_loc)
ply_locs = np.array(ply_locs)
lidar_timestamps = sorted(ply_locs[:, 1])
calib_path = f"{dataset_dir}/{log_id}/vehicle_calibration_info.json"
calib_data = read_json_file(calib_path)
ind = 0
for i in range(9):
if calib_data['camera_data_'][i]['key'] ==\
'image_raw_stereo_front_left':
ind = i
break
rotation = np.array(calib_data['camera_data_'][ind]['value']
['vehicle_SE3_camera_']['rotation']['coefficients'])
translation = np.array(calib_data['camera_data_'][ind]['value']
['vehicle_SE3_camera_']['translation'])
egovehicle_SE3_cam = SE3(rotation=quat2rotmat(rotation),
translation=translation)
for idx in range(len(lidar_timestamps)):
lidar_timestamp = lidar_timestamps[idx]
cam_timestamp = sdb.get_closest_cam_channel_timestamp(
lidar_timestamp, "stereo_front_left", str(log_id))
occupancy_map = np.zeros((256, 256))
pose_fpath = os.path.join(
dataset_dir, log_id, "poses",
"city_SE3_egovehicle_" + str(lidar_timestamp) + ".json")
if not Path(pose_fpath).exists():
continue
pose_data = read_json_file(pose_fpath)
rotation = np.array(pose_data["rotation"])
translation = np.array(pose_data["translation"])
xcenter = translation[0]
ycenter = translation[1]
city_to_egovehicle_se3 = SE3(rotation=quat2rotmat(rotation),
translation=translation)
ego_car_nearby_lane_ids = avm.get_lane_ids_in_xy_bbox(
xcenter, ycenter, city_name, 50.0)
occupancy_map = get_lane_bev(ego_car_nearby_lane_ids, avm, city_name,
city_to_egovehicle_se3,
egovehicle_SE3_cam, occupancy_map, res,
255)
output_loc = os.path.join(
output_dir, log_id, 'road_gt',
'stereo_front_left_' + str(cam_timestamp) + '.png')
cv2.imwrite(output_loc, occupancy_map)
