yolo

[19기_한지희]

hitter_trained_model.pt

5.95MB

import cv2
import torch
from ultralytics import YOLO

# YOLO 모델을 로드합니다 (이미 로드된 상태라면 이 부분은 건너뛰세요)
model = YOLO('/content/drive/MyDrive/batting/hitter_trained_model.pt')

# 클래스 이름 설정 (모델 학습 시 사용한 클래스 이름 목록을 여기에 입력하세요)
class_names = ["geonchang", "other_class"]  # 예시로 geonchang과 다른 클래스를 넣었습니다. 적절히 수정하세요.

# 동영상 파일 경로
video_path = '/content/drive/MyDrive/batting/test.mp4'  # 동영상 경로를 적절히 수정하세요
output_video_path = '/content/drive/MyDrive/batting/test_output.mp4'  # 출력될 동영상 경로

# 동영상을 불러옵니다
cap = cv2.VideoCapture(video_path)

# 동영상의 프레임 크기 및 FPS를 가져옵니다
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)

# 출력 동영상을 저장할 준비를 합니다
fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 코덱 설정
out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))

# Confidence와 IOU 임계값을 설정
confidence_threshold = 0.6  # 원하는 신뢰도 임계값으로 조정
iou_threshold = 0.5         # 원하는 IOU 임계값으로 조정

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break  # 동영상이 끝나면 루프를 종료합니다

    # YOLO 모델을 사용하여 현재 프레임에서 사물 검출을 수행합니다
    results = model(frame)

    # confidence_threshold 이상인 검출 결과만 그리도록 필터링
    for result in results:
        boxes = result.boxes  # 검출된 박스 가져오기
        for box in boxes:
            if box.conf >= confidence_threshold:  # 신뢰도 필터링
                # 박스 좌표 및 신뢰도 가져오기
                x1, y1, x2, y2 = map(int, box.xyxy[0])  # 좌표
                confidence = box.conf.item()            # 신뢰도
                class_id = int(box.cls.item())          # 클래스 인덱스
                label = f"{class_names[class_id]} {confidence:.2f}"  # 클래스 이름과 신뢰도 표시

                # 박스를 그립니다
                cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
                cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

    # 결과가 그려진 프레임을 출력 동영상에 저장합니다
    out.write(frame)

    # 결과를 실시간으로 확인하려면 아래 코드 사용 (원하는 경우)
    # cv2.imshow('Detected Frame', frame)
    # if cv2.waitKey(1) & 0xFF == ord('q'):
    #     break

# 리소스 해제
cap.release()
out.release()
# cv2.destroyAllWindows()

test.mp4

1.70MB

yk.mp4

670.12KB

[19기_한지희]

https://blog.testworks.co.kr/data-labeling/

[유연수]

https://github.com/wkentaro/labelme

[유연수]

https://labelbox.com/