ROBOFLOW - 使用 python 進行訓練和測試

Roboflow 是用於註解影像以用於物件偵測 AI 的平台。

我將這個平台用於 C2SMR c2smr.fr，我的海上救援電腦視覺協會。

在本文中，我將向您展示如何使用該平台並使用 python 訓練模型。

您可以在我的 github 上找到更多範例程式碼：https://github.com/C2SMR/Detector

---

I - 資料集

要建立資料集，請造訪 https://app.roboflow.com/ 並開始註解您的圖像，如下圖所示。

在這個例子中，我繞道所有游泳者來預測他們在未來影像中的位置。
為了獲得良好的結果，請裁剪所有游泳者並將邊界框放置在物件後面以正確包圍它。

您已經可以使用公共 roboflow 資料集進行此檢查 https://universe.roboflow.com/

二、培訓

在訓練階段，你可以直接使用 roboflow，但到了第三次你就需要付費了，這就是為什麼我向你展示如何使用筆記型電腦進行操作。

第一步是匯入資料集。為此，您可以匯入 Roboflow 庫。

pip install roboflow

要建立模型，您需要使用 YOLO 演算法，您可以使用 ultralytics 函式庫匯入該演算法。

pip install ultralytics

在我的腳本中，我使用以下命令：

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

您必須取得：

• 存取金鑰
• 工作區
• roboflow 項目名稱
• 專案資料集版本
• 訓練模型的紀元數
• 神經網路大小

最初，腳本會下載 yolov8-obb​​.pt，這是帶有鍛鍊前資料的預設 yolo 權重，以方便訓練。

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()

三、顯示

訓練完模型後，得到檔案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())

以上是ROBOFLOW - 使用 python 進行訓練和測試的詳細內容。更多資訊請關注PHP中文網其他相關文章！