-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathRoad.py
154 lines (125 loc) · 5.89 KB
/
Road.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
import pygame
import cv2, os
import threading
import numpy as np
import tensorflow as tf # assuming your YOLO model is in TensorFlow
import time
import pygame
from constants import *
global minLineLength, maxLineGap
minLineLength = 20
maxLineGap = 4
class Road:
def __init__(self):
# Defining some sample road properties
self.road_types = ["asphalt", "gravel", "dirt", "snow"]
self.roughness = {
"asphalt": 0.1, # Roughness scale: 0 (smooth) to 1 (very rough)
"gravel": 0.6,
"dirt": 0.8,
"snow": 0.7
}
self.friction = {
"asphalt": 0.9, # Friction scale: 0 (no friction) to 1 (high friction)
"gravel": 0.5,
"dirt": 0.4,
"snow": 0.3
}
def get_road_properties(self, position):
"""
This function returns road properties based on a given position.
For simplicity, the road type changes every 100 units of position.
You can modify the logic to be more sophisticated if needed.
"""
# Based on position, determine which type of road it is
road_type = self.road_types[position // 100 % len(self.road_types)]
return {
"road_type": road_type,
"roughness": self.roughness[road_type],
"friction": self.friction[road_type]
}
@staticmethod
def lane_decider(surface):
pygame_image = pygame.surfarray.array3d(surface)
return pygame_image
def calculate_lane_center_and_slope(self, lines):
if lines is None:
return None, None
x_center_accum = 0
y_center_accum = 0
slopes = []
for line in lines:
x1, y1, x2, y2 = line.reshape(4)
x_center_accum += (x1 + x2) / 2
y_center_accum += (y1 + y2) / 2
# Calculate slope and avoid division by zero
if x2 - x1 != 0:
slopes.append((y2 - y1) / (x2 - x1))
# Calculate average center and slope
lane_center = (int(x_center_accum / len(lines)), int(y_center_accum / len(lines)))
avg_slope = sum(slopes) / len(slopes) if slopes else 0
return lane_center, avg_slope
@staticmethod
def pygame_surface_to_opencv_image(surface):
"""Convert a pygame surface to an OpenCV image."""
opencv_image = pygame.surfarray.array3d(surface).transpose([1, 0, 2])
return opencv_image
def detect_lanes_pygame(self, image, minLineLength=95, maxLineGap=100, canny_thresholds=(255, 255),solid_line_threshold=200.0):
road_image = self.pygame_surface_to_opencv_image(image)
gray = cv2.cvtColor(road_image[:, :, ::-1], cv2.COLOR_BGR2GRAY)
contrasted = cv2.convertScaleAbs(gray, alpha=1.5, beta=0)
blurred = cv2.GaussianBlur(contrasted, (5, 5), 0)
canny = cv2.Canny(blurred, *canny_thresholds)
kernel = np.ones((3, 3), np.uint8)
dilated = cv2.dilate(canny, kernel, iterations=1)
height, width = road_image.shape[:2]
rectangle_top_left = (-50, height//2-30)
rectangle_bottom_right = (width+150, 2*height//3-20)
rectangle = np.array([[rectangle_top_left, (rectangle_bottom_right[0], rectangle_top_left[1]), rectangle_bottom_right, (rectangle_top_left[0], rectangle_bottom_right[1])]])
polygons = np.array(rectangle, dtype=np.int32)
mask = np.zeros_like(dilated)
cv2.fillPoly(mask, polygons, 255)
masked_image = cv2.bitwise_and(dilated, mask) # Use dilated here
lines = cv2.HoughLinesP(masked_image, 2, np.pi/180, 100, np.array([]), minLineLength=minLineLength, maxLineGap=maxLineGap)
line_image = np.zeros(road_image.shape, dtype=road_image.dtype)
if lines is not None:
for line in lines:
x1, y1, x2, y2 = line.reshape(4)
length = ((x2 - x1)**2 + (y2 - y1)**2)**0.5
if length > solid_line_threshold: # Solid line
cv2.line(line_image, (x1, y1), (x2, y2), (0, 255, 0), 10)
else: # Dashed line
cv2.line(line_image, (x1, y1), (x2, y2), (0, 0, 255), 10)
combined_image = cv2.addWeighted(road_image, 0.8, line_image, 1, 1)
# Draw the ROI polygon for debugging purposes
cv2.polylines(combined_image, [polygons], isClosed=True, color=(255, 255, 0), thickness=2)
return lines, combined_image
def car_position_in_lane(self,lines, car_image_np, combined_image):
lane_center,avg_slope=self.calculate_lane_center_and_slope(lines)
# If we have a valid center, draw the car there.
if lane_center:
# car_image_np = self.pygame_surface_to_opencv_image(car_scaled)
car_height, car_width = car_image_np.shape[:2]
top_left_corner = (lane_center[0] - car_width // 2, lane_center[1] - car_height // 2)
combined_image[top_left_corner[1]: top_left_corner[1] + car_height, top_left_corner[0]: top_left_corner[0] + car_width] = car_image_np
return combined_image
if __name__=="__main__":
current_dir = os.path.dirname(os.path.abspath(__file__))
assets_path = os.path.join(current_dir, "assets\\")
rd = Road()
pygame_image = pygame.image.load(assets_path + "city\\" + "4Wayroad800x600.png")
while True: # Loop until a break condition is met
if pygame_image.get_size() != (800, 600):
scaled_pygame_image = pygame.transform.scale(pygame_image, (WIDTH, HEIGHT))
com_img = rd.detect_lanes_pygame(scaled_pygame_image)
else:
com_img = rd.detect_lanes_pygame(pygame_image)
cv2.imshow('Image', com_img)
key = cv2.waitKey(1) & 0xFF
if key == 27: # Escape key to break
break
cv2.destroyAllWindows()
# Example:
# road = Road()
# position = 150 # Sample position
# print(road.get_road_properties(position))