RF-DETR: модель с открытым исходным кодом для обнаружения визуальных объектов в реальном времени

Общее введение

RF-DETR - это модель обнаружения объектов с открытым исходным кодом, разработанная командой Roboflow. Она основана на Трансформатор Архитектура, основной особенностью которой является эффективность в реальном времени. Модель впервые обнаружила более 60 точек доступа в режиме реального времени на наборе данных Microsoft COCO, а также показала превосходные результаты в эталонном тесте RF100-VL, адаптируясь к широкому спектру реальных сценариев. Модель доступна в двух версиях: RF-DETR-base (29 миллионов параметров) и RF-DETR-large (128 миллионов параметров). Модель имеет небольшой размер и подходит для развертывания на пограничных устройствах. Код и предварительно обученные веса лицензированы под лицензией Apache 2.0 и являются бесплатными и открытыми для использования сообществом. Пользователи могут получить доступ к ресурсам с GitHub для легкого обучения или развертывания.

Список функций

• Обнаружение объектов в реальном времени: быстрое распознавание объектов на изображениях или видео с низкой задержкой.
• Обучение на пользовательских наборах данных: поддержка настройки моделей на собственных данных.
• Работа на граничных устройствах: модель легкая и подходит для устройств с ограниченными ресурсами.
• Регулируемое разрешение: пользователи могут сбалансировать скорость и точность проверки.
• Поддержка предварительно обученных моделей: предоставляет предварительно обученные веса, основанные на наборе данных COCO.
• Обработка видеопотока: позволяет анализировать видео в режиме реального времени и выдавать результаты.
• Экспорт ONNX: поддерживает преобразование в формат ONNX для упрощения кросс-платформенного развертывания.
• Обучение на нескольких GPU: для ускорения процесса обучения можно использовать несколько видеокарт.

Использование помощи

Использование RF-DETR разделено на три части: установка, вывод и обучение. Ниже приведены подробные шаги, которые помогут вам быстро приступить к работе.

Процесс установки

1. Подготовка к защите окружающей среды
   Требуется Python 3.9 или выше и PyTorch 1.13.0 или выше. Если используется GPU, запустите nvidia-smi Проверьте привод.
   • Установите PyTorch:

     pip install torch>=1.13.0 torchvision>=0.14.0
     

   • Код для скачивания:

     git clone https://github.com/roboflow/rf-detr.git
     cd rf-detr
     

   • Установите зависимость:

     pip install rfdetr
     

     Это автоматически установит numpy, иsupervision и другие необходимые библиотеки.

2. Проверка установки
   Выполните следующий код:

   from rfdetr import RFDETRBase
   print("安装成功")

Если ошибок нет, установка завершена.

операция вывода

RF-DETR поставляется с предварительно обученной моделью набора данных COCO для непосредственного обнаружения изображений или видео.

1. обнаружение изображений
   • Код примера:

     import io
     import requests
     from PIL import Image
     from rfdetr import RFDETRBase
     import supervision as sv
     model = RFDETRBase()
     url = "https://media.roboflow.com/notebooks/examples/dog-2.jpeg"
     image = Image.open(io.BytesIO(requests.get(url).content))
     detections = model.predict(image, threshold=0.5)
     labels = [f"{class_id} {confidence:.2f}" for class_id, confidence in zip(detections.class_id, detections.confidence)]
     annotated_image = image.copy()
     annotated_image = sv.BoxAnnotator().annotate(annotated_image, detections)
     annotated_image = sv.LabelAnnotator().annotate(annotated_image, detections, labels)
     sv.plot_image(annotated_image)
     

   • Этот код обнаруживает объекты на изображении, маркирует ограничительную рамку и уровень доверия, а затем выводит результаты на экран.
2. Обнаружение видео
   • во-первых, установите opencv-python::

     pip install opencv-python
     

   • Код примера:

     import cv2
     from rfdetr import RFDETRBase
     import supervision as sv
     model = RFDETRBase()
     cap = cv2.VideoCapture("video.mp4")  # 替换为你的视频路径
     while cap.isOpened():
     ret, frame = cap.read()
     if not ret:
     break
     image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
     detections = model.predict(image, threshold=0.5)
     annotated_frame = sv.BoxAnnotator().annotate(frame, detections)
     cv2.imshow("RF-DETR Detection", annotated_frame)
     if cv2.waitKey(1) & 0xFF == ord('q'):
     break
     cap.release()
     cv2.destroyAllWindows()
     

   • Это позволит обнаруживать объекты на видео кадр за кадром и отображать их в режиме реального времени.
3. Настройка разрешения
   • Разрешение может быть установлено при инициализации (должно быть кратно 56):

     model = RFDETRBase(resolution=560)
     

   • Чем выше разрешение, тем выше точность, но это будет медленнее.

Обучение пользовательских моделей

RF-DETR поддерживает тонкую настройку с помощью собственного набора данных, но набор данных должен быть в формате COCO, содержащем train, иvalid ответить пением test Три подкаталога.

1. Подготовка набора данных
   • Пример структуры каталога:

     dataset/
     ├── train/
     │   ├── _annotations.coco.json
     │   ├── image1.jpg
     │   └── image2.jpg
     ├── valid/
     │   ├── _annotations.coco.json
     │   ├── image1.jpg
     │   └── image2.jpg
     └── test/
     ├── _annotations.coco.json
     ├── image1.jpg
     └── image2.jpg
     

   • Наборы данных в формате COCO могут быть созданы с помощью платформы Roboflow:

     from roboflow import Roboflow
     rf = Roboflow(api_key="你的API密钥")
     project = rf.workspace("rf-100-vl").project("mahjong-vtacs-mexax-m4vyu-sjtd")
     dataset = project.version(2).download("coco")
     

2. Начните обучение
   • Код примера:

     from rfdetr import RFDETRBase
     model = RFDETRBase()
     model.train(dataset_dir="./mahjong-vtacs-mexax-m4vyu-sjtd-2", epochs=10, batch_size=4, grad_accum_steps=4, lr=1e-4)
     

   • Для обучения рекомендуемый общий размер партии (batch_size * grad_accum_stepsНапример, в графических процессорах A100 используется batch_size=16, grad_accum_steps=1Графические процессоры T4 batch_size=4, grad_accum_steps=4.
3. Обучение работе с несколькими графическими процессорами
   • создать main.py Документация:

     from rfdetr import RFDETRBase
     model = RFDETRBase()
     model.train(dataset_dir="./dataset", epochs=10, batch_size=4, grad_accum_steps=4, lr=1e-4)
     

   • Запускается в терминале:

     python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py
     

   • главнокомандующий (военный) 8 Замените количество используемых графических процессоров. Обратите внимание на настройку batch_size чтобы сохранить стабильный размер партии.
4. Результаты тренировок с нагрузкой
   • После обучения генерируются два файла с весами: обычные веса и веса EMA (более стабильные). Метод загрузки:

     model = RFDETRBase(pretrain_weights="./output/model_ema.pt")
     detections = model.predict("image.jpg")
     

Экспорт ONNX

• Экспорт в формат ONNX для легкого развертывания на других платформах:

  from rfdetr import RFDETRBase
  model = RFDETRBase()
  model.export()
  

• Экспортированный файл сохраняется в папке output Каталог оптимальных выводов для краевых устройств.

сценарий применения

1. автоматическое вождение
   RF-DETR обнаруживает транспортные средства и пешеходов на дороге в режиме реального времени. Низкая задержка и высокая точность подходят для встраиваемых систем.
2. промышленный контроль качества
   RF-DETR быстро выявляет дефекты деталей на заводских сборочных линиях. Модель имеет небольшой вес и может работать непосредственно на оборудовании.
3. видеонаблюдение
   RF-DETR обрабатывает видео с камер наблюдения для обнаружения аномальных объектов или поведения в режиме реального времени. Он поддерживает потоковое видео и подходит для круглосуточной охраны.

QA

1. Какие форматы наборов данных поддерживаются?
   Поддерживается только формат COCO. Набор данных должен содержать train, иvalid ответить пением test подкаталоги, каждый из которых имеет соответствующий _annotations.coco.json Документация.
2. Как получить ключ API Roboflow?
   Войдите на https://app.roboflow.com, найдите ключ API в настройках вашей учетной записи, скопируйте его и установите в переменную окружения ROBOFLOW_API_KEY.
3. Как долго длится обучение?
   Зависит от аппаратного обеспечения и размера набора данных. На GPU T4 10 эпох могут занять несколько часов. Меньшие наборы данных можно обрабатывать на CPU, но это будет медленно.