yolov8烟火识别代码

以下是使用YOLOv5进行烟火识别的示例代码：

import cv2
import torch
from torch import nn
from torch.nn import functional as F
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords, xyxy2xywh
from utils.torch_utils import select_device


class YOLOv5(nn.Module):
    def __init__(self, weights, device):
        super(YOLOv5, self).__init__()
        self.model = attempt_load(weights, map_location=device)
        self.device = device

    def forward(self, x, conf_thres=0.5, iou_thres=0.5):
        # Inference
        img_size = x.shape[-2:]
        with torch.no_grad():
            pred = self.model(x)[0]
            pred = non_max_suppression(pred, conf_thres, iou_thres)[0]
        if pred is None:
            return None
        # Rescale boxes from img_size to im0 size
        pred[:, :4] = scale_coords(img_size, pred[:, :4], x.shape).round()
        # Convert xyxy to xywh
        pred[:, :4] = xyxy2xywh(pred[:, :4]).round()
        return pred


if __name__ == '__main__':
    device = select_device('0')
    model = YOLOv5('yolov5s.pt', device)

    # Open video file
    cap = cv2.VideoCapture('firework.mp4')

    while True:
        # Read a frame
        ret, frame = cap.read()
        if not ret:
            break
        # Convert BGR to RGB
        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        # Resize and pad
        img0 = letterbox(img, new_shape=640)[0]
        # Convert to Tensor
        img = torch.from_numpy(img0).to(device)
        img = img.permute(2, 0, 1).unsqueeze(0).float() / 255.0
        # Inference
        pred = model(img)
        if pred is not None:
            for x, y, w, h, conf, cls in pred:
                if int(cls) == 0:  # class 0 is firework
                    # Draw box and label
                    label = f'firework: {conf:.2f}'
                    cv2.rectangle(img0, (int(x), int(y)), (int(x + w), int(y + h)), (0, 255, 0), 2)
                    cv2.putText(img0, label, (int(x), int(y) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
        # Display the resulting frame
        cv2.imshow('frame', img0)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # Release video file and close windows
    cap.release()
    cv2.destroyAllWindows()

请注意，此示例代码假设您已经安装了YOLOv5。