Search code examples
pythonopencvobject-detectionyoloyolov8

Detection of object above line using Yolov8 and OpenCV


So basically I am using YOLOv8 for object detection. I have passed my RTSP URL of CCTV as my video path. So it takes the feed from the CCTV and detects objects in real time. Now what I want to do is create an imaginary line using OpenCV and detect objects only below that line. So the bounding boxes should come below the line only. All the objects that are above the line shouldn't be detected and filtered out

Below is the code

import cv2
import torch
import numpy as np
from ultralytics import YOLO

video_path ="rtsp://192.168.1.83/live/0/MAIN"


cap = cv2.VideoCapture(video_path)

model = YOLO('yolov8n.pt')

x_line=600

while cap.isOpened():
    # Read a frame from the video
    success, frame = cap.read()
    width = int(cap.get(3))
    height = int(cap.get(4))

    if success:
        
        # Run YOLOv8 inference on the frame
        resized_frame = cv2.resize(frame, (1280, 720), interpolation=cv2.INTER_LINEAR)
        cv2.line(resized_frame, (0, x_line), (width, x_line), (255, 0, 0), 10)

        # Visualize the results on the frame
        results = model(resized_frame, conf=0.6,classes=0)
        
        annotated_frame = results[0].plot()

        # Display the annotated frame
        cv2.imshow("YOLOv8 Inference", annotated_frame)

        # Break the loop if 'q' is pressed
        if cv2.waitKey(1) & 0xFF == ord("q"):
            break

    else:
        # Break the loop if the end of the video is reached
        break

cap.release()
cv2.destroyAllWindows()

Solution

  • Based on the discussion above you can simply filter the result set according to your region of interest:

    import cv2
    from ultralytics import YOLO
    from ultralytics.yolo.utils.plotting import Annotator
    
    
    model = YOLO('yolov8n.pt')
    x_line = 100
    
    img = cv2.imread('zidane.jpg')
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    results = model.predict(img, conf=0.5, classes=0)
    annotator = Annotator(img)
    
    for r in results:
        for box in r.boxes:
            b = box.xyxy[0]
            if b[1] > x_line:
                c = box.cls
                annotator.box_label(b, f"{r.names[int(c)]} {float(box.conf):.2}")
    
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    cv2.line(img, (0, x_line), (img.shape[1] - 1, x_line), (255, 0, 0), 2)
    
    cv2.imshow("YOLO", img)
    cv2.waitKey(0)
    
    cv2.destroyAllWindows()
    

    Result:

    enter image description here