Roboflow는 객체 감지 AI에 사용하기 위해 이미지에 주석을 달기 위한 플랫폼입니다.

저는 해상 구조를 위한 컴퓨터 비전 협회인 C2SMR c2smr.fr에 이 플랫폼을 사용합니다.

이 기사에서는 이 플랫폼을 사용하고 Python으로 모델을 교육하는 방법을 보여줍니다.

내 github에서 더 많은 샘플 코드를 찾을 수 있습니다: https://github.com/C2SMR/Detector

I - 데이터세트

데이터 세트를 생성하려면 https://app.roboflow.com/으로 이동하여 다음 이미지와 같이 이미지에 주석을 달기 시작하세요.

이 예에서는 모든 수영선수를 우회하여 미래 이미지에서 그들의 위치를 ​​예측합니다.
좋은 결과를 얻으려면 모든 수영 선수를 자르고 개체 바로 뒤에 경계 상자를 배치하여 개체를 올바르게 둘러쌉니다.

이 확인을 위해 이미 공개 roboflow 데이터세트를 사용할 수 있습니다 https://universe.roboflow.com/

II - 훈련

훈련 단계에서는 roboflow를 직접 사용할 수 있지만 세 번째부터는 비용을 지불해야 하므로 노트북으로 하는 방법을 보여드리겠습니다.

첫 번째 단계는 데이터세트를 가져오는 것입니다. 이를 위해 Roboflow 라이브러리를 가져올 수 있습니다.

pip install roboflow

모델을 생성하려면 Ultralytics 라이브러리로 가져올 수 있는 YOLO 알고리즘을 사용해야 합니다.

pip install ultralytics

내 스크립트에서는 다음 명령을 사용합니다.

py train.py api-key project-workspace project-name project-version nb-epoch size_model

다음을 획득해야 합니다:

• 액세스 키
• 작업공간
• roboflow 프로젝트 이름
• 프로젝트 데이터세트 버전
• 모델 학습을 위한 시대 수
• 신경망 크기

처음에 스크립트는 훈련을 용이하게 하기 위해 운동 전 데이터가 포함된 기본 욜로 웨이트인 yolov8-obb.pt를 다운로드합니다.

import sys
import os
import random
from roboflow import Roboflow
from ultralytics import YOLO
import yaml
import time


class Main:
    rf: Roboflow
    project: object
    dataset: object
    model: object
    results: object
    model_size: str

    def __init__(self):
        self.model_size = sys.argv[6]
        self.import_dataset()
        self.train()

    def import_dataset(self):
        self.rf = Roboflow(api_key=sys.argv[1])
        self.project = self.rf.workspace(sys.argv[2]).project(sys.argv[3])
        self.dataset = self.project.version(sys.argv[4]).download("yolov8-obb")

        with open(f'{self.dataset.location}/data.yaml', 'r') as file:
            data = yaml.safe_load(file)

        data['path'] = self.dataset.location

        with open(f'{self.dataset.location}/data.yaml', 'w') as file:
            yaml.dump(data, file, sort_keys=False)

    def train(self):
        list_of_models = ["n", "s", "m", "l", "x"]
        if self.model_size != "ALL" and self.model_size in list_of_models:

            self.model = YOLO(f"yolov8{self.model_size}-obb.pt")

            self.results = self.model.train(data=f"{self.dataset.location}/"
                                                 f"yolov8-obb.yaml",
                                            epochs=int(sys.argv[5]), imgsz=640)



        elif self.model_size == "ALL":
            for model_size in list_of_models:
                self.model = YOLO(f"yolov8{model_size}.pt")

                self.results = self.model.train(data=f"{self.dataset.location}"
                                                     f"/yolov8-obb.yaml",
                                                epochs=int(sys.argv[5]),
                                                imgsz=640)



        else:
            print("Invalid model size")



if __name__ == '__main__':
    Main()

III - 디스플레이

모델을 훈련한 후 가중치에 해당하는 best.py 및 last.py 파일을 얻습니다.

ultralytics 라이브러리를 사용하면 YOLO를 가져와 체중을 로드한 다음 테스트 비디오를 로드할 수도 있습니다.
이 예에서는 추적 기능을 사용하여 각 수영 선수의 ID를 가져옵니다.

import cv2
from ultralytics import YOLO
import sys


def main():
    cap = cv2.VideoCapture(sys.argv[1])

    model = YOLO(sys.argv[2])

    while True:
        ret, frame = cap.read()
        results = model.track(frame, persist=True)
        res_plotted = results[0].plot()
        cv2.imshow("frame", res_plotted)

        if cv2.waitKey(1) == 27:
            break

    cap.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()

예측을 분석하기 위해 다음과 같이 모델 json을 얻을 수 있습니다.

 results = model.track(frame, persist=True)
 results_json = json.loads(results[0].tojson())