RF-DETR: un modelo de código abierto para la detección visual de objetos en tiempo real

Introducción general

RF-DETR es un modelo de detección de objetos de código abierto desarrollado por el equipo Roboflow. Se basa en Transformador la característica principal es la eficiencia en tiempo real. El modelo logra la primera detección en tiempo real de más de 60 puntos de acceso en el conjunto de datos COCO de Microsoft, y también destaca en la prueba de referencia RF100-VL, adaptándose a una amplia gama de escenarios del mundo real. Está disponible en dos versiones: RF-DETR-base (29 millones de parámetros) y RF-DETR-large (128 millones de parámetros). El modelo es pequeño y adecuado para el despliegue de dispositivos periféricos. El código y los pesos preentrenados tienen licencia Apache 2.0 y son de uso libre y gratuito para la comunidad. Los usuarios pueden acceder a los recursos desde GitHub para facilitar el entrenamiento o el despliegue.

Lista de funciones

• Detección de objetos en tiempo real: reconocimiento rápido de objetos en imágenes o vídeos con baja latencia.
• Entrenamiento con conjuntos de datos personalizados: soporte para el ajuste de modelos con sus propios datos.
• Funcionamiento en dispositivos periféricos: el modelo es ligero y adecuado para dispositivos con recursos limitados.
• Resolución ajustable: los usuarios pueden equilibrar la velocidad de inspección y la precisión.
• Apoyo al modelo preentrenado: proporciona pesos preentrenados basados en el conjunto de datos COCO.
• Procesamiento de secuencias de vídeo: puede analizar vídeo en tiempo real y emitir resultados.
• Exportación a ONNX: admite la conversión a formato ONNX para facilitar la implantación multiplataforma.
• Entrenamiento multi-GPU: se pueden utilizar varias tarjetas gráficas para acelerar el proceso de entrenamiento.

Utilizar la ayuda

El uso de RF-DETR se divide en tres partes: instalación, inferencia y formación. A continuación se detallan los pasos para ayudarle a empezar rápidamente.

Proceso de instalación

1. Preparación medioambiental
   Requiere Python 3.9 o superior, y PyTorch 1.13.0 o superior. Si utiliza una GPU, ejecute nvidia-smi Comprueba la unidad.
   • Instala PyTorch:

     pip install torch>=1.13.0 torchvision>=0.14.0
     

   • Descargar código:

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

   • Instale la dependencia:

     pip install rfdetr
     

     Esto instalará automáticamente numpyysupervision y otras bibliotecas necesarias.

2. Verificar la instalación
   Ejecute el siguiente código:

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

Si no aparece ningún error, la instalación ha finalizado.

operación de inferencia

RF-DETR viene con un modelo preentrenado del conjunto de datos COCO para detectar imágenes o vídeos directamente.

1. detección de imágenes
   • Código de ejemplo:

     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)
     

   • Este código detecta objetos en la imagen, etiqueta el cuadro delimitador y el nivel de confianza y, a continuación, muestra los resultados.
2. Detección de vídeo
   • En primer lugar, instale opencv-python::

     pip install opencv-python
     

   • Código de ejemplo:

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

   • Esto detectará objetos en el vídeo fotograma a fotograma y los mostrará en tiempo real.
3. Ajuste de la resolución
   • La resolución puede fijarse en la inicialización (debe ser múltiplo de 56):

     model = RFDETRBase(resolution=560)
     

   • Cuanto mayor sea la resolución, mayor será la precisión, pero será más lenta.

Formación de modelos personalizados

RF-DETR admite el ajuste fino con su propio conjunto de datos, pero éste debe estar en formato COCO y contener trainyvalid responder cantando test Tres subdirectorios.

1. Preparación del conjunto de datos
   • Ejemplo de estructura de catálogo:

     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
     

   • Los conjuntos de datos en formato COCO pueden generarse utilizando la plataforma 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. Iniciar la formación
   • Código de ejemplo:

     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)
     

   • Para la formación, el tamaño total de lote recomendado (batch_size * grad_accum_stepsPor ejemplo, las GPUs A100 utilizan la tecnología batch_size=16, grad_accum_steps=1GPU T4 batch_size=4, grad_accum_steps=4.
3. Formación Multi-GPU
   • establecer main.py Documentación:

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

   • Se ejecuta en el terminal:

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

   • comandante en jefe (militar) 8 Sustitúyalo por el número de GPUs que esté utilizando. Tenga en cuenta el ajuste batch_size para mantener estable el tamaño total del lote.
4. Resultados del entrenamiento de carga
   • Tras el entrenamiento se generan dos archivos de pesos: pesos normales y pesos EMA (más estables). Método de carga:

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

Exportación ONNX

• Exporte a formato ONNX para facilitar la implantación en otras plataformas:

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

• El archivo exportado se guarda en la carpeta output Catálogo de inferencia óptima para dispositivos de borde.

escenario de aplicación

1. conducción automática
   RF-DETR detecta vehículos y peatones en la carretera en tiempo real. Su baja latencia y alta precisión son idóneas para sistemas empotrados.
2. control de calidad industrial
   RF-DETR identifica rápidamente los defectos de las piezas en las líneas de montaje de las fábricas. El modelo es ligero y puede funcionar directamente en el equipo.
3. videovigilancia
   RF-DETR procesa vídeo de vigilancia para detectar objetos o comportamientos anómalos en tiempo real. Admite streaming de vídeo y es adecuado para seguridad 24/7.

CONTROL DE CALIDAD

1. ¿Qué formatos de conjuntos de datos se admiten?
   Sólo se admite el formato COCO. El conjunto de datos debe contener trainyvalid responder cantando test subdirectorios, cada uno con su correspondiente _annotations.coco.json Documentación.
2. ¿Cómo obtener la clave API de Roboflow?
   Inicie sesión en https://app.roboflow.com, busque la clave API en la configuración de su cuenta, cópiela y establézcala en la variable de entorno ROBOFLOW_API_KEY.
3. ¿Cuánto dura la formación?
   Depende del hardware y del tamaño del conjunto de datos. En una GPU T4, 10 épocas pueden tardar unas horas. Los conjuntos de datos más pequeños pueden ejecutarse en una CPU, pero es lento.