[YOLOv5] 학습 파일로 동영상 객체탐지 테스트

테스트 순서

비디오에서 관심영역 정하기

 

파일 및 경로

체크한 것: 사용한 소스파일과 코드파일

 

"""video_frame_capture.py"""

import cv2

# 비디오 읽기
cap = cv2.VideoCapture('video.mp4')

# 캡쳐할 프레임 번호
frame_number = 50  # 이 값을 원하는 프레임 번호로 변경하세요

# 프레임 번호로 이동
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)

# 해당 프레임 읽기
ret, frame = cap.read()

# 좌표를 저장할 리스트
coords = []

# 마우스 콜백 함수 정의
def get_coordinates(event, x, y, flags, param):
    if event == cv2.EVENT_LBUTTONDOWN:
        print('Point Selected: ', x, y)
        coords.append((x, y))

# 윈도우 이름 설정 및 마우스 콜백 함수 등록
cv2.namedWindow("image")
cv2.setMouseCallback("image", get_coordinates)

# 이미지 표시
cv2.imshow("image", frame)
cv2.waitKey(0)

# 모든 좌표 출력
print("All selected coordinates: ", coords)

cap.release()
cv2.destroyAllWindows()

 

 

---

비디오에서 객체 탐지(관심영역 설정)

 

위에서 구한 관심영역 좌표로 코드 작성하였음

"""for_video_test_colab_train_Result_import.py"""

import torch
import cv2
import numpy as np

# YOLOv5 모델 로드
model_path = 'C:/Users/songs/PycharmProjects/mediapipe/yolov5/data/dataset_230515/230518_서영학습_테스트용 사진들/best_SY_230518.pt'
model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path)

# 비디오 읽기
cap = cv2.VideoCapture('video2.mp4')

# 저장할 비디오 설정
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(
    'C:/Users/songs/PycharmProjects/mediapipe/yolov5/data/dataset_230515/230518_서영학습_테스트용 사진들/video_detect2.mp4', fourcc,
    30.0, (int(cap.get(3)), int(cap.get(4))))

points = np.array([[353, 35]
,[458, 51]
,[585, 80]
,[712, 116]
,[844, 156]
,[931, 193]
,[1011, 274]
,[1076, 363]
,[1111, 439]
,[1007, 520]
,[880, 565]
,[761, 603]
,[657, 643]
,[544, 693]
,[514, 566]
,[482, 447]
,[453, 346]
,[413, 216]
,[351, 45]], dtype=np.int32)

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

    if not ret:
        break

    # 관심 영역 설정 (창문 안쪽을 까맣게 처리)
    mask = np.ones_like(frame, dtype=np.uint8) * 255
    cv2.fillPoly(mask, [points], (0, 0, 0))
    roi = cv2.bitwise_and(frame, mask)

    # 객체 탐지 실행
    results = model(roi[:, :, ::-1])  # YOLOv5는 RGB 이미지를 기대하므로 BGR에서 RGB로 변환

    # 결과 출력
    results.print()

    # 탐지된 객체에 대한 경계 상자와 클래스를 그림
    rendered_frame = results.render()[0]

    # 비디오에 프레임 쓰기
    out.write(rendered_frame[:, :, ::-1])  # RGB에서 BGR로 변환하여 OpenCV로 저장

    # 결과창 띄우기
    cv2.imshow('Detections', rendered_frame[:, :, ::-1])  # RGB에서 BGR로 변환하여 OpenCV로 표시

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
out.release()
cv2.destroyAllWindows()

 

 

돌려본 결과

video_detect_2.mp4

4.14MB