Spaces:
Sleeping
Sleeping
File size: 4,646 Bytes
c897073 |
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 |
import cv2
import numpy as np
from ultralytics import YOLO
import gradio as gr
import tempfile
import os
model = YOLO("yolov8n.pt")
def detect_lanes(image):
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
blurred = cv2.GaussianBlur(gray, (5, 5), 0)
edges = cv2.Canny(blurred, 50, 150)
height, width = edges.shape
roi = np.array([[
(int(0.2 * width), height),
(int(0.8 * width), height),
(int(0.6 * width), int(0.6 * height)),
(int(0.4 * width), int(0.6 * height))
]], dtype=np.int32)
mask = np.zeros_like(edges)
cv2.fillPoly(mask, roi, 255)
masked_edges = cv2.bitwise_and(edges, mask)
# Use Hough Line Transform to detect lines
lines = cv2.HoughLinesP(
masked_edges,
rho=1,
theta=np.pi / 180,
threshold=50,
minLineLength=50,
maxLineGap=200
)
# Create an image to draw the lines on
line_image = np.zeros_like(image)
# Filter and average the lines for left and right lanes
left_lines = []
right_lines = []
if lines is not None:
for line in lines:
x1, y1, x2, y2 = line[0]
slope = (y2 - y1) / (x2 - x1) if x2 != x1 else 0
if abs(slope) < 0.5: # Ignore near-horizontal lines
continue
if slope < 0: # Left lane
left_lines.append((x1, y1, x2, y2))
else: # Right lane
right_lines.append((x1, y1, x2, y2))
def average_lines(lines):
if len(lines) == 0:
return None
x_coords = []
y_coords = []
for line in lines:
x_coords += [line[0], line[2]]
y_coords += [line[1], line[3]]
poly_fit = np.polyfit(y_coords, x_coords, 1)
y1 = height
y2 = int(height * 0.6)
x1 = int(np.polyval(poly_fit, y1))
x2 = int(np.polyval(poly_fit, y2))
return x1, y1, x2, y2
left_lane = average_lines(left_lines)
right_lane = average_lines(right_lines)
# Draw the lanes
if left_lane:
cv2.line(line_image, (left_lane[0], left_lane[1]), (left_lane[2], left_lane[3]), (0, 255, 0), 5)
if right_lane:
cv2.line(line_image, (right_lane[0], right_lane[1]), (right_lane[2], right_lane[3]), (0, 255, 0), 5)
# Overlay the lines on the original image
result = cv2.addWeighted(image, 0.8, line_image, 1, 0)
return result
# YOLO Object Detection Function
def detect_objects(image):
results = model(image)
annotated_frame = results[0].plot() # Annotate the frame with bounding boxes
return annotated_frame
def process_image(image):
image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
detected_objects = detect_objects(image)
lane_detected = detect_lanes(image)
return cv2.cvtColor(detected_objects, cv2.COLOR_BGR2RGB), cv2.cvtColor(lane_detected, cv2.COLOR_BGR2RGB)
def process_video(video):
cap = cv2.VideoCapture(video.name)
output_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(cap.get(cv2.CAP_PROP_FPS))
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
detected_objects = detect_objects(frame)
# Lane detection
lane_detected = detect_lanes(frame)
# Combine detections and lane detection
combined_frame = cv2.addWeighted(detected_objects, 0.5, lane_detected, 0.5, 0)
out.write(combined_frame)
cap.release()
out.release()
return output_path
with gr.Blocks() as demo:
gr.Markdown("# YOLO Object Detection and Lane Detection")
with gr.Tab("Image Upload"):
image_input = gr.Image(label="Upload Image", type="pil")
object_output = gr.Image(label="Object Detection Output")
lane_output = gr.Image(label="Lane Detection Output")
image_button = gr.Button("Process Image")
image_button.click(process_image, inputs=image_input, outputs=[object_output, lane_output])
with gr.Tab("Video Upload"):
video_input = gr.Video(label="Upload Video")
video_output = gr.Video(label="Processed Video")
video_button = gr.Button("Process Video")
video_button.click(process_video, inputs=video_input, outputs=video_output)
demo.launch()
|