-
Notifications
You must be signed in to change notification settings - Fork 243
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
allow options on rendering centerlines, directly generating mp4, and … #170
Merged
Merged
Changes from 9 commits
Commits
Show all changes
21 commits
Select commit
Hold shift + click to select a range
f3e7cea
allow options on rendering centerlines, directly generating mp4, and …
johnwlambert 54923b6
reduce duplicated code, and add additional compression function for c…
johnwlambert dea83ae
clean up whitespace
johnwlambert 68dacac
start video writer class
johnwlambert e911cd1
reformat with black
c27dd37
reformat with black
3881748
do not dump frames if generate_video_only specified, and add missing …
1eaca1d
improve presentation of codec params
5b4ccc0
Fix typo
95811fe
update to non-deprecated loader method
johnwlambert 26d59e4
fix merge conflicts
878c061
Merge branch 'master' of https://github.com/argoai/argoverse-api into…
6bfd801
Merge branch 'master' of https://github.com/argoai/argoverse-api into…
e2d73b8
Merge branch 'master' of https://github.com/argoai/argoverse-api into…
d14204e
improve module docstring
d60c2c2
improve docstrings
40757e8
clean up docstrings, no need to include Returns: section for no retur…
e3147b5
improve docstrings
9baa77e
improve readability
dc6a6d8
ensure writer is not None for mypy's sake
98c40ba
reduce required precision when matching Sim(2) and SE(2) results to 1e-5
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,51 @@ | ||
#!/usr/bin/python3 | ||
|
||
import cv2 | ||
import numpy as np | ||
|
||
""" | ||
Python-based utilities to avoid blowing up the disk with images, as FFMPEG requires. | ||
Inspired by Detectron2 and MSeg: | ||
https://github.com/facebookresearch/detectron2/blob/bab413cdb822af6214f9b7f70a9b7a9505eb86c5/demo/demo.py | ||
https://github.com/mseg-dataset/mseg-semantic/blob/master/mseg_semantic/utils/cv2_video_utils.py | ||
See OpenCV documentation for more details: | ||
https://docs.opencv.org/2.4/modules/highgui/doc/reading_and_writing_images_and_video.html | ||
""" | ||
|
||
class VideoWriter: | ||
""" | ||
Lazy init, so that the user doesn't have to know width/height a priori. | ||
Our default codec is "mp4v", though you may prefer "x264", if available | ||
on your system | ||
""" | ||
def __init__(self, output_fpath: str, fps: int = 30) -> None: | ||
""" """ | ||
self.output_fpath = output_fpath | ||
self.fps = fps | ||
self.writer = None | ||
self.codec = "mp4v" | ||
|
||
def init_outf(self, height: int, width: int) -> None: | ||
""" """ | ||
self.writer = cv2.VideoWriter( | ||
filename=self.output_fpath, | ||
# some installation of opencv may not support x264 (due to its license), | ||
# you can try other format (e.g. MPEG) | ||
fourcc=cv2.VideoWriter_fourcc(*self.codec), | ||
fps=float(self.fps), | ||
frameSize=(width, height), | ||
isColor=True, | ||
) | ||
|
||
def add_frame(self, rgb_frame: np.ndarray) -> None: | ||
""" | ||
""" | ||
h, w, _ = rgb_frame.shape | ||
if self.writer is None: | ||
self.init_outf(height=h, width=w) | ||
bgr_frame = rgb_frame[:,:,::-1] | ||
self.writer.write(bgr_frame) | ||
|
||
def complete(self) -> None: | ||
""" """ | ||
self.writer.release() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -3,9 +3,9 @@ | |
import copy | ||
import glob | ||
import logging | ||
import multiprocessing | ||
import os | ||
import sys | ||
from multiprocessing import Pool | ||
from pathlib import Path | ||
from typing import Any, Iterable, List, Mapping, Sequence, Tuple, Union | ||
|
||
|
@@ -26,7 +26,8 @@ | |
from argoverse.utils.camera_stats import RING_CAMERA_LIST, STEREO_CAMERA_LIST | ||
from argoverse.utils.city_visibility_utils import clip_point_cloud_to_visible_region | ||
from argoverse.utils.cv2_plotting_utils import draw_clipped_line_segment | ||
from argoverse.utils.ffmpeg_utils import write_nonsequential_idx_video | ||
from argoverse.utils.cv2_video_utils import VideoWriter | ||
from argoverse.utils.ffmpeg_utils import ffmpeg_compress_video, write_nonsequential_idx_video | ||
from argoverse.utils.frustum_clipping import generate_frustum_planes | ||
from argoverse.utils.ply_loader import load_ply | ||
from argoverse.utils.se3 import SE3 | ||
|
@@ -39,11 +40,13 @@ | |
|
||
# jigger lane pixel values by [-10,10] range | ||
LANE_COLOR_NOISE = 20 | ||
STEREO_FPS = 5 | ||
RING_CAM_FPS = 30 | ||
|
||
|
||
def plot_lane_centerlines_in_img( | ||
lidar_pts: np.ndarray, | ||
city_to_egovehicle_se3: SE3, | ||
city_SE3_egovehicle: SE3, | ||
img: np.ndarray, | ||
city_name: str, | ||
avm: ArgoverseMap, | ||
|
@@ -54,7 +57,7 @@ def plot_lane_centerlines_in_img( | |
) -> np.ndarray: | ||
""" | ||
Args: | ||
city_to_egovehicle_se3: SE3 transformation representing egovehicle to city transformation | ||
city_SE3_egovehicle: SE(3) transformation representing egovehicle to city transformation | ||
img: Array of shape (M,N,3) representing updated image | ||
city_name: str, string representing city name, i.e. 'PIT' or 'MIA' | ||
avm: instance of ArgoverseMap | ||
|
@@ -70,7 +73,7 @@ def plot_lane_centerlines_in_img( | |
t = camera_config.extrinsic[:3, 3] | ||
cam_SE3_egovehicle = SE3(rotation=R, translation=t) | ||
|
||
query_x, query_y, _ = city_to_egovehicle_se3.translation | ||
query_x, query_y, _ = city_SE3_egovehicle.translation | ||
local_centerlines = avm.find_local_lane_centerlines(query_x, query_y, city_name) | ||
|
||
for centerline_city_fr in local_centerlines: | ||
|
@@ -81,7 +84,7 @@ def plot_lane_centerlines_in_img( | |
valid_idx = np.isnan(ground_heights) | ||
centerline_city_fr = centerline_city_fr[~valid_idx] | ||
|
||
centerline_egovehicle_fr = city_to_egovehicle_se3.inverse().transform_point_cloud(centerline_city_fr) | ||
centerline_egovehicle_fr = city_SE3_egovehicle.inverse().transform_point_cloud(centerline_city_fr) | ||
centerline_uv_cam = cam_SE3_egovehicle.transform_point_cloud(centerline_egovehicle_fr) | ||
|
||
# can also clip point cloud to nearest LiDAR point depth | ||
|
@@ -105,6 +108,8 @@ def dump_clipped_3d_cuboids_to_images( | |
data_dir: str, | ||
experiment_prefix: str, | ||
motion_compensate: bool = True, | ||
omit_centerlines: bool = False, | ||
generate_video_only: bool = False, | ||
) -> List[str]: | ||
""" | ||
We bring the 3D points into each camera coordinate system, and do the clipping there in 3D. | ||
|
@@ -115,24 +120,40 @@ def dump_clipped_3d_cuboids_to_images( | |
data_dir: path to dataset with the latest data | ||
experiment_prefix: Output directory | ||
motion_compensate: Whether to motion compensate when projecting | ||
omit_centerlines: whether to omit map vector lane geometry from rendering | ||
generate_video_only: whether to generate mp4s only without dumping individual frames | ||
|
||
Returns: | ||
saved_img_fpaths | ||
""" | ||
saved_img_fpaths = [] | ||
dl = SimpleArgoverseTrackingDataLoader(data_dir=data_dir, labels_dir=data_dir) | ||
avm = ArgoverseMap() | ||
if not omit_centerlines: | ||
avm = ArgoverseMap() | ||
fps_map = { | ||
cam_name: STEREO_FPS if "stereo" in cam_name else RING_CAM_FPS | ||
for cam_name in RING_CAMERA_LIST + STEREO_CAMERA_LIST | ||
} | ||
category_subdir = "amodal_labels" | ||
if not Path(f"{experiment_prefix}_{category_subdir}").exists(): | ||
os.makedirs(f"{experiment_prefix}_{category_subdir}") | ||
video_output_dir = f"{experiment_prefix}_{category_subdir}" | ||
|
||
for log_id in log_ids: | ||
save_dir = f"{experiment_prefix}_{log_id}" | ||
if not Path(save_dir).exists(): | ||
if not generate_video_only and not Path(save_dir).exists(): | ||
# JPG images will be dumped here, if requested by arguments | ||
os.makedirs(save_dir) | ||
|
||
city_name = dl.get_city_name(log_id) | ||
log_calib_data = dl.get_log_calibration_data(log_id) | ||
|
||
flag_done = False | ||
for cam_idx, camera_name in enumerate(RING_CAMERA_LIST + STEREO_CAMERA_LIST): | ||
fps = fps_map[camera_name] | ||
if generate_video_only: | ||
mp4_path = f"{video_output_dir}/{log_id}_{camera_name}_{fps}fps.mp4" | ||
video_writer = VideoWriter(mp4_path) | ||
cam_im_fpaths = dl.get_ordered_log_cam_fpaths(log_id, camera_name) | ||
for i, im_fpath in enumerate(cam_im_fpaths): | ||
if i % 50 == 0: | ||
|
@@ -154,8 +175,8 @@ def dump_clipped_3d_cuboids_to_images( | |
break | ||
continue | ||
|
||
city_to_egovehicle_se3 = dl.get_city_to_egovehicle_se3(log_id, cam_timestamp) | ||
if city_to_egovehicle_se3 is None: | ||
city_SE3_egovehicle = dl.get_city_to_egovehicle_se3(log_id, cam_timestamp) | ||
if city_SE3_egovehicle is None: | ||
continue | ||
|
||
lidar_timestamp = Path(ply_fpath).stem.split("_")[-1] | ||
|
@@ -170,15 +191,16 @@ def dump_clipped_3d_cuboids_to_images( | |
img = imageio.imread(im_fpath)[:, :, ::-1].copy() | ||
camera_config = get_calibration_config(log_calib_data, camera_name) | ||
planes = generate_frustum_planes(camera_config.intrinsic.copy(), camera_name) | ||
img = plot_lane_centerlines_in_img( | ||
lidar_pts, | ||
city_to_egovehicle_se3, | ||
img, | ||
city_name, | ||
avm, | ||
camera_config, | ||
planes, | ||
) | ||
if not omit_centerlines: | ||
img = plot_lane_centerlines_in_img( | ||
lidar_pts, | ||
city_SE3_egovehicle, | ||
img, | ||
city_name, | ||
avm, | ||
camera_config, | ||
planes, | ||
) | ||
|
||
for label_idx, label in enumerate(labels): | ||
obj_rec = json_label_dict_to_obj_record(label) | ||
|
@@ -216,40 +238,64 @@ def dump_clipped_3d_cuboids_to_images( | |
copy.deepcopy(camera_config), | ||
) | ||
|
||
cv2.imwrite(save_img_fpath, img) | ||
saved_img_fpaths += [save_img_fpath] | ||
if max_num_images_to_render != -1 and len(saved_img_fpaths) > max_num_images_to_render: | ||
if generate_video_only: | ||
video_writer.add_frame(img[:, :, ::-1]) | ||
else: | ||
cv2.imwrite(save_img_fpath, img) | ||
saved_img_fpaths += [save_img_fpath] | ||
if ( | ||
not generate_video_only | ||
and max_num_images_to_render != -1 | ||
and len(saved_img_fpaths) > max_num_images_to_render | ||
): | ||
flag_done = True | ||
break | ||
if generate_video_only: | ||
video_writer.complete() | ||
ffmpeg_compress_video(mp4_path, fps) | ||
if flag_done: | ||
break | ||
category_subdir = "amodal_labels" | ||
|
||
if not Path(f"{experiment_prefix}_{category_subdir}").exists(): | ||
os.makedirs(f"{experiment_prefix}_{category_subdir}") | ||
|
||
for cam_idx, camera_name in enumerate(RING_CAMERA_LIST + STEREO_CAMERA_LIST): | ||
# Write the cuboid video -- could also write w/ fps=20,30,40 | ||
if "stereo" in camera_name: | ||
fps = 5 | ||
else: | ||
fps = 30 | ||
img_wildcard = f"{save_dir}/{camera_name}_%*.jpg" | ||
output_fpath = f"{experiment_prefix}_{category_subdir}/{log_id}_{camera_name}_{fps}fps.mp4" | ||
write_nonsequential_idx_video(img_wildcard, output_fpath, fps) | ||
if not generate_video_only: | ||
for cam_idx, camera_name in enumerate(RING_CAMERA_LIST + STEREO_CAMERA_LIST): | ||
# Write the cuboid video from individual frames -- could also write w/ fps=20,30,40 | ||
fps = fps_map[camera_name] | ||
img_wildcard = f"{save_dir}/{camera_name}_%*.jpg" | ||
output_fpath = f"{video_output_dir}/{log_id}_{camera_name}_{fps}fps.mp4" | ||
write_nonsequential_idx_video(img_wildcard, output_fpath, fps) | ||
|
||
return saved_img_fpaths | ||
|
||
|
||
def main(args: Any): | ||
"""Run the example.""" | ||
log_ids = [log_id.strip() for log_id in args.log_ids.split(",")] | ||
dump_clipped_3d_cuboids_to_images( | ||
log_ids, | ||
args.max_num_images_to_render * 9, | ||
args.dataset_dir, | ||
args.experiment_prefix, | ||
) | ||
if args.use_multiprocessing: | ||
single_process_args = [ | ||
( | ||
[log_id], | ||
args.max_num_images_to_render * 9, | ||
args.dataset_dir, | ||
args.experiment_prefix, | ||
not args.no_motion_compensation, | ||
args.omit_centerlines, | ||
args.generate_video_only, | ||
) | ||
for log_id in log_ids | ||
] | ||
with Pool(os.cpu_count()) as p: | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Since it take awhile to run, would be nice to be able to specify number of cpu to use as well so we can do something else while rendering. |
||
accum = p.starmap(dump_clipped_3d_cuboids_to_images, single_process_args) | ||
|
||
else: | ||
dump_clipped_3d_cuboids_to_images( | ||
log_ids, | ||
args.max_num_images_to_render * 9, | ||
args.dataset_dir, | ||
args.experiment_prefix, | ||
not args.no_motion_compensation, | ||
args.omit_centerlines, | ||
args.generate_video_only, | ||
) | ||
|
||
|
||
if __name__ == "__main__": | ||
|
@@ -262,6 +308,26 @@ def main(args: Any): | |
help="number of images within which to render 3d cuboids", | ||
) | ||
parser.add_argument("--dataset-dir", type=str, required=True, help="path to the dataset folder") | ||
parser.add_argument( | ||
"--use-multiprocessing", | ||
action="store_true", | ||
help="uses multiprocessing only if arg is specified on command line, otherwise single process", | ||
) | ||
parser.add_argument( | ||
"--no-motion-compensation", | ||
action="store_true", | ||
help="motion compensate by default, unless arg is specified on command line to not do so", | ||
) | ||
parser.add_argument( | ||
"--omit-centerlines", | ||
action="store_true", | ||
help="renders centerlines by default, will omit them if arg is specified on command line", | ||
) | ||
parser.add_argument( | ||
"--generate-video-only", | ||
action="store_true", | ||
help="produces mp4 files only, without dumping any individual frames/images to JPGs", | ||
) | ||
parser.add_argument( | ||
"--log-ids", | ||
type=str, | ||
|
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Does this replace this dependency? https://github.com/argoai/argoverse-api/blob/1dfc03ee8dc9b338eac6c86e21d7b999b59e2498/argoverse/visualization/generate_sequence_videos.py#L108
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I didn't realize we had this dependency tbh (it's not in our dependency list for some reason), I'll take a look
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
What were your thoughts on this?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
haven't played around with moviepy before, but we should check it out more: https://zulko.github.io/moviepy/