Roboflow 是一个用于注释图像以用于对象检测 AI 的平台。
我将这个平台用于 C2SMR c2smr.fr，我的海上救援计算机视觉协会。
在本文中，我将向您展示如何使用该平台并使用 python 训练模型。
您可以在我的 github 上找到更多示例代码：https://github.com/C2SMR/Detector
要创建数据集，请访问 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
您必须获得：
最初，脚本会下载 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中文网其他相关文章！
