ROBOFLOW – Trainieren und Testen mit Python

Roboflow ist eine Plattform zum Kommentieren von Bildern zur Verwendung in der Objekterkennungs-KI.

Ich nutze diese Plattform für C2SMR c2smr.fr, meinen Computer-Vision-Verband für Seenotrettung.

In diesem Artikel zeige ich Ihnen, wie Sie diese Plattform nutzen und Ihr Modell mit Python trainieren.

Weiteren Beispielcode finden Sie auf meinem Github: https://github.com/C2SMR/detector

---

I – Datensatz

Um Ihren Datensatz zu erstellen, gehen Sie zu https://app.roboflow.com/ und beginnen Sie mit der Kommentierung Ihres Bildes, wie im folgenden Bild gezeigt.

In diesem Beispiel umgehe ich alle Schwimmer, um ihre Position in zukünftigen Bildern vorherzusagen.
Um ein gutes Ergebnis zu erzielen, schneiden Sie alle Schwimmer zu und platzieren Sie den Begrenzungsrahmen direkt hinter dem Objekt, um es korrekt zu umgeben.

Sie können bereits einen öffentlichen Roboflow-Datensatz verwenden, siehe https://universe.roboflow.com/

II – Ausbildung

Für die Schulungsphase können Sie roboflow direkt nutzen, beim dritten Mal müssen Sie jedoch bezahlen, weshalb ich Ihnen zeige, wie es mit Ihrem Laptop geht.

Der erste Schritt besteht darin, Ihren Datensatz zu importieren. Dazu können Sie die Roboflow-Bibliothek importieren.

pip install roboflow

Um ein Modell zu erstellen, müssen Sie den YOLO-Algorithmus verwenden, den Sie mit der Ultralytics-Bibliothek importieren können.

pip install ultralytics

In meinem Skript verwende ich den folgenden Befehl:

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

Sie müssen Folgendes erhalten:

• der Zugangsschlüssel
• Arbeitsbereich
• Roboflow-Projektname
• Projektdatensatzversion
• Anzahl der Epochen zum Trainieren des Modells
• Größe des neuronalen Netzwerks

Zunächst lädt das Skript yolov8-obb.pt herunter, das Standard-Yolo-Gewicht mit Pre-Workout-Daten, um das Training zu erleichtern.

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 – Anzeige

Nach dem Training des Modells erhalten Sie die Dateien best.py und last.py, die dem Gewicht entsprechen.

Mit der Ultralytics-Bibliothek können Sie auch YOLO importieren und Ihr Gewicht und dann Ihr Testvideo laden.
In diesem Beispiel verwende ich die Tracking-Funktion, um für jeden Schwimmer eine ID zu erhalten.

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

Um die Vorhersage zu analysieren, können Sie den Modell-JSON wie folgt abrufen.

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

Das obige ist der detaillierte Inhalt vonROBOFLOW – Trainieren und Testen mit Python. Für weitere Informationen folgen Sie bitte anderen verwandten Artikeln auf der PHP chinesischen Website!