131 lines
4.3 KiB
Python
131 lines
4.3 KiB
Python
import math
|
|
import numpy as np
|
|
from imgaug.augmentables.lines import LineString
|
|
from scipy.interpolate import InterpolatedUnivariateSpline
|
|
|
|
|
|
def lane_to_linestrings(lanes):
|
|
lines = []
|
|
for lane in lanes:
|
|
lines.append(LineString(lane))
|
|
|
|
return lines
|
|
|
|
|
|
def linestrings_to_lanes(lines):
|
|
lanes = []
|
|
for line in lines:
|
|
lanes.append(line.coords)
|
|
|
|
return lanes
|
|
|
|
|
|
def sample_lane(points, sample_ys, img_w):
|
|
# this function expects the points to be sorted
|
|
points = np.array(points)
|
|
if not np.all(points[1:, 1] < points[:-1, 1]):
|
|
raise Exception('Annotaion points have to be sorted')
|
|
x, y = points[:, 0], points[:, 1]
|
|
|
|
# interpolate points inside domain
|
|
assert len(points) > 1
|
|
interp = InterpolatedUnivariateSpline(
|
|
y[::-1], x[::-1], k=min(3, len(points) - 1))
|
|
domain_min_y = y.min()
|
|
domain_max_y = y.max()
|
|
sample_ys_inside_domain = sample_ys[(sample_ys >= domain_min_y) & (
|
|
sample_ys <= domain_max_y)]
|
|
assert len(sample_ys_inside_domain) > 0
|
|
interp_xs = interp(sample_ys_inside_domain)
|
|
|
|
# extrapolate lane to the bottom of the image with a straight line using the 2 points closest to the bottom
|
|
two_closest_points = points[:2]
|
|
extrap = np.polyfit(
|
|
two_closest_points[:, 1], two_closest_points[:, 0], deg=1)
|
|
extrap_ys = sample_ys[sample_ys > domain_max_y]
|
|
extrap_xs = np.polyval(extrap, extrap_ys)
|
|
all_xs = np.hstack((extrap_xs, interp_xs))
|
|
|
|
# separate between inside and outside points
|
|
inside_mask = (all_xs >= 0) & (all_xs < img_w)
|
|
xs_inside_image = all_xs[inside_mask]
|
|
xs_outside_image = all_xs[~inside_mask]
|
|
|
|
return xs_outside_image, xs_inside_image
|
|
|
|
|
|
def filter_lane(lane):
|
|
assert lane[-1][1] <= lane[0][1]
|
|
filtered_lane = []
|
|
used = set()
|
|
for p in lane:
|
|
if p[1] not in used:
|
|
filtered_lane.append(p)
|
|
used.add(p[1])
|
|
|
|
return filtered_lane
|
|
|
|
|
|
def transform_annotation(img_w, img_h, max_lanes, n_offsets, offsets_ys,
|
|
n_strips, strip_size, anno):
|
|
old_lanes = anno['lanes']
|
|
|
|
# removing lanes with less than 2 points
|
|
old_lanes = filter(lambda x: len(x) > 1, old_lanes)
|
|
# sort lane points by Y (bottom to top of the image)
|
|
old_lanes = [sorted(lane, key=lambda x: -x[1]) for lane in old_lanes]
|
|
# remove points with same Y (keep first occurrence)
|
|
old_lanes = [filter_lane(lane) for lane in old_lanes]
|
|
# normalize the annotation coordinates
|
|
old_lanes = [[[x * img_w / float(img_w), y * img_h / float(img_h)]
|
|
for x, y in lane] for lane in old_lanes]
|
|
# create tranformed annotations
|
|
lanes = np.ones(
|
|
(max_lanes, 2 + 1 + 1 + 2 + n_offsets), dtype=np.float32
|
|
) * -1e5 # 2 scores, 1 start_y, 1 start_x, 1 theta, 1 length, S+1 coordinates
|
|
lanes_endpoints = np.ones((max_lanes, 2))
|
|
# lanes are invalid by default
|
|
lanes[:, 0] = 1
|
|
lanes[:, 1] = 0
|
|
for lane_idx, lane in enumerate(old_lanes):
|
|
if lane_idx >= max_lanes:
|
|
break
|
|
|
|
try:
|
|
xs_outside_image, xs_inside_image = sample_lane(lane, offsets_ys,
|
|
img_w)
|
|
except AssertionError:
|
|
continue
|
|
if len(xs_inside_image) <= 1:
|
|
continue
|
|
all_xs = np.hstack((xs_outside_image, xs_inside_image))
|
|
lanes[lane_idx, 0] = 0
|
|
lanes[lane_idx, 1] = 1
|
|
lanes[lane_idx, 2] = len(xs_outside_image) / n_strips
|
|
lanes[lane_idx, 3] = xs_inside_image[0]
|
|
|
|
thetas = []
|
|
for i in range(1, len(xs_inside_image)):
|
|
theta = math.atan(
|
|
i * strip_size /
|
|
(xs_inside_image[i] - xs_inside_image[0] + 1e-5)) / math.pi
|
|
theta = theta if theta > 0 else 1 - abs(theta)
|
|
thetas.append(theta)
|
|
|
|
theta_far = sum(thetas) / len(thetas)
|
|
|
|
# lanes[lane_idx,
|
|
# 4] = (theta_closest + theta_far) / 2 # averaged angle
|
|
lanes[lane_idx, 4] = theta_far
|
|
lanes[lane_idx, 5] = len(xs_inside_image)
|
|
lanes[lane_idx, 6:6 + len(all_xs)] = all_xs
|
|
lanes_endpoints[lane_idx, 0] = (len(all_xs) - 1) / n_strips
|
|
lanes_endpoints[lane_idx, 1] = xs_inside_image[-1]
|
|
|
|
new_anno = {
|
|
'label': lanes,
|
|
'old_anno': anno,
|
|
'lane_endpoints': lanes_endpoints
|
|
}
|
|
return new_anno
|