Python 시작하기 - YOLO객체 감지

YOLO 객체 감지 결과를 XML 형식으로 저장하는 코드를 제공하겠습니다. 이 코드는 비디오에서 YOLOv8 객체 감지를 수행한 후, 감지된 객체 정보(클래스, 바운딩 박스 좌표 등)를 XML 파일로 저장합니다. 결과는 위에서 제공한 XML 형식과 유사하게 작성됩니다.

YOLO 객체 감지 코드:

import cv2
import torch
from ultralytics import YOLO
import xml.etree.ElementTree as ET

# YOLOv8 모델 로드
model = YOLO('yolo.pt')  # YOLO 모델 파일 로드 

# 비디오 파일 경로 설정
video_path = 'test_video.mp4'

# XML 파일로 저장하는 함수
def save_results_to_xml(detections, frame_id, xml_filename='yolo_output.xml'):
    root = ET.Element("YoloDetectionResults")
    library = ET.SubElement(root, "Library")

    # 기본 메타데이터 설정
    ET.SubElement(library, "Scenario").text = "ObjectDetection"
    ET.SubElement(library, "Dataset").text = "CustomDataset"
    ET.SubElement(library, "Libversion").text = "1.0"

    clip = ET.SubElement(library, "Clip")
    header = ET.SubElement(clip, "Header")
    ET.SubElement(header, "Filename").text = video_path
    ET.SubElement(header, "Stage").text = "1"
    ET.SubElement(header, "Duration").text = "00:00:00"
    ET.SubElement(header, "Location").text = "UrbanArea"

    # 감지 결과를 추가
    alarms = ET.SubElement(clip, "Alarms")

    for detection in detections:
        # YOLO 감지 결과에서 클래스명과 바운딩 박스 좌표 추출
        for box in detection.boxes:
            alarm = ET.SubElement(alarms, "Alarm")
            ET.SubElement(alarm, "StartTime").text = f"{frame_id // 30:02d}:{(frame_id % 30) * 2:02d}" # 대략적인 시작 시간
            ET.SubElement(alarm, "AlarmDescription").text = detection.names[int(box.cls)]  # 객체 클래스 이름

            # 바운딩 박스 좌표
            bbox = ET.SubElement(alarm, "BoundingBox")
            ET.SubElement(bbox, "xmin").text = str(int(box.xyxy[0][0]))
            ET.SubElement(bbox, "ymin").text = str(int(box.xyxy[0][1]))
            ET.SubElement(bbox, "xmax").text = str(int(box.xyxy[0][2]))
            ET.SubElement(bbox, "ymax").text = str(int(box.xyxy[0][3]))

    # XML 파일로 저장
    tree = ET.ElementTree(root)
    tree.write(xml_filename)

# 비디오 파일 읽기
cap = cv2.VideoCapture(video_path)
frame_id = 0

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # YOLOv8로 객체 감지 수행
    results = model(frame)

    # 감지 결과를 XML로 저장 (프레임 단위로 저장)
    save_results_to_xml(results, frame_id, xml_filename='yolo_output.xml')

    frame_id += 1

# 비디오 릴리즈
cap.release()

코드 설명:

1. YOLO 모델 로드: YOLOv8 모델을 로드하여 비디오 프레임에서 객체 감지를 수행합니다.
2. 비디오 프레임 읽기: OpenCV를 사용해 비디오를 프레임 단위로 읽어들이고, 각 프레임마다 YOLO으로 객체 감지를 수행합니다.
3. XML로 저장: 감지된 객체 정보를 save_results_to_xml() 함수로 XML 형식으로 저장합니다. 감지된 객체의 클래스, 바운딩 박스 좌표를 XML에 기록하고, 감지된 시간과 지속 시간도 저장합니다.
4. 프레임 번호: frame_id는 대략적인 감지 시간을 계산하기 위해 사용됩니다(60 FPS 기준).

결과:

이 코드를 실행하면 비디오에서 감지된 YOLOv8 객체 정보를 기반으로 XML 파일(yolo_output.xml)이 생성됩니다. XML 파일은 감지된 객체의 클래스, 바운딩 박스 좌표, 감지된 시간 등의 정보를 포함하게 됩니다.