Roboflow é uma plataforma para anotação de imagens para uso em IA de detecção de objetos.

Eu uso esta plataforma para C2SMR c2smr.fr, minha associação de visão computacional para resgate marítimo.

Neste artigo mostro como usar esta plataforma e treinar seu modelo com python.

Você pode encontrar mais exemplos de código em meu github: https://github.com/C2SMR/detector

I - Conjunto de dados

Para criar seu conjunto de dados, acesse https://app.roboflow.com/ e comece a anotar sua imagem conforme mostrado na imagem a seguir.

Neste exemplo, desvio todos os nadadores para prever sua posição em imagens futuras.
Para obter um bom resultado, corte todos os nadadores e coloque a caixa delimitadora logo após o objeto para cercá-lo corretamente.

Você já pode usar um conjunto de dados roboflow público, para isso verifique https://universe.roboflow.com/

II - Treinamento

Para a etapa de treinamento, você pode usar o roboflow diretamente, mas na terceira vez você terá que pagar, por isso estou mostrando como fazer isso com seu laptop.

A primeira etapa é importar seu conjunto de dados. Para fazer isso, você pode importar a biblioteca Roboflow.

pip install roboflow

Para criar um modelo, você precisa usar o algoritmo YOLO, que pode ser importado com a biblioteca ultralítica.

pip install ultralytics

No meu script, eu uso o seguinte comando:

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

Você deve obter:

• a chave de acesso
• espaço de trabalho
• nome do projeto roboflow
• versão do conjunto de dados do projeto
• número de épocas para treinar o modelo
• tamanho da rede neural

Inicialmente, o script baixa yolov8-obb.pt, o peso padrão do yolo com dados pré-treino, para facilitar o treinamento.

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 - Exibição

Após treinar o modelo, você obtém os arquivos best.py e last.py, que correspondem ao peso.

Com a biblioteca ultralítica, você também pode importar YOLO e carregar seu peso e depois seu vídeo de teste.
Neste exemplo, estou usando a função de rastreamento para obter um ID para 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 analisar a previsão, você pode obter o modelo json da seguinte maneira.

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