134 lines
4.1 KiB
Python
134 lines
4.1 KiB
Python
|
from pathlib import Path
|
||
|
|
from sys import exit
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
from configargparse import (
|
||
|
|
ArgParser,
|
||
|
|
ArgumentDefaultsRawHelpFormatter,
|
||
|
|
YAMLConfigFileParser,
|
||
|
|
)
|
||
|
|
from numpy.lib import recfunctions as rfn
|
||
|
|
from rich.progress import Progress
|
||
|
|
from rosbags.highlevel import AnyReader
|
||
|
|
|
||
|
|
from util import existing_path
|
||
|
|
|
||
|
|
# Mapping of PointField datatypes to NumPy dtypes
|
||
|
|
POINTFIELD_DATATYPES = {
|
||
|
|
1: np.int8, # INT8
|
||
|
|
2: np.uint8, # UINT8
|
||
|
|
3: np.int16, # INT16
|
||
|
|
4: np.uint16, # UINT16
|
||
|
|
5: np.int32, # INT32
|
||
|
|
6: np.uint32, # UINT32
|
||
|
|
7: np.float32, # FLOAT32
|
||
|
|
8: np.float64, # FLOAT64
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def read_pointcloud(msg):
|
||
|
|
# Build the dtype dynamically from the fields
|
||
|
|
dtype_fields = {}
|
||
|
|
column_names = []
|
||
|
|
current_offset = 0
|
||
|
|
|
||
|
|
for field in msg.fields:
|
||
|
|
np_dtype = POINTFIELD_DATATYPES.get(field.datatype)
|
||
|
|
if np_dtype is None:
|
||
|
|
raise ValueError(
|
||
|
|
f"Unsupported datatype {field.datatype} for field {field.name}"
|
||
|
|
)
|
||
|
|
|
||
|
|
if field.offset > current_offset:
|
||
|
|
gap_size = field.offset - current_offset
|
||
|
|
gap_field_name = f"unused_{current_offset}"
|
||
|
|
dtype_fields[gap_field_name] = (
|
||
|
|
f"V{gap_size}",
|
||
|
|
current_offset,
|
||
|
|
) # Raw bytes as filler
|
||
|
|
current_offset += gap_size
|
||
|
|
|
||
|
|
dtype_fields[field.name] = (np_dtype, field.offset)
|
||
|
|
column_names.append(field.name)
|
||
|
|
current_offset = field.offset + np_dtype().itemsize
|
||
|
|
|
||
|
|
if current_offset < msg.point_step:
|
||
|
|
gap_size = msg.point_step - current_offset
|
||
|
|
gap_field_name = f"unused_{current_offset}"
|
||
|
|
dtype_fields[gap_field_name] = (f"V{gap_size}", current_offset)
|
||
|
|
|
||
|
|
dtype = np.dtype(dtype_fields)
|
||
|
|
return np.frombuffer(msg.data, dtype=dtype)
|
||
|
|
|
||
|
|
|
||
|
|
def clean_pointcloud(points):
|
||
|
|
valid_fields = [
|
||
|
|
name for name in points.dtype.names if not name.startswith("unused_")
|
||
|
|
]
|
||
|
|
cleaned_points = rfn.repack_fields(points[valid_fields])
|
||
|
|
return cleaned_points
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> int:
|
||
|
|
parser = ArgParser(
|
||
|
|
config_file_parser_class=YAMLConfigFileParser,
|
||
|
|
default_config_files=["data_analyze_config.yaml"],
|
||
|
|
formatter_class=ArgumentDefaultsRawHelpFormatter,
|
||
|
|
description="Analyse data from a rosbag or mcap file and output additional data",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--config-file", is_config_file=True, help="yaml config file path"
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--input-experiment-path",
|
||
|
|
required=True,
|
||
|
|
type=existing_path,
|
||
|
|
help="path to experiment. (directly to bag file, to parent folder for mcap)",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--pointcloud-topic",
|
||
|
|
default="/ouster/points",
|
||
|
|
type=str,
|
||
|
|
help="topic in the ros/mcap bag file containing the point cloud data",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--output-path",
|
||
|
|
default=Path("./output"),
|
||
|
|
type=Path,
|
||
|
|
help="path augmented dataset should be written to",
|
||
|
|
)
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
output_path = args.output_path / args.input_experiment_path.stem
|
||
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
parser.write_config_file(
|
||
|
|
parser.parse_known_args()[0],
|
||
|
|
output_file_paths=[(output_path / "config.yaml").as_posix()],
|
||
|
|
)
|
||
|
|
|
||
|
|
with AnyReader([args.input_experiment_path]) as reader:
|
||
|
|
connections = reader.connections
|
||
|
|
topics = dict(sorted({conn.topic: conn for conn in connections}.items()))
|
||
|
|
assert (
|
||
|
|
args.pointcloud_topic in topics
|
||
|
|
), f"Topic {args.pointcloud_topic} not found"
|
||
|
|
topic = topics[args.pointcloud_topic]
|
||
|
|
|
||
|
|
with Progress() as progress:
|
||
|
|
task = progress.add_task("Analyzing data", total=topic.msgcount)
|
||
|
|
for connection, timestamp, rawdata in reader.messages(connections=[topic]):
|
||
|
|
pointcloud_msg = reader.deserialize(rawdata, connection.msgtype)
|
||
|
|
original_pointcloud = read_pointcloud(pointcloud_msg)
|
||
|
|
cleaned_pointcloud = clean_pointcloud(original_pointcloud)
|
||
|
|
|
||
|
|
# todo add analysis here
|
||
|
|
|
||
|
|
progress.advance(task)
|
||
|
|
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
exit(main())
|