""" Red Light Violation Detection Test Script """ import cv2 import numpy as np import os import sys import time import argparse # Add parent directory to path for imports parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.append(parent_dir) # Import utilities for crosswalk detection from qt_app_pyside.utils.crosswalk_utils import ( detect_and_draw_crosswalk, # New advanced function with visualization detect_crosswalk, detect_stop_line, draw_violation_line, check_vehicle_violation ) # Import traffic light utilities from qt_app_pyside.utils.traffic_light_utils import detect_traffic_light_color, draw_traffic_light_status def process_test_video(video_path): """ Process a test video to demonstrate red light violation detection. Args: video_path: Path to the test video file """ # Open the video file cap = cv2.VideoCapture(video_path) if not cap.isOpened(): print(f"Error: Could not open video file {video_path}") return # Get video properties width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) fps = cap.get(cv2.CAP_PROP_FPS) print(f"Video loaded: {width}x{height} @ {fps}fps") # Create output directory for results output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_results") os.makedirs(output_dir, exist_ok=True) # Create output video writer output_path = os.path.join(output_dir, "violation_detection_output.avi") fourcc = cv2.VideoWriter_fourcc(*'XVID') out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) # Detection state frame_count = 0 violation_line_y = None traffic_light_color = "unknown" tracked_vehicles = {} violations = [] # Main processing loop while True: ret, frame = cap.read() if not ret: break # Make a copy for annotation annotated_frame = frame.copy() # Every 50 frames, attempt to detect crosswalk/stop line if frame_count % 50 == 0 or violation_line_y is None: # Use advanced function that visualizes the crosswalk annotated_frame, crosswalk_bbox, crosswalk_contours = detect_and_draw_crosswalk(frame) if crosswalk_bbox: violation_line_y = crosswalk_bbox[1] - 10 # 10px before crosswalk print(f"Detected crosswalk at y={violation_line_y}") else: # Try to detect stop line stop_line_y = detect_stop_line(frame) if stop_line_y: violation_line_y = stop_line_y - 10 # 10px before stop line print(f"Detected stop line at y={violation_line_y}") # If still no violation line, use default if violation_line_y is None: violation_line_y = int(height * 0.75) # Default at 75% of height # Draw violation line (make it always thick, visible, and labeled) line_color = (0, 0, 255) if traffic_light_color == "red" else (0, 255, 0) annotated_frame = draw_violation_line( annotated_frame, violation_line_y, color=line_color, thickness=10, style='solid', label=f"Violation Line: y={violation_line_y}") print(f"[DEBUG] Violation line drawn at y={violation_line_y}, color={line_color}, thickness=10") # Demo traffic light detection # In a real app, you would get traffic light bbox from your detector # For this demo, we'll create a fake traffic light region in the corner # Create a demo traffic light bounding box (top-right corner) traffic_light_bbox = [width-100, 50, width-20, 200] # Every 10 frames, simulate traffic light detection # In a real app, you would detect the color from the video if frame_count % 10 == 0: # Alternate between colors for demonstration if traffic_light_color == "red": traffic_light_color = "green" elif traffic_light_color == "green": traffic_light_color = "yellow" elif traffic_light_color == "yellow": traffic_light_color = "red" else: traffic_light_color = "red" # Start with red # Draw a sample traffic light for visualization light_height = traffic_light_bbox[3] - traffic_light_bbox[1] light_width = traffic_light_bbox[2] - traffic_light_bbox[0] # Draw traffic light housing cv2.rectangle(annotated_frame, (traffic_light_bbox[0], traffic_light_bbox[1]), (traffic_light_bbox[2], traffic_light_bbox[3]), (100, 100, 100), -1) # Draw the active light based on current color if traffic_light_color == "red": cv2.circle(annotated_frame, (traffic_light_bbox[0] + light_width//2, traffic_light_bbox[1] + light_height//4), light_width//3, (0, 0, 255), -1) elif traffic_light_color == "yellow": cv2.circle(annotated_frame, (traffic_light_bbox[0] + light_width//2, traffic_light_bbox[1] + light_height//2), light_width//3, (0, 255, 255), -1) elif traffic_light_color == "green": cv2.circle(annotated_frame, (traffic_light_bbox[0] + light_width//2, traffic_light_bbox[1] + 3*light_height//4), light_width//3, (0, 255, 0), -1) # Use our improved function to visualize traffic light status annotated_frame = draw_traffic_light_status(annotated_frame, traffic_light_bbox, traffic_light_color) # Display traffic light color cv2.putText( annotated_frame, f"Traffic Light: {traffic_light_color.upper()}", (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255) if traffic_light_color == "red" else (0, 255, 255) if traffic_light_color == "yellow" else (0, 255, 0), 2 ) # Every 5 frames, simulate vehicle detection if frame_count % 5 == 0: # Simulate vehicle moving from top to bottom vehicle_y = int((frame_count / 500.0) * height) vehicle_x = width // 2 vehicle_width = 100 vehicle_height = 80 # Create bounding box [x1, y1, x2, y2] bbox = [ vehicle_x - vehicle_width // 2, vehicle_y - vehicle_height // 2, vehicle_x + vehicle_width // 2, vehicle_y + vehicle_height // 2 ] # Draw vehicle bbox cv2.rectangle( annotated_frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2 ) # Check for violation if (traffic_light_color == "red" and check_vehicle_violation(bbox, violation_line_y)): # Mark violation cv2.putText( annotated_frame, "RED LIGHT VIOLATION!", (bbox[0], bbox[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2 ) # Re-draw vehicle bbox in red cv2.rectangle( annotated_frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 0, 255), 3 ) # Save violation frame violation_path = os.path.join(output_dir, f"violation_{len(violations)}.jpg") cv2.imwrite(violation_path, frame) violations.append({ "frame": frame_count, "bbox": bbox, "path": violation_path }) print(f"Violation detected at frame {frame_count}") # Write the frame to output video out.write(annotated_frame) # Display frame cv2.imshow('Red Light Violation Detection Test', annotated_frame) # Check for exit key (q) if cv2.waitKey(1) & 0xFF == ord('q'): break frame_count += 1 # Clean up cap.release() out.release() cv2.destroyAllWindows() print(f"Processing complete. {len(violations)} violations detected.") print(f"Output video saved to: {output_path}") if __name__ == "__main__": # Parse command-line arguments parser = argparse.ArgumentParser(description='Test red light violation detection') parser.add_argument('--video', type=str, help='Path to test video file') args = parser.parse_args() # If video path is provided, use it; otherwise download a sample video_path = args.video if not video_path or not os.path.exists(video_path): # Try to find a sample video in the workspace sample_paths = [ "sample_data/traffic.mp4", "../sample_data/traffic.mp4", "test_videos/traffic_light.mp4", "../test_videos/traffic_light.mp4" ] for path in sample_paths: if os.path.exists(path): video_path = path break if not video_path: print("Error: No video file specified. Please provide a path with --video") sys.exit(1) process_test_video(video_path)