Skip to content
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

Speedup semantic rasterizer #140

Open
wants to merge 5 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 8 additions & 4 deletions l5kit/l5kit/data/map_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -110,7 +110,7 @@ def get_lane_coords(self, element_id: str) -> dict:
element_id (str): lane element id

Returns:
dict: a dict with the two boundaries coordinates as (Nx3) XYZ arrays
dict: a dict with the boundaries coordinates as (3xN) XYZ arrays
"""
element = self[element_id]
assert self.is_lane(element)
Expand All @@ -132,7 +132,10 @@ def get_lane_coords(self, element_id: str) -> dict:
lane.geo_frame,
)

return {"xyz_left": xyz_left, "xyz_right": xyz_right}
xyz = np.vstack((xyz_left, np.flip(xyz_right, 0)))
xyz[:, -1] = 1.0

return {"xyz": xyz.T}

@staticmethod
@no_type_check
Expand Down Expand Up @@ -161,7 +164,7 @@ def get_crosswalk_coords(self, element_id: str) -> dict:
element_id (str): crosswalk element id

Returns:
dict: a dict with the polygon coordinates as an (Nx3) XYZ array
dict: a dict with the polygon coordinates as an (3xN) XYZ array
"""
element = self[element_id]
assert self.is_crosswalk(element)
Expand All @@ -174,7 +177,8 @@ def get_crosswalk_coords(self, element_id: str) -> dict:
traffic_element.geo_frame,
)

return {"xyz": xyz}
xyz[:, -1] = 1.0
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why is this required? We ignore the z coordinates in the following so this shouldn't make any difference right?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the current implementation, first you cut the z-coordinate lane_coords["xyz_right"][:, :2] and later in the transform method you stack it back np.vstack((points[:num_dims, :], np.ones(points.shape[1]))). With this we can save the one cut, the np.ones creation, the stack.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Another way of doing this would be to only use the first 2x2 of the matrix (only XY) in the semantic rasterizer. My issue with setting 1 here is that we're removing information in a very hidden part of the code, which may render debugging problematic in the future (also considering there is a cache system in between)

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see. You can move it to the semantic rasterizer method; before the dot product calls.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like this idea also :) Today was quite busy, but I should be able to work on this tomorrow (hopefully)

return {"xyz": xyz.T}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can see the point of caching this to avoid stack at runtime, but we may still want to include either the two lanes or at least the length of the first, so that we can always unpack the two apart (I'm thinking about centre line support right now)

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The cheapest/easiest solution if you add the length of the first line. I think map_api should handle the centerline calculation as well (cachable).

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah, I agree on that and I'm fine with having the length included

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One note: I changed the returned xyz format to xyz.T.


def is_traffic_face_colour(self, element_id: str, colour: str) -> bool:
"""
Expand Down
28 changes: 13 additions & 15 deletions l5kit/l5kit/rasterization/semantic_rasterizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

from ..data.filter import filter_tl_faces_by_status
from ..data.map_api import MapAPI
from ..geometry import rotation33_as_yaw, transform_point, transform_points, world_to_image_pixels_matrix
from ..geometry import rotation33_as_yaw, transform_point, world_to_image_pixels_matrix
from .rasterizer import Rasterizer

# sub-pixel drawing precision constants
Expand Down Expand Up @@ -95,20 +95,20 @@ def get_bounds(self) -> dict:

if self.proto_API.is_lane(element):
lane = self.proto_API.get_lane_coords(element_id)
x_min = min(np.min(lane["xyz_left"][:, 0]), np.min(lane["xyz_right"][:, 0]))
y_min = min(np.min(lane["xyz_left"][:, 1]), np.min(lane["xyz_right"][:, 1]))
x_max = max(np.max(lane["xyz_left"][:, 0]), np.max(lane["xyz_right"][:, 0]))
y_max = max(np.max(lane["xyz_left"][:, 1]), np.max(lane["xyz_right"][:, 1]))
x_min = np.min(lane["xyz"][0, :])
y_min = np.min(lane["xyz"][1, :])
x_max = np.max(lane["xyz"][0, :])
y_max = np.max(lane["xyz"][1, :])

lanes_bounds = np.append(lanes_bounds, np.asarray([[[x_min, y_min], [x_max, y_max]]]), axis=0)
lanes_ids.append(element_id)

if self.proto_API.is_crosswalk(element):
crosswalk = self.proto_API.get_crosswalk_coords(element_id)
x_min = np.min(crosswalk["xyz"][:, 0])
y_min = np.min(crosswalk["xyz"][:, 1])
x_max = np.max(crosswalk["xyz"][:, 0])
y_max = np.max(crosswalk["xyz"][:, 1])
x_min = np.min(crosswalk["xyz"][0, :])
y_min = np.min(crosswalk["xyz"][1, :])
x_max = np.max(crosswalk["xyz"][0, :])
y_max = np.max(crosswalk["xyz"][1, :])

crosswalks_bounds = np.append(
crosswalks_bounds, np.asarray([[[x_min, y_min], [x_max, y_max]]]), axis=0,
Expand Down Expand Up @@ -176,12 +176,10 @@ def render_semantic_map(

# get image coords
lane_coords = self.proto_API.get_lane_coords(self.bounds_info["lanes"]["ids"][idx])
xy_left = cv2_subpixel(transform_points(lane_coords["xyz_left"][:, :2], world_to_image_space))
xy_right = cv2_subpixel(transform_points(lane_coords["xyz_right"][:, :2], world_to_image_space))
lanes_area = np.vstack((xy_left, np.flip(xy_right, 0))) # start->end left then end->start right
lanes_xy = cv2_subpixel(world_to_image_space.dot(lane_coords["xyz"]).T[:, :2])
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not using transform or transform_transpose functions here?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seemed a bit complicated. Transposes, vstack, np.ones, cut z-axis, etc. I did not want to break other parts of the code by modifying those methods.

(Feel free to correct anything you want.)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah you're perfectly right about that. Let me see if we can also handle this case there (same len for matrix and points)

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I did these modifications only to save as many CPU cycles as I possibly can. I did not consider long term usability, I only care about speeding up for the competition, so it is fine if you drop all of or part of this PR.


# Note(lberg): this called on all polygons skips some of them, don't know why
cv2.fillPoly(img, [lanes_area], (17, 17, 31), lineType=cv2.LINE_AA, shift=CV2_SHIFT)
cv2.fillPoly(img, [lanes_xy], (17, 17, 31), lineType=cv2.LINE_AA, shift=CV2_SHIFT)

lane_type = "default" # no traffic light face is controlling this lane
lane_tl_ids = set([MapAPI.id_as_str(la_tc) for la_tc in lane.traffic_controls])
Expand All @@ -193,7 +191,7 @@ def render_semantic_map(
elif self.proto_API.is_traffic_face_colour(tl_id, "yellow"):
lane_type = "yellow"

lanes_lines[lane_type].extend([xy_left, xy_right])
lanes_lines[lane_type].extend([lanes_xy])

cv2.polylines(img, lanes_lines["default"], False, (255, 217, 82), lineType=cv2.LINE_AA, shift=CV2_SHIFT)
cv2.polylines(img, lanes_lines["green"], False, (0, 255, 0), lineType=cv2.LINE_AA, shift=CV2_SHIFT)
Expand All @@ -205,7 +203,7 @@ def render_semantic_map(
for idx in elements_within_bounds(center_world, self.bounds_info["crosswalks"]["bounds"], raster_radius):
crosswalk = self.proto_API.get_crosswalk_coords(self.bounds_info["crosswalks"]["ids"][idx])

xy_cross = cv2_subpixel(transform_points(crosswalk["xyz"][:, :2], world_to_image_space))
xy_cross = cv2_subpixel(world_to_image_space.dot(crosswalk["xyz"]).T[:, :2])
crosswalks.append(xy_cross)

cv2.polylines(img, crosswalks, True, (255, 117, 69), lineType=cv2.LINE_AA, shift=CV2_SHIFT)
Expand Down