fedex/mt
Public Access
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

added pickle file for training with lidar range data to output options

Browse Source
This commit is contained in:
Jan Kowalczyk
2024-06-11 15:56:33 +02:00
parent e92c6a0a84
commit 2eb1bf2e05
2 changed files with 164 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files

220
tools/render2d.py
View File

@@ -16,7 +16,9 @@ matplotlib.use("Agg")
import matplotlib.pyplot as plt
from util import (
angle, angle_width, positive_int,
angle,
angle_width,
positive_int,
load_dataset,
existing_path,
create_video_from_images,
@@ -25,6 +27,70 @@ from util import (
)
def fill_sparse_data(data: DataFrame, horizontal_resolution: int) -> DataFrame:
complete_original_ids = DataFrame(
{
"original_id": np.arange(
0,
(data["ring"].max() + 1) * horizontal_resolution,
dtype=np.uint32,
)
}
)
data = complete_original_ids.merge(data, on="original_id", how="left")
data["ring"] = data["original_id"] // horizontal_resolution
data["horizontal_position"] = data["original_id"] % horizontal_resolution
return data
def crop_lidar_data_to_roi(
data: DataFrame,
roi_angle_start: float,
roi_angle_width: float,
horizontal_resolution: int,
) -> tuple[DataFrame, int]:
if roi_angle_width == 360:
return data, horizontal_resolution
roi_index_start = int(horizontal_resolution / 360 * roi_angle_start)
roi_index_width = int(horizontal_resolution / 360 * roi_angle_width)
roi_index_end = roi_index_start + roi_index_width
if roi_index_end < horizontal_resolution:
cropped_data = data.iloc[:, roi_index_start:roi_index_end]
else:
roi_index_end = roi_index_end - horizontal_resolution
cropped_data = data.iloc[:, roi_index_end:roi_index_start]
return cropped_data, roi_index_width
def create_projection_data(
dataset: Dataset,
horizontal_resolution: int,
roi_angle_start: float,
roi_angle_width: float,
) -> list[Path]:
converted_lidar_frames = []
for i, pc in track(
enumerate(dataset, 1), description="Rendering images...", total=len(dataset)
):
lidar_data = fill_sparse_data(pc.data, 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, _ = crop_lidar_data_to_roi(
lidar_data, roi_angle_start, roi_angle_width, horizontal_resolution
)
converted_lidar_frames.append(lidar_data.to_numpy())
return np.stack(converted_lidar_frames, axis=0)
def create_2d_projection(
df: DataFrame,
output_file_path: Path,
@@ -35,7 +101,9 @@ def create_2d_projection(
horizontal_resolution: int,
vertical_resolution: int,
):
fig, ax = plt.subplots(figsize=(float(horizontal_resolution) / 100, float(vertical_resolution) / 100))
fig, ax = plt.subplots(
figsize=(float(horizontal_resolution) / 100, float(vertical_resolution) / 100)
)
ax.imshow(
df,
cmap=get_colormap_with_special_missing_color(
@@ -48,60 +116,50 @@ def create_2d_projection(
plt.savefig(tmp_file_path, dpi=100, bbox_inches="tight", pad_inches=0)
plt.close()
img = Image.open(tmp_file_path)
img_resized = img.resize((horizontal_resolution, vertical_resolution), Image.LANCZOS)
img_resized = img.resize(
(horizontal_resolution, vertical_resolution), Image.LANCZOS
)
img_resized.save(output_file_path)
tmp_file_path.unlink()
def render_2d_images(
dataset: Dataset,
output_images_path: Path,
image_pattern_prefix: str,
tmp_files_path: Path,
output_path: Path,
colormap_name: str,
missing_data_color: str,
reverse_colormap: bool,
horizontal_resolution: int,
roi_angle_start: float,
roi_angle_width: float,
vertical_scale: int,
horizontal_scale: int,
roi_angle_start: float,
roi_angle_width: float,
) -> list[Path]:
rendered_images = []
for i, pc in track(
enumerate(dataset, 1), description="Rendering images...", total=len(dataset)
):
complete_original_ids = DataFrame({'original_id': np.arange(0, (pc.data['ring'].max() + 1) * horizontal_resolution, dtype=np.uint32)})
pc.data = complete_original_ids.merge(pc.data, on='original_id', how='left')
pc.data['ring'] = (pc.data['original_id'] // horizontal_resolution)
pc.data["horizontal_position"] = pc.data["original_id"] % horizontal_resolution
image_data = pc.data.pivot(
image_data = fill_sparse_data(pc.data, horizontal_resolution).pivot(
index="ring", columns="horizontal_position", values="range"
)
if roi_angle_width != 360:
roi_index_start = int(horizontal_resolution / 360 * roi_angle_start)
roi_index_width = int(horizontal_resolution / 360 * roi_angle_width)
roi_index_end = roi_index_start + roi_index_width
if roi_index_end < horizontal_resolution:
image_data = image_data.iloc[:, roi_index_start:roi_index_end]
else:
roi_index_end = roi_index_end - horizontal_resolution
image_data = image_data.iloc[:, roi_index_end:roi_index_start]
image_data, output_horizontal_resolution = crop_lidar_data_to_roi(
image_data, roi_angle_start, roi_angle_width, horizontal_resolution
)
normalized_data = (image_data - image_data.min().min()) / (
image_data.max().max() - image_data.min().min()
)
image_path = create_2d_projection(
normalized_data,
output_images_path / f"{image_pattern_prefix}_frame_{i:04d}.png",
tmp_files_path / "tmp.png",
output_path / f"frame_{i:04d}.png",
output_path / "tmp.png",
colormap_name,
missing_data_color,
reverse_colormap,
horizontal_resolution=(roi_index_width if roi_angle_width != 360 else horizontal_resolution) * horizontal_scale,
vertical_resolution=(pc.data['ring'].max() + 1) * vertical_scale
horizontal_resolution=output_horizontal_resolution * horizontal_scale,
vertical_resolution=(pc.data["ring"].max() + 1) * vertical_scale,
)
rendered_images.append(image_path)
@@ -120,13 +178,22 @@ def main() -> int:
"--render-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)"
"--input-experiment-path",
required=True,
type=existing_path,
help="path to experiment. (directly to bag file, to parent folder for mcap)",
)
parser.add_argument(
"--tmp-files-path",
default=Path("./tmp"),
"--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 temporary files will be written to",
help="path rendered frames should be written to",
)
parser.add_argument(
"--output-images",
@@ -134,12 +201,6 @@ def main() -> int:
default=True,
help="if rendered frames should be outputted as images",
)
parser.add_argument(
"--output-images-path",
default=Path("./output"),
type=Path,
help="path rendered frames should be written to",
)
parser.add_argument(
"--output-video",
type=bool,
@@ -147,16 +208,16 @@ def main() -> int:
help="if rendered frames should be outputted as a video",
)
parser.add_argument(
"--output-video-path",
default=Path("./output/2d_render.mp4"),
type=Path,
help="path rendered video should be written to",
"--output-pickle",
default=True,
type=bool,
help="if the processed data should be saved as a pickle file",
)
parser.add_argument(
"--output-images-prefix",
default="2d_render",
type=str,
help="filename prefix for output",
"--skip-existing",
default=True,
type=bool,
help="if true will skip rendering existing files",
)
parser.add_argument(
"--colormap-name",
@@ -176,30 +237,12 @@ def main() -> int:
type=bool,
help="if colormap should be reversed",
)
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(
"--horizontal-resolution",
default=2048,
type=positive_int,
help="number of horizontal lidar data points",
)
parser.add_argument(
"--roi-angle-start",
default=0,
type=angle,
help="angle where roi starts",
)
parser.add_argument(
"--roi-angle-width",
default=360,
type=angle_width,
help="width of roi in degrees",
)
parser.add_argument(
"--vertical-scale",
default=1,
@@ -212,48 +255,67 @@ def main() -> int:
type=positive_int,
help="multiplier for horizontal scale, for better visualization",
)
parser.add_argument(
"--roi-angle-start",
default=0,
type=angle,
help="angle where roi starts",
)
parser.add_argument(
"--roi-angle-width",
default=360,
type=angle_width,
help="width of roi in degrees",
)
args = parser.parse_args()
if args.output_images:
args.output_images_path.mkdir(parents=True, exist_ok=True)
args.tmp_files_path = args.output_images_path
else:
args.tmp_files_path.mkdir(parents=True, exist_ok=True)
output_path = args.output_path / args.input_experiment_path.stem
output_path.mkdir(parents=True, exist_ok=True)
if args.output_video:
args.output_video_path.parent.mkdir(parents=True, exist_ok=True)
# Create temporary folder for images, if outputting images we use the output folder itself as temp folder
tmp_path = output_path / "frames" if args.output_images else output_path / "tmp"
tmp_path.mkdir(parents=True, exist_ok=True)
dataset = load_dataset(args.input_experiment_path, args.pointcloud_topic)
images = render_2d_images(
dataset,
args.tmp_files_path,
args.output_images_prefix,
args.tmp_files_path,
tmp_path,
args.colormap_name,
args.missing_data_color,
args.reverse_colormap,
args.horizontal_resolution,
args.roi_angle_start,
args.roi_angle_width,
args.vertical_scale,
args.horizontal_scale,
args.roi_angle_start,
args.roi_angle_width,
)
if args.output_video:
input_images_pattern = (
f"{args.tmp_files_path / args.output_images_prefix}_frame_%04d.png"
if args.output_pickle:
output_pickle_path = (
output_path / args.input_experiment_path.stem
).with_suffix(".pkl")
processed_range_data = create_projection_data(
dataset,
args.horizontal_resolution,
args.roi_angle_start,
args.roi_angle_width,
)
processed_range_data.dump(output_pickle_path)
if args.output_video:
input_images_pattern = f"{tmp_path}/frame_%04d.png"
create_video_from_images(
input_images_pattern,
args.output_video_path,
(output_path / args.input_experiment_path.stem).with_suffix(".mp4"),
calculate_average_frame_rate(dataset),
)
if not args.output_images:
for image in images:
image.unlink()
tmp_path.rmdir()
return 0

30
tools/util.py
View File

@@ -8,14 +8,17 @@ from matplotlib.colors import Colormap
from matplotlib import colormaps
def load_dataset(bag_file_path: Path, pointcloud_topic: str = "/ouster/points") -> 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 calculate_average_frame_rate(dataset: Dataset):
timestamps = dataset.timestamps
time_deltas = [timestamps[i + 1] - timestamps[i] for i in range(len(timestamps) - 1)]
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
@@ -38,39 +41,52 @@ def existing_folder(path_string: str) -> Path:
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)")
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]")
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 = (
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:
def create_video_from_images(
input_images_pattern: str, output_file: Path, frame_rate: int
) -> None:
# Construct the ffmpeg command
command = [
"ffmpeg",