115 lines
3.1 KiB
Python
115 lines
3.1 KiB
Python
from argparse import ArgumentTypeError
|
|
from datetime import timedelta
|
|
from pathlib import Path
|
|
from subprocess import run
|
|
|
|
from matplotlib import colormaps
|
|
from matplotlib.colors import Colormap
|
|
from pointcloudset import Dataset
|
|
|
|
|
|
def load_dataset(
|
|
bag_file_path: Path, pointcloud_topic: str = "/ouster/points"
|
|
) -> Dataset:
|
|
return Dataset.from_file(bag_file_path, topic=pointcloud_topic)
|
|
|
|
|
|
def save_dataset(dataset: Dataset, output_file_path: Path):
|
|
if not output_file_path.is_dir():
|
|
raise ArgumentTypeError(f"{output_file_path} has to be a valid folder!")
|
|
dataset.to_file(output_file_path)
|
|
|
|
|
|
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)
|