112 lines
3.0 KiB
Python
112 lines
3.0 KiB
Python
import os
|
|
import cv2
|
|
import numpy as np
|
|
from scipy.interpolate import InterpolatedUnivariateSpline
|
|
|
|
|
|
class Lane:
|
|
def __init__(self, points=None, invalid_value=-2., metadata=None):
|
|
super(Lane, self).__init__()
|
|
self.curr_iter = 0
|
|
self.points = points
|
|
self.invalid_value = invalid_value
|
|
self.function = InterpolatedUnivariateSpline(
|
|
points[:, 1], points[:, 0], k=min(3, len(points) - 1))
|
|
self.min_y = points[:, 1].min() - 0.01
|
|
self.max_y = points[:, 1].max() + 0.01
|
|
self.metadata = metadata or {}
|
|
|
|
def __repr__(self):
|
|
return '[Lane]\n' + str(self.points) + '\n[/Lane]'
|
|
|
|
def __call__(self, lane_ys):
|
|
lane_xs = self.function(lane_ys)
|
|
|
|
lane_xs[(lane_ys < self.min_y) | (lane_ys > self.max_y
|
|
)] = self.invalid_value
|
|
return lane_xs
|
|
|
|
def to_array(self, sample_y_range, img_w, img_h):
|
|
self.sample_y = range(sample_y_range[0], sample_y_range[1],
|
|
sample_y_range[2])
|
|
sample_y = self.sample_y
|
|
img_w, img_h = img_w, img_h
|
|
ys = np.array(sample_y) / float(img_h)
|
|
xs = self(ys)
|
|
valid_mask = (xs >= 0) & (xs < 1)
|
|
lane_xs = xs[valid_mask] * img_w
|
|
lane_ys = ys[valid_mask] * img_h
|
|
lane = np.concatenate(
|
|
(lane_xs.reshape(-1, 1), lane_ys.reshape(-1, 1)), axis=1)
|
|
return lane
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def __next__(self):
|
|
if self.curr_iter < len(self.points):
|
|
self.curr_iter += 1
|
|
return self.points[self.curr_iter - 1]
|
|
self.curr_iter = 0
|
|
raise StopIteration
|
|
|
|
|
|
COLORS = [
|
|
(255, 0, 0),
|
|
(0, 255, 0),
|
|
(0, 0, 255),
|
|
(255, 255, 0),
|
|
(255, 0, 255),
|
|
(0, 255, 255),
|
|
(128, 255, 0),
|
|
(255, 128, 0),
|
|
(128, 0, 255),
|
|
(255, 0, 128),
|
|
(0, 128, 255),
|
|
(0, 255, 128),
|
|
(128, 255, 255),
|
|
(255, 128, 255),
|
|
(255, 255, 128),
|
|
(60, 180, 0),
|
|
(180, 60, 0),
|
|
(0, 60, 180),
|
|
(0, 180, 60),
|
|
(60, 0, 180),
|
|
(180, 0, 60),
|
|
(255, 0, 0),
|
|
(0, 255, 0),
|
|
(0, 0, 255),
|
|
(255, 255, 0),
|
|
(255, 0, 255),
|
|
(0, 255, 255),
|
|
(128, 255, 0),
|
|
(255, 128, 0),
|
|
(128, 0, 255),
|
|
]
|
|
|
|
|
|
def imshow_lanes(img, lanes, show=False, out_file=None, width=4):
|
|
lanes_xys = []
|
|
for _, lane in enumerate(lanes):
|
|
xys = []
|
|
for x, y in lane:
|
|
if x <= 0 or y <= 0:
|
|
continue
|
|
x, y = int(x), int(y)
|
|
xys.append((x, y))
|
|
lanes_xys.append(xys)
|
|
lanes_xys.sort(key=lambda xys: xys[0][0] if len(xys) > 0 else 0)
|
|
|
|
for idx, xys in enumerate(lanes_xys):
|
|
for i in range(1, len(xys)):
|
|
cv2.line(img, xys[i - 1], xys[i], COLORS[idx], thickness=width)
|
|
|
|
if show:
|
|
cv2.imshow('view', img)
|
|
cv2.waitKey(0)
|
|
|
|
if out_file:
|
|
if not os.path.exists(os.path.dirname(out_file)):
|
|
os.makedirs(os.path.dirname(out_file))
|
|
cv2.imwrite(out_file, img)
|