Roboflow is a platform for annotating images for use in object detection AI.
I use this platform for C2SMR c2smr.fr, my computer vision association for maritime rescue.
In this article I show you how to use this platform and train your model with python.
You can find more sample code on my github : https://github.com/C2SMR/detector
I - Dataset
To create your dataset, go to https://app.roboflow.com/ and start annotating your image as shown in the following image.
In this example, I detour all the swimmers to predict their position in future images.
To get a good result, crop all the swimmers and place the bounding box just after the object to surround it correctly.
You can already use a public roboflow dataset, for this check https://universe.roboflow.com/
II - Training
For the training stage, you can use roboflow directly, but by the third time you'll have to pay, which is why I'm showing you how to do it with your laptop.
The first step is to import your dataset. To do this, you can import the Roboflow library.
pip install roboflow
To create a model, you need to use the YOLO algorithm, which you can import with the ultralytics library.
pip install ultralytics
In my script, I use the following command :
py train.py api-key project-workspace project-name project-version nb-epoch size_model
You must obtain :
- the access key
- workspace
- roboflow project name
- project dataset version
- number of epochs to train the model
- neural network size
Initially, the script downloads yolov8-obb.pt, the default yolo weight with pre-workout data, to facilitate training.
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 - Display
After training the model, you get the files best.py and last.py, which correspond to the weight.
With ultralytics library, you can also import YOLO and load your weight and then your test video.
In this example, I'm using the tracking function to get an ID for each swimmer.
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()
To analyze the prediction, you can obtain the model json as follows.
results = model.track(frame, persist=True)
results_json = json.loads(results[0].tojson())
Top comments (0)