Files
mt/tools/util.py
2024-06-11 11:07:05 +02:00

93 lines
2.9 KiB
Python

from pointcloudset import Dataset
from pointcloudset.io.dataset.ros import dataset_from_ros
from pathlib import Path
from argparse import ArgumentTypeError
from subprocess import run
from datetime import timedelta
from matplotlib.colors import Colormap
from matplotlib import colormaps
def load_dataset(bag_file_path: Path, pointcloud_topic: str = "/ouster/points") -> Dataset:
return Dataset.from_file(bag_file_path, topic=pointcloud_topic)
def calculate_average_frame_rate(dataset: Dataset):
timestamps = dataset.timestamps
time_deltas = [timestamps[i + 1] - timestamps[i] for i in range(len(timestamps) - 1)]
average_delta = sum(time_deltas, timedelta()) / len(time_deltas)
average_frame_rate = 1 / average_delta.total_seconds()
return average_frame_rate
def existing_file(path_string: str) -> Path:
path = Path(path_string)
if not path.exists():
raise ArgumentTypeError(f"{path} does not exist!")
if not path.is_file():
raise ArgumentTypeError(f"{path} is not a valid file!")
return path
def existing_folder(path_string: str) -> Path:
path = Path(path_string)
if not path.exists():
raise ArgumentTypeError(f"{path} does not exist!")
if not path.is_dir():
raise ArgumentTypeError(f"{path} is not a valid folder!")
return path
def existing_path(path_string: str) -> Path:
path = Path(path_string)
if not path.exists():
raise ArgumentTypeError(f"{path} does not exist!")
return path
def positive_int(number_str: str) -> int:
number_val = int(number_str)
if number_val < 0:
raise ArgumentTypeError(f"{number_val} is not a positive integer!")
return number_val
def angle(angle_str: str) -> float:
angle_val = float(angle_str)
if angle_val < 0 or angle_val >= 360:
raise ArgumentTypeError(f"{angle_val} is not a valid angle! Needs to be in [0, 360)")
return angle_val
def angle_width(angle_str: str) -> float:
angle_val = float(angle_str)
if angle_val < 0 or angle_val > 360:
raise ArgumentTypeError(f"{angle_val} is not a valid angle width! Needs to be in [0, 360]")
return angle_val
def get_colormap_with_special_missing_color(
colormap_name: str, missing_data_color: str = "black", reverse: bool = False
) -> Colormap:
colormap = colormaps[colormap_name] if not reverse else colormaps[f"{colormap_name}_r"]
colormap.set_bad(missing_data_color)
return colormap
def create_video_from_images(input_images_pattern: str, output_file: Path, frame_rate: int) -> None:
# Construct the ffmpeg command
command = [
"ffmpeg",
"-y",
"-framerate",
str(frame_rate),
"-i",
input_images_pattern,
"-c:v",
"libx264",
"-profile:v",
"high",
"-crf",
"20",
"-pix_fmt",
"yuv420p",
output_file.as_posix(),
]
run(command, check=True)