Roboflow es una plataforma para anotar imágenes para su uso en IA de detección de objetos.

Utilizo esta plataforma para C2SMR c2smr.fr, mi asociación de visión por computadora para rescate marítimo.

En este artículo te muestro cómo usar esta plataforma y entrenar tu modelo con python.

Puedes encontrar más código de muestra en mi github: https://github.com/C2SMR/detector

I - conjunto de datos

Para crear su conjunto de datos, vaya a https://app.roboflow.com/ y comience a anotar su imagen como se muestra en la siguiente imagen.

En este ejemplo, desvío a todos los nadadores para predecir su posición en imágenes futuras.
Para obtener un buen resultado, recorte a todos los nadadores y coloque el cuadro delimitador justo después del objeto para rodearlo correctamente.

Ya puedes usar un conjunto de datos público de roboflow; para esto, consulta https://universe.roboflow.com/

II - Formación

Para la etapa de capacitación, puedes usar roboflow directamente, pero para la tercera vez tendrás que pagar, por eso te muestro cómo hacerlo con tu computadora portátil.

El primer paso es importar su conjunto de datos. Para hacer esto, puedes importar la biblioteca Roboflow.

pip install roboflow

Para crear un modelo, necesitas usar el algoritmo YOLO, que puedes importar con la biblioteca ultralytics.

pip install ultralytics

En mi script, uso el siguiente comando:

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

Debes obtener:

• la clave de acceso
• espacio de trabajo
• nombre del proyecto roboflow
• versión del conjunto de datos del proyecto
• número de épocas para entrenar el modelo
• tamaño de la red neuronal

Inicialmente, el script descarga yolov8-obb.pt, el peso predeterminado de yolo con datos previos al entrenamiento, para facilitar el entrenamiento.

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 - Pantalla

Después de entrenar el modelo, obtienes los archivos best.py y last.py, que corresponden al peso.

Con la biblioteca de ultralíticos, también puedes importar YOLO y cargar tu peso y luego tu video de prueba.
En este ejemplo, estoy usando la función de seguimiento para obtener una identificación de cada nadador.

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

Para analizar la predicción, puede obtener el modelo json de la siguiente manera.

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