Added motion estimation with outputs
Added Recorder pose outputs to compare calculations Added dataset parsing library
This commit is contained in:
1
data_parser/__init__.py
Normal file
1
data_parser/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
from data_parser.data_parser import create_tables, Dataset, Timestamp, Camera, CameraView
|
||||
409
data_parser/data_parser.py
Normal file
409
data_parser/data_parser.py
Normal file
@@ -0,0 +1,409 @@
|
||||
import json
|
||||
from os import PathLike
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
from scipy.spatial.transform import RigidTransform, Rotation
|
||||
|
||||
from sqlalchemy import Column, Engine, Float, Integer, Enum, ForeignKey, String
|
||||
from sqlalchemy.orm import DeclarativeBase, Session, relationship
|
||||
|
||||
import re
|
||||
|
||||
class ModelBase(DeclarativeBase):
|
||||
pass
|
||||
|
||||
metadata = ModelBase.metadata
|
||||
|
||||
def create_tables(engine:Engine):
|
||||
metadata.create_all(engine)
|
||||
|
||||
ts_pose_re = re.compile(r"city_SE3_egovehicle_(?P<timestamp>\d*)")
|
||||
cam_name_re = re.compile(r"image_raw_(?P<name>.*)")
|
||||
camview_fn_re = re.compile(r"(?P<cam_name>[^\d]*)_(?P<timestamp>\d*)")
|
||||
|
||||
|
||||
def json_to_pose(data: dict) -> RigidTransform:
|
||||
""" Creates a pose from an Argoverse pose json
|
||||
|
||||
Args:
|
||||
data (dict): json contents
|
||||
|
||||
Returns:
|
||||
RigidTransform: pose in the json data
|
||||
"""
|
||||
# TODO improve json validation
|
||||
|
||||
if 'rotation' in data:
|
||||
if 'coefficients' in data['rotation']:
|
||||
rot_data = data['rotation']['coefficients']
|
||||
else:
|
||||
rot_data = data['rotation']
|
||||
else:
|
||||
raise RuntimeError('No rotation present in pose data')
|
||||
|
||||
if 'translation' in data:
|
||||
trans_data = data['translation']
|
||||
else:
|
||||
raise RuntimeError('No translation present in pose data')
|
||||
|
||||
|
||||
return RigidTransform.from_components(
|
||||
np.array(trans_data),
|
||||
Rotation.from_quat(rot_data))
|
||||
|
||||
|
||||
class Dataset(ModelBase):
|
||||
""" Represents a dataset """
|
||||
__tablename__ = "datasets"
|
||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
root_path = Column(String(256))
|
||||
key = Column(String(36))
|
||||
city = Column(String(3), default="")
|
||||
|
||||
timestamps = relationship("Timestamp", back_populates='dataset')
|
||||
cameras = relationship("Camera", back_populates='dataset')
|
||||
camera_views = relationship("CameraView", back_populates='dataset')
|
||||
|
||||
def __init__(self, root_path: PathLike, key: str, city: str = ''):
|
||||
""" Constructor
|
||||
|
||||
Args:
|
||||
key (str): key of the dataset
|
||||
root_path (PathLike): Root of the data directory
|
||||
"""
|
||||
self.key = key
|
||||
self.city = city if len(city) <= 3 else city[0:3]
|
||||
self.root_path = str(root_path)
|
||||
|
||||
@staticmethod
|
||||
def import_dataset(root_path: PathLike, session:Session) -> 'Dataset':
|
||||
""" Imports a dataset
|
||||
|
||||
Args:
|
||||
root_path (PathLike): root path to the dataset
|
||||
session (Session): SQLAlchemy session to use
|
||||
|
||||
Returns:
|
||||
Dataset: Reference to the dataset
|
||||
"""
|
||||
|
||||
root_path = Path(root_path)
|
||||
key = root_path.stem
|
||||
city = ''
|
||||
|
||||
city_info_path = root_path / "city_info.json"
|
||||
calib_path = root_path / "vehicle_calibration_info.json"
|
||||
pose_root = root_path / "poses"
|
||||
|
||||
# Get City Information
|
||||
if city_info_path.exists():
|
||||
with open(city_info_path, 'r') as f:
|
||||
data = json.load(f)
|
||||
city = data['city'] if 'city' in data else ''
|
||||
|
||||
# Create the new dataset
|
||||
dataset = Dataset(root_path, key, city)
|
||||
|
||||
# Create timestamps from poses
|
||||
for path in pose_root.glob("*.json"):
|
||||
# TODO Validate json
|
||||
timestamp = Timestamp.from_json(path, dataset, session)
|
||||
|
||||
# Create Cameras
|
||||
with open(calib_path, 'r') as f:
|
||||
data = json.load(f)
|
||||
|
||||
# TODO Validate file
|
||||
cameras = []
|
||||
for camera_data in data['camera_data_']:
|
||||
camera = Camera.from_json(camera_data, dataset, session)
|
||||
|
||||
if camera is not None:
|
||||
cameras.append(camera)
|
||||
|
||||
# Create Camera Views
|
||||
for camera in cameras:
|
||||
camera_dir = root_path / camera.name
|
||||
|
||||
for path in camera_dir.glob("*.*"):
|
||||
CameraView.from_path(path, dataset, session)
|
||||
|
||||
return dataset
|
||||
|
||||
|
||||
class Timestamp(ModelBase):
|
||||
""" Correlates data at a given timestamp """
|
||||
__tablename__ = "timestamps"
|
||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
time = Column(Integer)
|
||||
|
||||
vehicle_pose_x = Column(Float) # x component
|
||||
vehicle_pose_y = Column(Float) # y component
|
||||
vehicle_pose_z = Column(Float) # z component
|
||||
vehicle_pose_q0 = Column(Float) # Quaternion component 0
|
||||
vehicle_pose_q1 = Column(Float) # Quaternion component 1
|
||||
vehicle_pose_q2 = Column(Float) # Quaternion component 2
|
||||
vehicle_pose_q3 = Column(Float) # Quaternion component 3
|
||||
|
||||
dataset_id = Column(Integer, ForeignKey('datasets.id'), nullable=False)
|
||||
dataset = relationship('Dataset', back_populates='timestamps')
|
||||
|
||||
camera_views = relationship("CameraView", back_populates="timestamp")
|
||||
|
||||
def __init__(self, time: int, vehicle_pose: RigidTransform, dataset: Dataset):
|
||||
""" Constructor
|
||||
|
||||
Args:
|
||||
time (int): nanosecond time for the timestamp
|
||||
vehicle_pose (RigidTransform): vehicle pose in world coordinates
|
||||
dataset (Dataset): dataset this timestamp belongs to
|
||||
|
||||
"""
|
||||
|
||||
self.time = time
|
||||
|
||||
self.set_vehicle_pose(vehicle_pose)
|
||||
|
||||
self.dataset = dataset
|
||||
|
||||
def set_vehicle_pose(self, vehicle_pose: RigidTransform):
|
||||
""" Sets the vehicle's pose
|
||||
|
||||
Args:
|
||||
vehicle_pose (RigidTransform): vehicle pose in world coordinates
|
||||
"""
|
||||
t, R = vehicle_pose.as_components()
|
||||
quat = R.as_quat()
|
||||
|
||||
self.vehicle_pose_x = t[0]
|
||||
self.vehicle_pose_y = t[1]
|
||||
self.vehicle_pose_z = t[2]
|
||||
self.vehicle_pose_q0 = quat[0]
|
||||
self.vehicle_pose_q1 = quat[0]
|
||||
self.vehicle_pose_q2 = quat[0]
|
||||
self.vehicle_pose_q3 = quat[0]
|
||||
|
||||
def get_vehicle_pose(self) -> RigidTransform:
|
||||
""" Returns the vehicle pose as a RigidTransform object """
|
||||
return RigidTransform.from_components(
|
||||
translation=np.array(
|
||||
[self.vehicle_pose_x, self.vehicle_pose_y, self.vehicle_pose_z]),
|
||||
rotation=Rotation.from_quat([self.vehicle_pose_q0, self.vehicle_pose_q1, self.vehicle_pose_q2, self.vehicle_pose_q3]))
|
||||
|
||||
@staticmethod
|
||||
def from_json(path: PathLike, dataset: Dataset, session:Session) -> Optional["Timestamp"]:
|
||||
""" Creates a timestamp object from a pose file path and dataset object
|
||||
|
||||
Args:
|
||||
path (PathLike): Path to pose file
|
||||
dataset (Dataset): Dataset object to which this timestamp
|
||||
session (Session): SQLAlchemy session to use
|
||||
|
||||
Returns:
|
||||
Optional[Timestamp]: new timestamp object if it is successfully created, None otherwise
|
||||
"""
|
||||
result = None
|
||||
path = Path(path)
|
||||
ts_pose_match = ts_pose_re.match(str(path.stem))
|
||||
if ts_pose_match is not None:
|
||||
time = int(ts_pose_match.group("timestamp"))
|
||||
|
||||
with open(path, 'r') as f:
|
||||
# TODO Add exception handling
|
||||
result = Timestamp(time, json_to_pose(json.load(f)), dataset)
|
||||
session.add(result)
|
||||
session.commit()
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class Camera(ModelBase):
|
||||
""" Represents a camera in the dataset """
|
||||
__tablename__ = "cameras"
|
||||
|
||||
id = Column(Integer, primary_key=True)
|
||||
name = Column(String(20), nullable=False)
|
||||
|
||||
fx = Column(Float, nullable=False)
|
||||
fy = Column(Float, nullable=False)
|
||||
cx = Column(Float, nullable=False)
|
||||
cy = Column(Float, nullable=False)
|
||||
|
||||
extrinsics_x = Column(Float) # x component
|
||||
extrinsics_y = Column(Float) # y component
|
||||
extrinsics_z = Column(Float) # z component
|
||||
extrinsics_q0 = Column(Float) # Quaternion component 0
|
||||
extrinsics_q1 = Column(Float) # Quaternion component 1
|
||||
extrinsics_q2 = Column(Float) # Quaternion component 2
|
||||
extrinsics_q3 = Column(Float) # Quaternion component 3
|
||||
|
||||
dataset_id = Column(Integer, ForeignKey('datasets.id'), nullable=False)
|
||||
dataset = relationship('Dataset', back_populates='cameras')
|
||||
|
||||
camera_views = relationship("CameraView", back_populates="camera")
|
||||
|
||||
def __init__(self, name: str, intrinsics: np.ndarray, extrinsics: RigidTransform, dataset: Dataset):
|
||||
""" Constructor
|
||||
|
||||
Args:
|
||||
name (str): Name of the camera
|
||||
intrinsics (np.ndarray): Camera intrinsics matrix
|
||||
extrinsics (RigidTransform): Camera extrinsics transform
|
||||
dataset (Dataset): dataset this camera belongs to
|
||||
"""
|
||||
self.name = name
|
||||
self.set_intrinsics(intrinsics)
|
||||
self.set_extrinsics(extrinsics)
|
||||
self.dataset = dataset
|
||||
|
||||
def set_intrinsics(self, intrinsics: np.ndarray):
|
||||
""" Sets the camera intrinsics matrix
|
||||
|
||||
Args:
|
||||
intrinsics (np.ndarray): Camera intrinsics matrix
|
||||
"""
|
||||
self.fx = intrinsics[0, 0]
|
||||
self.fy = intrinsics[1, 1]
|
||||
self.cx = intrinsics[0, 2]
|
||||
self.cy = intrinsics[1, 2]
|
||||
|
||||
def set_extrinsics(self, extrinsics: RigidTransform):
|
||||
""" Sets the camera extrinsics
|
||||
|
||||
Args:
|
||||
extrinsics (RigidTransform): camera extrinsics in vehicle coordinates
|
||||
"""
|
||||
t, R = extrinsics.as_components()
|
||||
quat = R.as_quat()
|
||||
|
||||
self.extrinsics_x = t[0]
|
||||
self.extrinsics_y = t[1]
|
||||
self.extrinsics_z = t[2]
|
||||
self.extrinsics_q0 = quat[0]
|
||||
self.extrinsics_q0 = quat[0]
|
||||
self.extrinsics_q0 = quat[0]
|
||||
self.extrinsics_q0 = quat[0]
|
||||
|
||||
def get_intrinsics(self) -> np.ndarray:
|
||||
""" Gets the camera intrinsics matrix
|
||||
|
||||
Returns:
|
||||
np.ndarray: 3 x 3 camera intrinsics matrix
|
||||
|
||||
"""
|
||||
|
||||
return np.array([[self.fx, 0, self.cx],
|
||||
[0, self.fy, self.cy],
|
||||
[0, 0, 1]])
|
||||
|
||||
def get_extrinsics(self) -> RigidTransform:
|
||||
""" Returns the extrinsics as a RigidTransform object """
|
||||
return RigidTransform.from_components(
|
||||
translation=np.array(
|
||||
[self.extrinsics_x, self.extrinsics_y, self.extrinsics_z]),
|
||||
rotation=Rotation.from_quat([self.extrinsics_q0, self.extrinsics_q1, self.extrinsics_q2, self.extrinsics_q3]))
|
||||
|
||||
@staticmethod
|
||||
def from_json(data: dict, dataset: Dataset, session:Session) -> Optional["Camera"]:
|
||||
""" Creates a camera from json data
|
||||
|
||||
Args:
|
||||
data (dict): camera data dict
|
||||
dataset (Dataset): dataset object to which the camera belongs
|
||||
session (Session): SQLAlchemy session to use
|
||||
|
||||
Returns:
|
||||
Camera: Camera object if the data is valid, None otherwise
|
||||
|
||||
"""
|
||||
result = None
|
||||
# TODO Validate data
|
||||
|
||||
name_match = cam_name_re.match(data['key'])
|
||||
|
||||
if name_match is not None:
|
||||
name = name_match.group('name')
|
||||
|
||||
values = data['value']
|
||||
|
||||
intrinsics = np.array([[values['focal_length_x_px_'], 0, values['focal_center_x_px_']],
|
||||
[0, values['focal_length_y_px_'], values['focal_center_x_px_']],
|
||||
[0, 0, 1]])
|
||||
extrinsics = json_to_pose(values['vehicle_SE3_camera_'])
|
||||
|
||||
result = Camera(name, intrinsics, extrinsics, dataset)
|
||||
session.add(result)
|
||||
session.commit()
|
||||
|
||||
return result
|
||||
|
||||
class CameraView(ModelBase):
|
||||
""" Represents a camera view at a given timestamp """
|
||||
__tablename__ = "camera_views"
|
||||
|
||||
id = Column(Integer, primary_key=True)
|
||||
|
||||
ext = Column(String(3), nullable=False)
|
||||
|
||||
dataset_id = Column(Integer, ForeignKey('datasets.id'), nullable=False)
|
||||
dataset = relationship('Dataset', back_populates='camera_views')
|
||||
|
||||
timestamp_id = Column(Integer, ForeignKey('timestamps.id'), nullable=False)
|
||||
timestamp = relationship('Timestamp', back_populates='camera_views')
|
||||
|
||||
camera_id = Column(Integer, ForeignKey('cameras.id'), nullable=False)
|
||||
camera = relationship('Camera', back_populates='camera_views', order_by=Timestamp.time)
|
||||
|
||||
def __init__(self, ext: str, dataset: Dataset, timestamp: Timestamp, camera: Camera):
|
||||
""" Constructor
|
||||
|
||||
Args:
|
||||
dataset (Dataset): dataset this camera view belongs to
|
||||
timestamp (Timestamp): timestamp this camera view was taken
|
||||
camera (Camera): camera that took this camera view
|
||||
"""
|
||||
self.ext = ext
|
||||
self.dataset = dataset
|
||||
self.timestamp = timestamp
|
||||
self.camera = camera
|
||||
|
||||
def get_path(self) -> Path:
|
||||
""" Get the path to the camera view image file
|
||||
|
||||
Returns:
|
||||
Path: Path to the camera view image file
|
||||
"""
|
||||
return Path(str(self.dataset.root_path)) / str(self.camera.name) / f"{self.camera.name}_{self.timestamp.time}.{self.ext}"
|
||||
|
||||
@staticmethod
|
||||
def from_path(path:PathLike, dataset:Dataset, session:Session) -> Optional['CameraView']:
|
||||
""" Creates a camera view from a file path
|
||||
|
||||
Args:
|
||||
path (PathLike): Path to the camera view image file
|
||||
dataset (Dataset): dataset this camera view belongs to
|
||||
|
||||
Returns:
|
||||
Optional[CameraView]: CameraView object if the path is valid, None otherwise
|
||||
"""
|
||||
result = None
|
||||
|
||||
path = Path(path)
|
||||
|
||||
camview_fn_match = camview_fn_re.match(path.stem)
|
||||
if camview_fn_match is not None:
|
||||
time = int(camview_fn_match.group('timestamp'))
|
||||
cam_name = camview_fn_match.group('cam_name')
|
||||
|
||||
timestamp = session.query(Timestamp).filter_by(time=time).first()
|
||||
camera = session.query(Camera).filter_by(name=cam_name).first()
|
||||
|
||||
if camera is not None and timestamp is not None:
|
||||
result = CameraView(path.suffix[1:], dataset, timestamp, camera)
|
||||
session.add(result)
|
||||
session.commit()
|
||||
|
||||
return result
|
||||
|
||||
Reference in New Issue
Block a user