split subter implementation (training + inference)

This commit is contained in:
Jan Kowalczyk
2024-10-17 08:36:18 +02:00
parent 5014c41b24
commit ddf4e4aa36
11 changed files with 699 additions and 145 deletions

17
tools/poetry.lock generated
View File

@@ -1435,6 +1435,21 @@ toolz = "*"
[package.extras]
complete = ["blosc", "numpy (>=1.20.0)", "pandas (>=1.3)", "pyzmq"]
[[package]]
name = "pathvalidate"
version = "3.2.0"
description = "pathvalidate is a Python library to sanitize/validate a string such as filenames/file-paths/etc."
optional = false
python-versions = ">=3.7"
files = [
{file = "pathvalidate-3.2.0-py3-none-any.whl", hash = "sha256:cc593caa6299b22b37f228148257997e2fa850eea2daf7e4cc9205cef6908dee"},
{file = "pathvalidate-3.2.0.tar.gz", hash = "sha256:5e8378cf6712bff67fbe7a8307d99fa8c1a0cb28aa477056f8fc374f0dff24ad"},
]
[package.extras]
docs = ["Sphinx (>=2.4)", "sphinx-rtd-theme (>=1.2.2)", "urllib3 (<2)"]
test = ["Faker (>=1.0.8)", "allpairspy (>=2)", "click (>=6.2)", "pytest (>=6.0.1)", "pytest-discord (>=0.1.4)", "pytest-md-report (>=0.4.1)"]
[[package]]
name = "pillow"
version = "10.3.0"
@@ -2416,4 +2431,4 @@ cffi = ["cffi (>=1.11)"]
[metadata]
lock-version = "2.0"
python-versions = "^3.11"
content-hash = "607634c6605566c6bec22c1e8f41d056d2920de45108297193592fe57b10507b"
content-hash = "cd1f01135a813fbbc06da686dfa67b0981d26ccb3038f7d0e4a17359326f2c8d"

View File

@@ -19,6 +19,7 @@ matplotlib = "^3.8.4"
dask = "^2024.4.2"
dask-expr = "^1.1.3"
pandas = "^2.2.2"
pathvalidate = "^3.2.0"
[build-system]
requires = ["poetry-core"]

View File

@@ -1,18 +1,22 @@
from configargparse import (
ArgParser,
YAMLConfigFileParser,
ArgumentDefaultsRawHelpFormatter,
)
from sys import exit
from pathlib import Path
from pointcloudset import Dataset
from rich.progress import track
from pandas import DataFrame
from PIL import Image
from math import pi
from typing import Optional
from multiprocessing import Pool
from pathlib import Path
from sys import exit
from typing import Optional, Tuple
import matplotlib
import numpy as np
from configargparse import (
ArgParser,
ArgumentDefaultsRawHelpFormatter,
YAMLConfigFileParser,
)
from pandas import DataFrame
from pathvalidate import sanitize_filename
from PIL import Image
from pointcloudset import Dataset
from rich.progress import track
from rosbags.highlevel import AnyReader
matplotlib.use("Agg")
import matplotlib.pyplot as plt
@@ -20,16 +24,44 @@ import matplotlib.pyplot as plt
from util import (
angle,
angle_width,
positive_int,
load_dataset,
existing_path,
create_video_from_images,
calculate_average_frame_rate,
create_video_from_images,
existing_path,
get_colormap_with_special_missing_color,
load_dataset,
positive_int,
)
def crop_lidar_data_to_roi(
def save_image_topics(
bag_file_path: Path, output_folder_path: Path, topic_names: list[str]
) -> list[Path]:
with AnyReader([bag_file_path]) as reader:
topic_paths = []
for topic_name in topic_names:
connections = [x for x in reader.connections if x.topic == topic_name]
frame_count = 0
topic_output_folder_path = output_folder_path / sanitize_filename(
topic_name
)
topic_output_folder_path.mkdir(exist_ok=True, parents=True)
topic_paths.append(topic_output_folder_path)
for connection, timestamp, rawdata in reader.messages(
connections=connections
):
img_msg = reader.deserialize(rawdata, connection.msgtype)
img_data = np.frombuffer(
img_msg.data,
dtype=np.uint8 if "8" in img_msg.encoding else np.uint16,
).reshape((img_msg.height, img_msg.width))
img = Image.fromarray(img_data, mode="L")
frame_count += 1
img.save(topic_output_folder_path / f"frame_{frame_count:04d}.png")
return topic_paths
def crop_projection_data_to_roi(
data: DataFrame,
roi_angle_start: float,
roi_angle_width: float,
@@ -83,103 +115,196 @@ def create_2d_projection(
tmp_file_path.unlink()
def process_frame(args) -> Tuple[int, np.ndarray, Optional[Path]]:
(
i,
pc,
output_path,
colormap_name,
missing_data_color,
reverse_colormap,
vertical_resolution,
horizontal_resolution,
vertical_scale,
horizontal_scale,
roi_angle_start,
roi_angle_width,
render_images,
) = args
lidar_data = pc.data.copy()
lidar_data["horizontal_position"] = (
lidar_data["original_id"] % horizontal_resolution
)
lidar_data["horizontal_position_yaw_f"] = (
0.5
* horizontal_resolution
* (np.arctan2(lidar_data["y"], lidar_data["x"]) / pi + 1.0)
)
lidar_data["horizontal_position_yaw"] = np.floor(
lidar_data["horizontal_position_yaw_f"]
)
lidar_data["vertical_position"] = np.floor(
lidar_data["original_id"] / horizontal_resolution
)
# fov = 32 * pi / 180
# fov_down = 17 * pi / 180
fov = 31.76 * pi / 180
fov_down = 17.3 * pi / 180
lidar_data["vertical_angle"] = np.arcsin(
lidar_data["z"]
/ np.sqrt(lidar_data["x"] ** 2 + lidar_data["y"] ** 2 + lidar_data["z"] ** 2)
)
lidar_data["vertical_angle_degree"] = lidar_data["vertical_angle"] * 180 / pi
lidar_data["vertical_position_pitch_f"] = vertical_resolution * (
1 - ((lidar_data["vertical_angle"] + fov_down) / fov)
)
lidar_data["vertical_position_pitch"] = np.floor(
lidar_data["vertical_position_pitch_f"]
)
duplicates = lidar_data[
lidar_data.duplicated(
subset=["vertical_position_pitch", "horizontal_position_yaw"],
keep=False,
)
].sort_values(by=["vertical_position_pitch", "horizontal_position_yaw"])
lidar_data["normalized_range"] = 1 / np.sqrt(
lidar_data["x"] ** 2 + lidar_data["y"] ** 2 + lidar_data["z"] ** 2
)
projection_data = lidar_data.pivot(
index="vertical_position_pitch",
columns="horizontal_position_yaw",
values="normalized_range",
)
projection_data = projection_data.reindex(
columns=range(horizontal_resolution), fill_value=0
)
projection_data = projection_data.reindex(
index=range(vertical_resolution), fill_value=0
)
projection_data, output_horizontal_resolution = crop_projection_data_to_roi(
projection_data, roi_angle_start, roi_angle_width, horizontal_resolution
)
if render_images:
image_path = create_2d_projection(
projection_data,
output_path / f"frame_{i:04d}.png",
output_path / f"tmp_{i:04d}.png",
colormap_name,
missing_data_color,
reverse_colormap,
horizontal_resolution=output_horizontal_resolution * horizontal_scale,
vertical_resolution=vertical_resolution * vertical_scale,
)
else:
image_path = None
return (
i,
projection_data.to_numpy(),
image_path,
lidar_data["vertical_position_pitch_f"].min(),
lidar_data["vertical_position_pitch_f"].max(),
)
# Adjusted to use a generator for args_list
def create_projection_data(
dataset: Dataset,
output_path: Path,
colormap_name: str,
missing_data_color: str,
reverse_colormap: bool,
vertical_resolution: int,
horizontal_resolution: int,
vertical_scale: int,
horizontal_scale: int,
roi_angle_start: float,
roi_angle_width: float,
render_images: bool,
) -> (np.ndarray, Optional[list[Path]]):
) -> Tuple[np.ndarray, Optional[list[Path]]]:
rendered_images = []
converted_lidar_frames = []
# Generator for args_list
# def args_generator():
# for i, pc in enumerate(dataset, 1):
# yield (
# i,
# pc,
# output_path,
# colormap_name,
# missing_data_color,
# reverse_colormap,
# vertical_resolution,
# horizontal_resolution,
# vertical_scale,
# horizontal_scale,
# roi_angle_start,
# roi_angle_width,
# render_images,
# )
# results = []
# with Pool() as pool:
# results_gen = pool.imap(process_frame, args_generator())
# for result in track(
# results_gen, description="Processing...", total=len(dataset)
# ):
# results.append(result)
# results.sort(key=lambda x: x[0])
for i, pc in track(
enumerate(dataset, 1), description="Creating projections...", total=len(dataset)
enumerate(dataset, 1), description="Processing...", total=len(dataset)
):
vertical_resolution = int(pc.data["ring"].max() + 1)
# Angle calculation implementation
# projected_data = pc.data.copy()
# projected_data["arctan"] = np.arctan2(projected_data["y"], projected_data["x"])
# projected_data["arctan_normalized"] = 0.5 * (projected_data["arctan"] / pi + 1.0)
# projected_data["arctan_scaled"] = projected_data["arctan_normalized"] * horizontal_resolution
# #projected_data["horizontal_position"] = np.floor(projected_data["arctan_scaled"])
# projected_data["horizontal_position"] = np.round(projected_data["arctan_scaled"])
# projected_data["normalized_range"] = 1 / np.sqrt(
# projected_data["x"] ** 2 + projected_data["y"] ** 2 + projected_data["z"] ** 2
# )
# duplicates = projected_data[projected_data.duplicated(subset=['ring', 'horizontal_position'], keep=False)].sort_values(by=['ring', 'horizontal_position'])
# sorted = projected_data.sort_values(by=['ring', 'horizontal_position'])
# FIXME: following pivot fails due to duplicates in the data, some points (x, y) are mapped to the same pixel in the projection, have to decide how to handles
# these cases
# projected_image_data = projected_data.pivot(
# index="ring", columns="horizontal_position", values="normalized_range"
# )
# projected_image_data = projected_image_data.reindex(columns=range(horizontal_resolution), fill_value=0)
# projected_image_data, output_horizontal_resolution = crop_lidar_data_to_roi(
# projected_image_data, roi_angle_start, roi_angle_width, horizontal_resolution
# )
# create_2d_projection(
# projected_image_data,
# output_path / f"frame_{i:04d}_projection.png",
# output_path / "tmp.png",
# colormap_name,
# missing_data_color,
# reverse_colormap,
# horizontal_resolution=output_horizontal_resolution * horizontal_scale,
# vertical_resolution=vertical_resolution * vertical_scale,
# )
lidar_data = pc.data.copy()
lidar_data["horizontal_position"] = (
lidar_data["original_id"] % horizontal_resolution
)
lidar_data["normalized_range"] = 1 / np.sqrt(
lidar_data["x"] ** 2 + lidar_data["y"] ** 2 + lidar_data["z"] ** 2
)
lidar_data = lidar_data.pivot(
index="ring", columns="horizontal_position", values="normalized_range"
)
lidar_data = lidar_data.reindex(
columns=range(horizontal_resolution), fill_value=0
)
lidar_data = lidar_data.reindex(index=range(vertical_resolution), fill_value=0)
lidar_data, output_horizontal_resolution = crop_lidar_data_to_roi(
lidar_data, roi_angle_start, roi_angle_width, horizontal_resolution
)
converted_lidar_frames.append(lidar_data.to_numpy())
if render_images:
image_path = create_2d_projection(
lidar_data,
output_path / f"frame_{i:04d}.png",
output_path / "tmp.png",
results = process_frame(
(
i,
pc,
output_path,
colormap_name,
missing_data_color,
reverse_colormap,
horizontal_resolution=output_horizontal_resolution * horizontal_scale,
vertical_resolution=vertical_resolution * vertical_scale,
vertical_resolution,
horizontal_resolution,
vertical_scale,
horizontal_scale,
roi_angle_start,
roi_angle_width,
render_images,
)
)
converted_lidar_frames.append(results[1])
if render_images:
rendered_images.append(results[2])
rendered_images.append(image_path)
# min_all = 100
# max_all = 0
# for _, data, img_path, min_frame, max_frame in results:
# converted_lidar_frames.append(data)
# if img_path:
# rendered_images.append(img_path)
# if min_frame < min_all:
# min_all = min_frame
# if max_frame > max_all:
# max_all = max_frame
# print(f"{min_all=}, {max_all=}")
projection_data = np.stack(converted_lidar_frames, axis=0)
if render_images:
return rendered_images, projection_data
return projection_data, rendered_images
else:
return projection_data
return (projection_data,)
def main() -> int:
@@ -204,6 +329,12 @@ def main() -> int:
type=str,
help="topic in the ros/mcap bag file containing the point cloud data",
)
parser.add_argument(
"--image-topics",
default=[],
nargs="+",
help="topics in the ros/mcap bag file containing the image data",
)
parser.add_argument(
"--output-path",
default=Path("./output"),
@@ -248,11 +379,17 @@ def main() -> int:
type=bool,
help="if colormap should be reversed",
)
parser.add_argument(
"--vertical-resolution",
default=32,
type=positive_int,
help="number of vertical lidar data point rows",
)
parser.add_argument(
"--horizontal-resolution",
default=2048,
type=positive_int,
help="number of horizontal lidar data points",
help="number of horizontal lidar data point columns",
)
parser.add_argument(
"--vertical-scale",
@@ -296,6 +433,7 @@ def main() -> int:
dataset = load_dataset(args.input_experiment_path, args.pointcloud_topic)
images = []
topic_paths = []
if not args.output_no_images or not args.output_no_video:
if not args.force_generation and all(
@@ -312,6 +450,7 @@ def main() -> int:
args.colormap_name,
args.missing_data_color,
args.reverse_colormap,
args.vertical_resolution,
args.horizontal_resolution,
args.vertical_scale,
args.horizontal_scale,
@@ -319,6 +458,10 @@ def main() -> int:
args.roi_angle_width,
render_images=True,
)
if args.image_topics:
topic_paths = save_image_topics(
args.input_experiment_path, output_path, args.image_topics
)
output_numpy_path = (output_path / args.input_experiment_path.stem).with_suffix(
".npy"
@@ -336,6 +479,7 @@ def main() -> int:
args.colormap_name,
args.missing_data_color,
args.reverse_colormap,
args.vertical_resolution,
args.horizontal_resolution,
args.vertical_scale,
args.horizontal_scale,
@@ -358,12 +502,20 @@ def main() -> int:
f"Skipping video generation for {args.input_experiment_path} as {output_path / args.input_experiment_path.stem}.mp4 already exists"
)
else:
frame_rate = calculate_average_frame_rate(dataset)
input_images_pattern = f"{tmp_path}/frame_%04d.png"
create_video_from_images(
input_images_pattern,
(output_path / args.input_experiment_path.stem).with_suffix(".mp4"),
calculate_average_frame_rate(dataset),
frame_rate,
)
for topic_path in topic_paths:
input_images_pattern = f"{topic_path}/frame_%04d.png"
create_video_from_images(
input_images_pattern,
(output_path / topic_path.stem).with_suffix(".mp4"),
frame_rate,
)
if args.output_no_images:
for image in images:

View File

@@ -1,11 +1,11 @@
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 pathlib import Path
from subprocess import run
from matplotlib import colormaps
from matplotlib.colors import Colormap
from pointcloudset import Dataset
def load_dataset(