|
|
import gradio as gr |
|
|
import cv2 |
|
|
import numpy as np |
|
|
from ultralytics import YOLO |
|
|
import os |
|
|
from datetime import datetime |
|
|
import matplotlib.pyplot as plt |
|
|
from utils import create_percentage_circle, generate_pdf_report |
|
|
from reportlab.lib.pagesizes import letter |
|
|
from reportlab.pdfgen import canvas |
|
|
|
|
|
|
|
|
for directory in ["screenshots", "logs", "reports"]: |
|
|
try: |
|
|
os.makedirs(directory, exist_ok=True) |
|
|
except Exception as e: |
|
|
print(f"Warning: Could not create {directory} directory: {e}") |
|
|
|
|
|
|
|
|
model = YOLO("yolov8m.pt") |
|
|
|
|
|
def analyze_video(video_path, confidence=0.3): |
|
|
""" |
|
|
Process video, detect objects, log results, capture screenshots, and generate PDF report. |
|
|
Args: |
|
|
video_path: Path to input video. |
|
|
confidence: Confidence threshold for detections. |
|
|
Returns: |
|
|
Tuple of (annotated video path, summary text, screenshot paths, percentage circle path, PDF path). |
|
|
""" |
|
|
if not video_path: |
|
|
return None, "Error: No video uploaded.", [], None, None |
|
|
|
|
|
cap = cv2.VideoCapture(video_path) |
|
|
if not cap.isOpened(): |
|
|
return None, "Error: Could not open video.", [], None, None |
|
|
|
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) |
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
|
|
fps = int(cap.get(cv2.CAP_PROP_FPS)) |
|
|
|
|
|
output_path = "output_annotated.mp4" |
|
|
fourcc = cv2.VideoWriter_fourcc(*"mp4v") |
|
|
out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) |
|
|
|
|
|
log_file = f"logs/analysis_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt" |
|
|
vehicle_count = 0 |
|
|
pedestrian_count = 0 |
|
|
sign_count = 0 |
|
|
alerts = [] |
|
|
screenshots = [] |
|
|
confidences = [] |
|
|
|
|
|
with open(log_file, "w") as log: |
|
|
log.write(f"Analysis started at {datetime.now()}\n") |
|
|
log.write(f"Video: {video_path}\n") |
|
|
log.write(f"Confidence threshold: {confidence}\n\n") |
|
|
|
|
|
frame_count = 0 |
|
|
while cap.isOpened(): |
|
|
ret, frame = cap.read() |
|
|
if not ret: |
|
|
break |
|
|
frame_count += 1 |
|
|
|
|
|
results = model(frame, conf=confidence) |
|
|
frame_alerts = [] |
|
|
detected_objects = [] |
|
|
for result in results: |
|
|
boxes = result.boxes |
|
|
for box in boxes: |
|
|
x1, y1, x2, y2 = map(int, box.xyxy[0]) |
|
|
cls = int(box.cls[0]) |
|
|
conf = float(box.conf[0]) |
|
|
label = f"{model.names[cls]} {conf:.2f}" |
|
|
confidences.append(conf) |
|
|
detected_objects.append((model.names[cls], conf)) |
|
|
|
|
|
if model.names[cls] == "person": |
|
|
color = (0, 255, 0) |
|
|
elif model.names[cls] in ["car", "truck", "bus"]: |
|
|
color = (255, 0, 0) |
|
|
else: |
|
|
color = (0, 0, 255) |
|
|
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2) |
|
|
cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) |
|
|
|
|
|
if model.names[cls] in ["car", "truck", "bus"]: |
|
|
vehicle_count += 1 |
|
|
elif model.names[cls] == "person": |
|
|
pedestrian_count += 1 |
|
|
elif "sign" in model.names[cls].lower(): |
|
|
sign_count += 1 |
|
|
|
|
|
if model.names[cls] == "person": |
|
|
for other_box in boxes: |
|
|
other_cls = int(other_box.cls[0]) |
|
|
if model.names[other_cls] in ["car", "truck", "bus"]: |
|
|
px1, py1, px2, py2 = map(int, box.xyxy[0]) |
|
|
vx1, vy1, vx2, vy2 = map(int, other_box.xyxy[0]) |
|
|
distance = np.sqrt((px1 - vx1)**2 + (py1 - vy1)**2) |
|
|
if distance < 100: |
|
|
alert = f"Frame {frame_count}: Pedestrian near vehicle (distance: {distance:.1f}px, confidence: {conf:.2f})" |
|
|
frame_alerts.append(alert) |
|
|
screenshot_path = f"screenshots/frame_{frame_count}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg" |
|
|
cv2.imwrite(screenshot_path, frame) |
|
|
screenshots.append(screenshot_path) |
|
|
|
|
|
log.write(f"Frame {frame_count}: Detected {len(boxes)} objects\n") |
|
|
for obj, conf in detected_objects: |
|
|
log.write(f" - {obj} (confidence: {conf:.2f})\n") |
|
|
if frame_alerts: |
|
|
log.write(f" Alerts: {', '.join(frame_alerts)}\n") |
|
|
alerts.extend(frame_alerts) |
|
|
|
|
|
out.write(frame) |
|
|
|
|
|
cap.release() |
|
|
out.release() |
|
|
|
|
|
avg_confidence = np.mean(confidences) * 100 if confidences else 0 |
|
|
circle_path = create_percentage_circle(avg_confidence, "Average Detection Confidence") |
|
|
|
|
|
summary = f""" |
|
|
Analysis Complete: |
|
|
- Total Frames Processed: {frame_count} |
|
|
- Average Vehicles per Frame: {vehicle_count / frame_count:.2f} |
|
|
- Average Pedestrians per Frame: {pedestrian_count / frame_count:.2f} |
|
|
- Traffic Signs Detected: {sign_count} |
|
|
- Safety Alerts: {len(alerts)} |
|
|
{chr(10).join(alerts) if alerts else "No critical incidents detected."} |
|
|
- Log File: {log_file} |
|
|
- Screenshots Saved: {len(screenshots)} |
|
|
""" |
|
|
|
|
|
log.write(f"\nSummary:\n{summary}") |
|
|
|
|
|
|
|
|
pdf_path = f"reports/report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf" |
|
|
generate_pdf_report(summary, screenshots, pdf_path) |
|
|
|
|
|
return output_path, summary, screenshots, circle_path, pdf_path |
|
|
|
|
|
css = """ |
|
|
.gradio-container {background-color: #f0f4f8; font-family: Arial;} |
|
|
.footer {display: none !important;} |
|
|
button {background-color: #27ae60; color: white; border-radius: 5px; padding: 10px 20px;} |
|
|
button:hover {background-color: #219653;} |
|
|
.gradio-gallery, .gradio-video, .gradio-textbox {border-radius: 5px; box-shadow: 0 2px 4px rgba(0,0,0,0.1);} |
|
|
""" |
|
|
|
|
|
with gr.Blocks(css=css) as iface: |
|
|
gr.Markdown("# Road Safety AI Video Analysis") |
|
|
gr.Markdown("Upload a traffic video for analysis and download a detailed PDF report.") |
|
|
|
|
|
with gr.Row(): |
|
|
video_input = gr.Video(label="Upload Video") |
|
|
confidence_input = gr.Slider(0.1, 1.0, value=0.3, step=0.1, label="Confidence Threshold") |
|
|
|
|
|
analyze_button = gr.Button("Analyze Video") |
|
|
|
|
|
with gr.Row(): |
|
|
with gr.Column(): |
|
|
video_output = gr.Video(label="Annotated Video") |
|
|
summary_output = gr.Textbox(label="Analysis Summary") |
|
|
with gr.Column(): |
|
|
screenshot_output = gr.Gallery(label="Incident Screenshots") |
|
|
circle_output = gr.Image(label="Confidence Metric") |
|
|
pdf_output = gr.File(label="Download Report") |
|
|
|
|
|
analyze_button.click( |
|
|
fn=analyze_video, |
|
|
inputs=[video_input, confidence_input], |
|
|
outputs=[video_output, summary_output, screenshot_output, circle_output, pdf_output] |
|
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
iface.launch() |
