Spaces:
Sleeping
Sleeping
| import cv2 | |
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image, ImageDraw | |
| from ultralytics import YOLO | |
| import logging | |
| import threading | |
| import queue | |
| import time | |
| # Set up logging | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| # Global variables for line coordinates and line equation | |
| start_point = None | |
| end_point = None | |
| line_params = None # Stores (slope, intercept, start_point, end_point) | |
| # Low-resolution for inference | |
| LOW_RES = (320, 180) | |
| # Frame queue for processed frames | |
| frame_queue = queue.Queue(maxsize=30) # Adjust queue size based on memory constraints | |
| # Thread control flag | |
| processing_active = True | |
| def extract_first_frame(stream_url): | |
| """ | |
| Extracts the first available frame from the IP camera stream and returns it as a PIL image. | |
| """ | |
| logger.info("Attempting to extract the first frame from the stream...") | |
| cap = cv2.VideoCapture(stream_url) | |
| if not cap.isOpened(): | |
| logger.error("Error: Could not open stream.") | |
| return None, "Error: Could not open stream." | |
| ret, frame = cap.read() | |
| cap.release() | |
| if not ret: | |
| logger.error("Error: Could not read the first frame.") | |
| return None, "Error: Could not read the first frame." | |
| # Convert the frame to a PIL image | |
| frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
| pil_image = Image.fromarray(frame_rgb) | |
| logger.info("First frame extracted successfully.") | |
| return pil_image, "First frame extracted successfully." | |
| def update_line(image, evt: gr.SelectData): | |
| """ | |
| Updates the line based on user interaction (click and drag). | |
| """ | |
| global start_point, end_point, line_params | |
| if start_point is None: | |
| start_point = (evt.index[0], evt.index[1]) | |
| draw = ImageDraw.Draw(image) | |
| draw.ellipse((start_point[0] - 5, start_point[1] - 5, start_point[0] + 5, start_point[1] + 5), fill="blue", outline="blue") | |
| return image, f"Line Coordinates:\nStart: {start_point}, End: None" | |
| end_point = (evt.index[0], evt.index[1]) | |
| if start_point[0] != end_point[0]: # Avoid division by zero | |
| slope = (end_point[1] - start_point[1]) / (end_point[0] - start_point[0]) | |
| intercept = start_point[1] - slope * start_point[0] | |
| line_params = (slope, intercept, start_point, end_point) | |
| else: | |
| line_params = (float('inf'), start_point[0], start_point, end_point) | |
| draw = ImageDraw.Draw(image) | |
| draw.line([start_point, end_point], fill="red", width=2) | |
| draw.ellipse((end_point[0] - 5, end_point[1] - 5, end_point[0] + 5, end_point[1] + 5), fill="green", outline="green") | |
| line_info = f"Line Coordinates:\nStart: {start_point}, End: {end_point}\nLine Equation: y = {line_params[0]:.2f}x + {line_params[1]:.2f}" | |
| start_point = None | |
| end_point = None | |
| return image, line_info | |
| def reset_line(): | |
| """ | |
| Resets the line coordinates. | |
| """ | |
| global start_point, end_point, line_params | |
| start_point = None | |
| end_point = None | |
| line_params = None | |
| return None, "Line reset. Click to draw a new line." | |
| def is_object_crossing_line(box, line_params): | |
| """ | |
| Determines if an object's bounding box is fully intersected by the user-drawn line. | |
| """ | |
| _, _, line_start, line_end = line_params | |
| x1, y1, x2, y2 = box | |
| box_edges = [((x1, y1), (x2, y1)), ((x2, y1), (x2, y2)), ((x2, y2), (x1, y2)), ((x1, y2), (x1, y1))] | |
| intersection_count = 0 | |
| for edge_start, edge_end in box_edges: | |
| if intersect(line_start, line_end, edge_start, edge_end): | |
| intersection_count += 1 | |
| return intersection_count >= 2 | |
| def intersect(A, B, C, D): | |
| """ | |
| Determines if two line segments AB and CD intersect. | |
| """ | |
| def ccw(A, B, C): | |
| return (C[1] - A[1]) * (B[0] - A[0]) - (B[1] - A[1]) * (C[0] - A[0]) | |
| def on_segment(A, B, C): | |
| return min(A[0], B[0]) <= C[0] <= max(A[0], B[0]) and min(A[1], B[1]) <= C[1] <= max(A[1], B[1]) | |
| ccw1 = ccw(A, B, C) | |
| ccw2 = ccw(A, B, D) | |
| ccw3 = ccw(C, D, A) | |
| ccw4 = ccw(C, D, B) | |
| return ((ccw1 * ccw2 < 0) and (ccw3 * ccw4 < 0)) or (ccw1 == 0 and on_segment(A, B, C)) or (ccw2 == 0 and on_segment(A, B, D)) or (ccw3 == 0 and on_segment(C, D, A)) or (ccw4 == 0 and on_segment(C, D, B)) | |
| def process_frames(stream_url, confidence_threshold, selected_classes): | |
| """ | |
| Processes frames in a separate thread and adds them to the frame queue. | |
| """ | |
| global processing_active, frame_queue | |
| cap = cv2.VideoCapture(stream_url) | |
| model = YOLO(model="yolo11n.pt") | |
| crossed_objects = {} | |
| while processing_active and cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| # Perform detection on low-res frame | |
| low_res_frame = cv2.resize(frame, LOW_RES) | |
| results = model.track(low_res_frame, persist=True, conf=confidence_threshold) | |
| # Scale bounding boxes to high-res | |
| scale_x = frame.shape[1] / LOW_RES[0] | |
| scale_y = frame.shape[0] / LOW_RES[1] | |
| for detection in results[0].boxes.data: | |
| x1, y1, x2, y2, conf, cls = detection | |
| x1, y1, x2, y2 = int(x1 * scale_x), int(y1 * scale_y), int(x2 * scale_x), int(y2 * scale_y) | |
| if is_object_crossing_line((x1, y1, x2, y2), line_params): | |
| crossed_objects[results[0].boxes.id.int().cpu().tolist()[0]] = True | |
| # Draw bounding boxes and line on the frame | |
| annotated_frame = results[0].plot() | |
| if line_params: | |
| draw_angled_line(annotated_frame, line_params, color=(0, 255, 0), thickness=2) | |
| # Add frame to the queue | |
| if not frame_queue.full(): | |
| frame_queue.put(annotated_frame) | |
| cap.release() | |
| def draw_angled_line(image, line_params, color=(0, 255, 0), thickness=2): | |
| """ | |
| Draws the user-defined line on the frame. | |
| """ | |
| _, _, start_point, end_point = line_params | |
| cv2.line(image, start_point, end_point, color, thickness) | |
| def display_frames(): | |
| """ | |
| Displays frames from the queue at a consistent frame rate. | |
| """ | |
| while processing_active: | |
| if not frame_queue.empty(): | |
| frame = frame_queue.get() | |
| yield cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), "" | |
| else: | |
| time.sleep(0.03) # Wait for the next frame | |
| # Define the Gradio interface | |
| with gr.Blocks() as demo: | |
| gr.Markdown("<h1>Real-time monitoring, object tracking, and line-crossing detection for CCTV camera streams.</h1></center>") | |
| gr.Markdown("## https://github.com/SanshruthR/CCTV_SENTRY_YOLO11") | |
| # Step 1: Enter the IP Camera Stream URL | |
| stream_url = gr.Textbox(label="Enter IP Camera Stream URL", value="https://s104.ipcamlive.com/streams/68idokwtondsqpmkr/stream.m3u8", visible=False) | |
| # Step 1: Extract the first frame from the stream | |
| gr.Markdown("### Step 1: Click on the frame to draw a line, the objects crossing it would be counted in real-time.") | |
| first_frame, status = extract_first_frame(stream_url.value) | |
| if first_frame is None: | |
| gr.Markdown(f"**Error:** {status}") | |
| else: | |
| image = gr.Image(value=first_frame, label="First Frame of Stream", type="pil") | |
| line_info = gr.Textbox(label="Line Coordinates", value="Line Coordinates:\nStart: None, End: None") | |
| image.select(update_line, inputs=image, outputs=[image, line_info]) | |
| # Step 2: Select classes to detect | |
| gr.Markdown("### Step 2: Select Classes to Detect") | |
| model = YOLO(model="yolo11n.pt") | |
| class_names = list(model.names.values()) | |
| selected_classes = gr.CheckboxGroup(choices=class_names, label="Select Classes to Detect") | |
| # Step 3: Adjust confidence threshold | |
| gr.Markdown("### Step 3: Adjust Confidence Threshold (Optional)") | |
| confidence_threshold = gr.Slider(minimum=0.0, maximum=1.0, value=0.2, label="Confidence Threshold") | |
| # Process the stream | |
| process_button = gr.Button("Process Stream") | |
| output_image = gr.Image(label="Processed Frame", streaming=True) | |
| error_box = gr.Textbox(label="Errors/Warnings", interactive=False) | |
| # Event listener for processing the video | |
| process_button.click( | |
| fn=lambda: (setattr(globals(), "processing_active", True), threading.Thread(target=process_frames, args=(stream_url.value, confidence_threshold.value, selected_classes.value)).start()), | |
| outputs=None | |
| ) | |
| # Display frames | |
| demo.load(display_frames, inputs=None, outputs=[output_image, error_box], every=0.03) | |
| # Launch the interface | |
| demo.launch(debug=True) |