YOLOv8检测图片和视频

一、检测图片

Python

import cv2
from ultralytics import YOLO
import torch

model_path = 'object_detection/best.pt'  # Change this to your YOLOv8 model's path
image_path = 'object_detection/32.jpg'  # Change this to your video's path

# Load the trained YOLOv8 model
model = YOLO(model_path)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print("Using device: %s" % device)
model.to(device)


# Process video frames
image = cv2.imread(image_path)
width, height, _ = image.shape
new_shape = [32*int(height/128), 32*int(width/128)]
image = cv2.resize(image, new_shape)

with torch.no_grad():
    results = model.predict(image)
    for result in results:
        for box in result.boxes:
            x1, y1, x2, y2 = box.xyxy[0]
            # Draw the bounding box on the BGR frame
            cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
            # Add a label above the box
            cv2.putText(image, result.names[int(box.cls)], (int(x1) - 30, int(y1) + 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)

    cv2.imshow('Video', image)

    cv2.waitKey(0)


cv2.destroyAllWindows()

二、检测视频

Python

import cv2
from ultralytics import YOLO
import torch

model_path = 'object_detection/best.pt'  # Change this to your YOLOv8 model's path
video_path = 'object_detection/物块.mp4'  # Change this to your video's path

# Load the trained YOLOv8 model
model = YOLO(model_path)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print("Using device: %s" % device)
model.to(device)
batch_size = 8
frames_rgb = []
frames = []
cap = cv2.VideoCapture(video_path)

if not cap.isOpened():
    print("Error: Could not open video.")
    exit()

# Process video frames
while True:
    ret, frame = cap.read()
    if not ret:
        print("Finished processing video.")
        break
    width, height, _ = frame.shape
    new_shape = [32*int(height/64), 32*int(width/64)]
    frame = cv2.resize(frame, new_shape)
    frames.append(frame)
    # YOLOv8 expects RGB images
    if len(frames) == batch_size:
        with torch.no_grad():
            results = model.predict(frames)

        # Process each detection
        for i, result in enumerate(results):
            for box in result.boxes:
                print(box.conf)
                if float(box.conf) > 0.9:
                    x1, y1, x2, y2 = box.xyxy[0]
                    # Draw the bounding box on the BGR frame
                    cv2.rectangle(frames[i], (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
                    # Add a label above the box
                    cv2.putText(frames[i], result.names[int(box.cls)], (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

            cv2.imshow('Video', frames[i])

            if cv2.waitKey(1) & 0xFF == ord('q'):
                cap.release()
                cv2.destroyAllWindows()
                exit()
        frames.clear()
        frames_rgb.clear()
cap.release()
cv2.destroyAllWindows()

使用了sahi的视频检测

import argparse
import sys
import cv2
from sahi import AutoDetectionModel
from sahi.predict import get_sliced_prediction
import imageio
import numpy as np


def run(weights="yolov8n.pt", source="test.mp4", view_img=False):
    """
    Run object detection on a video using YOLOv8 and SAHI.

    Args:
        weights (str): Model weights path.
        source (str): Video file path.
        view_img (bool): Show results.
    """

    yolov8_model_path = weights
    detection_model = AutoDetectionModel.from_pretrained(
        model_type="yolov8", model_path=yolov8_model_path, confidence_threshold=0.3, device="cuda:0"
    )
    videocapture = cv2.VideoCapture(0)

    new_shape = 32 * int(videocapture.get(3) / 64), 32 * int(videocapture.get(4) / 64)
    writer = imageio.get_writer("object_detection/object_detection.mp4", fps=1 / 0.025)

    while videocapture.isOpened():
        success, frame = videocapture.read()
        if not success:
            break
        frame = cv2.resize(frame, new_shape)
        image = frame.copy()
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        results = get_sliced_prediction(
            frame, detection_model, slice_height=512, slice_width=512, overlap_height_ratio=0.2, overlap_width_ratio=0.2
        )
        object_prediction_list = results.object_prediction_list

        boxes_list = []
        clss_list = []

        for ind, _ in enumerate(object_prediction_list):
            print(object_prediction_list[ind].score.value)
            if float(object_prediction_list[ind].score.value) > 0.85:
                boxes = (
                    object_prediction_list[ind].bbox.minx,
                    object_prediction_list[ind].bbox.miny,
                    object_prediction_list[ind].bbox.maxx,
                    object_prediction_list[ind].bbox.maxy,
                )
                clss = object_prediction_list[ind].category.name
                boxes_list.append(boxes)
                clss_list.append(clss)

        for box, cls in zip(boxes_list, clss_list):
            x1, y1, x2, y2 = box
            cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (56, 56, 255), 2)
            label = str(cls)
            t_size = cv2.getTextSize(label, 0, fontScale=0.6, thickness=1)[0]
            cv2.rectangle(
                image, (int(x1), int(y1) - t_size[1] - 3), (int(x1) + t_size[0], int(y1) + 3), (56, 56, 255), -1
            )
            cv2.putText(
                image, label, (int(x1), int(y1) - 2), 0, 0.6, [255, 255, 255], thickness=1, lineType=cv2.LINE_AA
            )

        if view_img:
            cv2.imshow("result", image)
            frame = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            writer.append_data((np.asarray(frame)).astype(np.uint8))

        if cv2.waitKey(1) == ord("q"):
            videocapture.release()
            cv2.destroyAllWindows()
            sys.exit()
    writer.close()


def parse_opt():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser()
    parser.add_argument("--weights", type=str, default="object_detection/best.pt", help="initial weights path")
    parser.add_argument("--source", type=str, default="object_detection/物块.mp4", help="video file path")
    parser.add_argument("--view-img", type=bool, default=True, help="show results")
    return parser.parse_args()


def main(options):
    """Main function."""
    run(**vars(options))


if __name__ == "__main__":
    opt = parse_opt()
    main(opt)