74 lines
3.1 KiB
Python
74 lines
3.1 KiB
Python
import cv2
|
|
import numpy as np
|
|
import math
|
|
from sklearn import linear_model
|
|
|
|
def lineCalc(vx, vy, x0, y0):
|
|
scale = 10
|
|
x1 = x0 + scale * vx
|
|
y1 = y0 + scale * vy
|
|
m = (y1 - y0) / (x1 - x0)
|
|
b = y1 - m * x1
|
|
return m, b
|
|
|
|
def lineIntersect(m1, b1, m2, b2):
|
|
a_1 = -m1
|
|
b_1 = 1
|
|
c_1 = b1
|
|
a_2 = -m2
|
|
b_2 = 1
|
|
c_2 = b2
|
|
d = a_1 * b_2 - a_2 * b_1
|
|
dx = c_1 * b_2 - c_2 * b_1
|
|
dy = a_1 * c_2 - a_2 * c_1
|
|
intersectionX = dx / d
|
|
intersectionY = dy / d
|
|
return intersectionX, intersectionY
|
|
|
|
def detect_crosswalk(frame):
|
|
'''Detects crosswalk/zebra lines and vanishing point in a BGR frame.'''
|
|
H, W = frame.shape[:2]
|
|
radius = 250
|
|
bw_width = 170
|
|
lower = np.array([170, 170, 170])
|
|
upper = np.array([255, 255, 255])
|
|
mask = cv2.inRange(frame, lower, upper)
|
|
erodeSize = int(H / 30)
|
|
erodeStructure = cv2.getStructuringElement(cv2.MORPH_RECT, (erodeSize, 1))
|
|
erode = cv2.erode(mask, erodeStructure, (-1, -1))
|
|
contours, _ = cv2.findContours(erode, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
|
|
bxbyLeftArray, bxbyRightArray = [], []
|
|
for cnt in contours:
|
|
bx, by, bw, bh = cv2.boundingRect(cnt)
|
|
if bw > bw_width:
|
|
cv2.line(frame, (bx, by), (bx + bw, by), (0, 255, 0), 2)
|
|
bxbyLeftArray.append([bx, by])
|
|
bxbyRightArray.append([bx + bw, by])
|
|
cv2.circle(frame, (int(bx), int(by)), 5, (0, 250, 250), 2)
|
|
cv2.circle(frame, (int(bx + bw), int(by)), 5, (250, 250, 0), 2)
|
|
if len(bxbyLeftArray) < 2 or len(bxbyRightArray) < 2:
|
|
return None, None, frame
|
|
medianL = np.median(bxbyLeftArray, axis=0)
|
|
medianR = np.median(bxbyRightArray, axis=0)
|
|
boundedLeft = [i for i in bxbyLeftArray if ((medianL[0] - i[0]) ** 2 + (medianL[1] - i[1]) ** 2) < radius ** 2]
|
|
boundedRight = [i for i in bxbyRightArray if ((medianR[0] - i[0]) ** 2 + (medianR[1] - i[1]) ** 2) < radius ** 2]
|
|
if len(boundedLeft) < 2 or len(boundedRight) < 2:
|
|
return None, None, frame
|
|
bxLeft = np.asarray([pt[0] for pt in boundedLeft]).reshape(-1, 1)
|
|
byLeft = np.asarray([pt[1] for pt in boundedLeft])
|
|
bxRight = np.asarray([pt[0] for pt in boundedRight]).reshape(-1, 1)
|
|
byRight = np.asarray([pt[1] for pt in boundedRight])
|
|
modelL = linear_model.RANSACRegressor().fit(bxLeft, byLeft)
|
|
modelR = linear_model.RANSACRegressor().fit(bxRight, byRight)
|
|
vx, vy, x0, y0 = cv2.fitLine(np.array(boundedLeft), cv2.DIST_L2, 0, 0.01, 0.01)
|
|
vx_R, vy_R, x0_R, y0_R = cv2.fitLine(np.array(boundedRight), cv2.DIST_L2, 0, 0.01, 0.01)
|
|
m_L, b_L = lineCalc(vx, vy, x0, y0)
|
|
m_R, b_R = lineCalc(vx_R, vy_R, x0_R, y0_R)
|
|
intersectionX, intersectionY = lineIntersect(m_R, b_R, m_L, b_L)
|
|
m = radius * 10
|
|
if intersectionY < H / 2:
|
|
cv2.circle(frame, (int(intersectionX), int(intersectionY)), 10, (0, 0, 255), 15)
|
|
cv2.line(frame, (int(x0 - m * vx), int(y0 - m * vy)), (int(x0 + m * vx), int(y0 + m * vy)), (255, 0, 0), 3)
|
|
cv2.line(frame, (int(x0_R - m * vx_R), int(y0_R - m * vy_R)), (int(x0_R + m * vx_R), int(y0_R + m * vy_R)), (255, 0, 0), 3)
|
|
return (int(intersectionX), int(intersectionY)), [list(medianL) + list(medianR)], frame
|