69 lines
1.7 KiB
Python
69 lines
1.7 KiB
Python
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
def region_of_interest(img, vertices):
|
||
|
|
mask = np.zeros_like(img)
|
||
|
|
cv2.fillPoly(mask, vertices, 255)
|
||
|
|
masked_image = cv2.bitwise_and(img, mask)
|
||
|
|
return masked_image
|
||
|
|
|
||
|
|
def draw_lines(img, lines):
|
||
|
|
if lines is None:
|
||
|
|
return
|
||
|
|
img = np.copy(img)
|
||
|
|
blank_image = np.zeros((img.shape[0], img.shape[1], 3), dtype=np.uint8)
|
||
|
|
|
||
|
|
for line in lines:
|
||
|
|
for x1, y1, x2, y2 in line:
|
||
|
|
cv2.line(blank_image, (x1, y1), (x2, y2), (0, 255, 0), thickness=5)
|
||
|
|
|
||
|
|
img = cv2.addWeighted(img, 0.8, blank_image, 1, 0.0)
|
||
|
|
return img
|
||
|
|
|
||
|
|
def process(image):
|
||
|
|
height = image.shape[0]
|
||
|
|
width = image.shape[1]
|
||
|
|
|
||
|
|
region_of_interest_vertices = [
|
||
|
|
(0, height),
|
||
|
|
(width // 2, height // 2),
|
||
|
|
(width, height)
|
||
|
|
]
|
||
|
|
|
||
|
|
gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||
|
|
canny_image = cv2.Canny(gray_image, 100, 200)
|
||
|
|
cropped_image = region_of_interest(canny_image,
|
||
|
|
np.array([region_of_interest_vertices], np.int32))
|
||
|
|
|
||
|
|
lines = cv2.HoughLinesP(cropped_image,
|
||
|
|
rho=2,
|
||
|
|
theta=np.pi / 180,
|
||
|
|
threshold=50,
|
||
|
|
lines=np.array([]),
|
||
|
|
minLineLength=40,
|
||
|
|
maxLineGap=100)
|
||
|
|
|
||
|
|
image_with_lines = draw_lines(image, lines)
|
||
|
|
return image_with_lines
|
||
|
|
|
||
|
|
def main():
|
||
|
|
cap = cv2.VideoCapture(0)
|
||
|
|
|
||
|
|
while cap.isOpened():
|
||
|
|
ret, frame = cap.read()
|
||
|
|
if not ret:
|
||
|
|
break
|
||
|
|
|
||
|
|
frame = process(frame)
|
||
|
|
|
||
|
|
cv2.imshow('frame', frame)
|
||
|
|
|
||
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
||
|
|
break
|
||
|
|
|
||
|
|
cap.release()
|
||
|
|
cv2.destroyAllWindows()
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|