57 lines
1.9 KiB
Python
57 lines
1.9 KiB
Python
import cv2
|
|
import numpy as np
|
|
|
|
def calculate_angle(left_line, right_line, width):
|
|
left_slope = np.inf if left_line[0][0] == left_line[1][0] else (left_line[1][1] - left_line[0][1]) / (left_line[1][0] - left_line[0][0])
|
|
right_slope = np.inf if right_line[0][0] == right_line[1][0] else (right_line[1][1] - right_line[0][1]) / (right_line[1][0] - right_line[0][0])
|
|
|
|
center_x = width // 2
|
|
left_y = int(left_slope * (center_x - left_line[0][0]) + left_line[0][1])
|
|
right_y = int(right_slope * (center_x - right_line[0][0]) + right_line[0][1])
|
|
|
|
cv2.line(frame, (center_x, 0), (center_x, 100), (255, 0, 0), 2)
|
|
cv2.line(frame, (center_x, left_y), (center_x, right_y), (0, 255, 0), 2)
|
|
|
|
angle_radians = np.arctan((right_y - left_y) / (center_x))
|
|
angle_degrees = np.degrees(angle_radians)
|
|
|
|
return angle_degrees
|
|
|
|
cap = cv2.VideoCapture(1)
|
|
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
|
edges = cv2.Canny(gray, 50, 150)
|
|
|
|
lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=50, minLineLength=50, maxLineGap=30)
|
|
|
|
left_lines = []
|
|
right_lines = []
|
|
|
|
if lines is not None:
|
|
for line in lines:
|
|
x1, y1, x2, y2 = line[0]
|
|
|
|
if x1 < frame.shape[1] // 2 and x2 < frame.shape[1] // 2:
|
|
left_lines.append([[x1, y1], [x2, y2]])
|
|
elif x1 > frame.shape[1] // 2 and x2 > frame.shape[1] // 2:
|
|
right_lines.append([[x1, y1], [x2, y2]])
|
|
|
|
if left_lines and right_lines:
|
|
left_line = np.mean(left_lines, axis=0)
|
|
right_line = np.mean(right_lines, axis=0)
|
|
|
|
angle = calculate_angle(left_line.squeeze(), right_line.squeeze(), frame.shape[1])
|
|
print(angle)
|
|
|
|
cv2.imshow('Video', frame)
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
break
|
|
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|