split subter implementation (training + inference)
This commit is contained in:
17
tools/poetry.lock
generated
17
tools/poetry.lock
generated
@@ -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"
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user