영상인식 기초, pickle 저장하기, yolo11x.pt, player검출하기, __init__.py 역할

오늘은 yolo11x.pt를 가지고 테니스 코트의 player을 검출하는 방법을 알아본다. 해당 글에서는 사실 player는 아니고 사람은 전부다 검출한다. pickle을 이용해서 훈련시간을 줄인다. 저장하고 저장한 것을 활용하는 방법을 소개한다. 

 

play_trackers.py

from ultralytics import YOLO
import cv2
import pickle

class PlayerTracker:
    def __init__(self, model_path):
        self.model = YOLO(model_path)

    def detect_frame(self, frame):
        results = self.model.track(frame, persist=True)[0]
        class_names = results.names
        player_dict = {}
        for box in results.boxes:
            track_id = int(box.id.tolist()[0])
            result = box.xyxy.tolist()[0]
            class_ids = box.cls.tolist()[0]
            det_class_names = class_names[class_ids]
            if det_class_names =="person":
                player_dict[track_id] = result
        return player_dict

    def detect_frames(self, frames, read_from_stub=False, stub_path=None):
        player_detections = []
        if read_from_stub and stub_path is not None:
            with open(stub_path, 'rb') as f:
                player_detections = pickle.load(f)
            return player_detections

        for frame in frames:
            player_dict = self.detect_frame(frame)
            player_detections.append(player_dict)

        if stub_path is not None:
            with open (stub_path, 'wb') as f:
                pickle.dump(player_detections, f)

        return player_detections

    def draw_bboxes(self, video_frames, player_detections):
        output_video_frames=[]
        for frame, player_dict in zip(video_frames, player_detections):
            for track_id, bbox in player_dict.items():
                x1,y1,x2,y2 = bbox
                cv2.putText(frame, f"Player ID: {track_id}", (int(x1), int(y1)-10), cv2.FONT_HERSHEY_COMPLEX_SMALL, 0.9, (0,0,255),2)
                cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0,0,255),2)
            output_video_frames.append(frame)
        return output_video_frames

 

 

 

1. __init__

    def __init__(self, model_path):
        self.model = YOLO(model_path)

 

 

2. frame에서 사물 검출하기

    def detect_frame(self, frame):
        results = self.model.track(frame, persist=True)[0]
        class_names = results.names
        player_dict = {}
        for box in results.boxes:
            track_id = int(box.id.tolist()[0])
            result = box.xyxy.tolist()[0]
            class_ids = box.cls.tolist()[0]
            det_class_names = class_names[class_ids]
            if det_class_names =="person":
                player_dict[track_id] = result
        return player_dict

 

 

3. frame에서 사람의 좌표 모으기 + pickle

    def detect_frames(self, frames, read_from_stub=False, stub_path=None):
        player_detections = []
        if read_from_stub and stub_path is not None:
            with open(stub_path, 'rb') as f:
                player_detections = pickle.load(f)
            return player_detections

        for frame in frames:
            player_dict = self.detect_frame(frame)
            player_detections.append(player_dict)

        if stub_path is not None:
            with open (stub_path, 'wb') as f:
                pickle.dump(player_detections, f)

        return player_detections

pickle을 이용해서 오랜 시간 좌표를 수집하는데 걸리는 시간을 단축한다. 

따로 저장해 놓는다. 

다음에 불러오고 싶을 때는 read_from_stub= True로 만들기만 하면 끝!

 

4. box 그리기

    def draw_bboxes(self, video_frames, player_detections):
        output_video_frames=[]
        for frame, player_dict in zip(video_frames, player_detections):
            for track_id, bbox in player_dict.items():
                x1,y1,x2,y2 = bbox
                cv2.putText(frame, f"Player ID: {track_id}", (int(x1), int(y1)-10), cv2.FONT_HERSHEY_COMPLEX_SMALL, 0.9, (0,0,255),2)
                cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0,0,255),2)
            output_video_frames.append(frame)
        return output_video_frames

보통 bounding box는 [x1,y1,x2,y2]의 형태로 나타냄.

x1, y1은 왼쪽 상단, x2,y2는 오른쪽 하단 좌표

 

cv2.rectangle에서 좌표를 가지고 직사각형을 그린다. 

 

 

mian.py

import cv2
from ultralytics import YOLO
from utils import (read_video, save_video)
from trackers import PlayerTracker

def main():
    print("hello")
    input_video_path = "input_videos/input_video.mp4"
    video_frames = read_video(input_video_path)

    #Detect players
    player_tracker =PlayerTracker(model_path="yolo11x.pt")
    player_detections= player_tracker.detect_frames(video_frames,  read_from_stub=False, stub_path = "trackers_stub/player_detection.pkl")

    #Draw Output

    #Draw Player Bounding Boxes
    output_video_frames = player_tracker.draw_bboxes(video_frames, player_detections)

    #Save the Output Video
    save_video(output_video_frames, "output_videos/output.avi")

if __name__ == "__main__":
    main()

 

 

__init__.py

from .player_tracker import PlayerTracker

파이썬 패키지에서 __init__.py는 패키지를 정의하고 초기화하는데 사용한다. 이 파일을 패키지 디렉터리에 넣음으로써 해당 디렉터리를 파이썬 모듈로 인식하게 된다. 이 파일이 없다면 해당 디렉터리는 단순한 폴더일 뿐 파이썬 패키지로 취금되지 않으며, 해당 디렉터리 내의 파일들을 임포트할 수 없다.