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()