RF-DETR: um modelo de código aberto para detecção de objetos visuais em tempo real

Introdução geral

O RF-DETR é um modelo de detecção de objetos de código aberto desenvolvido pela equipe do Roboflow. Ele é baseado em Transformador O recurso principal é a eficiência em tempo real. O modelo consegue a primeira detecção em tempo real de mais de 60 APs no conjunto de dados Microsoft COCO e também se destaca no teste de benchmark RF100-VL, adaptando-se a uma ampla variedade de cenários do mundo real. Ele está disponível em duas versões: RF-DETR-base (29 milhões de parâmetros) e RF-DETR-large (128 milhões de parâmetros). O modelo é pequeno e adequado para a implantação de dispositivos de borda. O código e os pesos pré-treinados são licenciados sob a licença Apache 2.0 e são gratuitos e abertos para uso da comunidade. Os usuários podem acessar os recursos do GitHub para facilitar o treinamento ou a implementação.

Lista de funções

• Detecção de objetos em tempo real: reconhecimento rápido de objetos em imagens ou vídeos com baixa latência.
• Treinamento de conjunto de dados personalizado: suporte para ajustar modelos com seus próprios dados.
• Execução em dispositivos de borda: o modelo é leve e adequado para dispositivos com recursos limitados.
• Resolução ajustável: os usuários podem equilibrar a velocidade e a precisão da inspeção.
• Suporte a modelos pré-treinados: fornece pesos pré-treinados com base no conjunto de dados COCO.
• Processamento de fluxo de vídeo: pode analisar vídeos em tempo real e gerar resultados.
• Exportação ONNX: suporta a conversão para o formato ONNX para facilitar a implementação em várias plataformas.
• Treinamento multi-GPU: várias placas gráficas podem ser usadas para acelerar o processo de treinamento.

Usando a Ajuda

O uso do RF-DETR é dividido em três partes: instalação, inferência e treinamento. Abaixo estão as etapas detalhadas para ajudá-lo a começar rapidamente.

Processo de instalação

1. Preparação ambiental
   Requer Python 3.9 ou superior e PyTorch 1.13.0 ou superior. Se estiver usando uma GPU, execute nvidia-smi Verifique a unidade.
   • Instale o PyTorch:

     pip install torch>=1.13.0 torchvision>=0.14.0
     

   • Código de download:

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

   • Instale a dependência:

     pip install rfdetr
     

     Isso instalará automaticamente numpyesupervision e outras bibliotecas necessárias.

2. Verificar a instalação
   Execute o seguinte código:

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

Se nenhum erro for relatado, a instalação estará concluída.

operação de inferência

O RF-DETR vem com um modelo pré-treinado do conjunto de dados COCO para detectar imagens ou vídeos diretamente.

1. detecção de imagens
   • Código de amostra:

     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)
     

   • Esse código detecta objetos na imagem, rotula a caixa delimitadora e o nível de confiança e, em seguida, exibe os resultados.
2. Detecção de vídeo
   • Em primeiro lugar, instale opencv-python::

     pip install opencv-python
     

   • Código de amostra:

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

   • Isso detectará objetos no vídeo quadro a quadro e os exibirá em tempo real.
3. Ajuste da resolução
   • A resolução pode ser definida na inicialização (deve ser um múltiplo de 56):

     model = RFDETRBase(resolution=560)
     

   • Quanto maior for a resolução, melhor será a precisão, mas será mais lenta.

Treinamento de modelos personalizados

O RF-DETR suporta o ajuste fino com seu próprio conjunto de dados, mas o conjunto de dados precisa estar no formato COCO, contendo trainevalid responder cantando test Três subdiretórios.

1. Preparação do conjunto de dados
   • Exemplo de estrutura 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
     

   • Os conjuntos de dados no formato COCO podem ser gerados usando a 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 o treinamento
   • Código de amostra:

     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 treinamento, o tamanho total recomendado do lote (batch_size * grad_accum_stepsPor exemplo, as GPUs A100 usam o batch_size=16, grad_accum_steps=1GPUs T4 batch_size=4, grad_accum_steps=4.
3. Treinamento em multi-GPU
   • estabelecer main.py Documentação:

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

   • É executado no terminal:

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

   • comandante-em-chefe (militar) 8 Substitua pelo número de GPUs que você está usando. Observe o ajuste batch_size para manter o tamanho total do lote estável.
4. Resultados do treinamento de carga
   • Dois arquivos de peso são gerados após o treinamento: pesos regulares e pesos EMA (mais estáveis). Método de carregamento:

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

Exportação ONNX

• Exporte para o formato ONNX para facilitar a implementação em outras plataformas:

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

• O arquivo exportado é salvo na pasta output Catálogo para inferência ideal para dispositivos de borda.

cenário do aplicativo

1. direção automática
   O RF-DETR detecta veículos e pedestres na estrada em tempo real. Sua baixa latência e alta precisão são adequadas para sistemas incorporados.
2. controle de qualidade industrial
   O RF-DETR identifica rapidamente os defeitos das peças nas linhas de montagem das fábricas. O modelo é leve e pode ser executado diretamente no equipamento.
3. vigilância por vídeo
   O RF-DETR processa vídeos de vigilância para detectar objetos ou comportamentos anormais em tempo real. Ele suporta streaming de vídeo e é adequado para segurança 24 horas por dia, 7 dias por semana.

QA

1. Quais formatos de conjunto de dados são compatíveis?
   Somente o formato COCO é compatível. O conjunto de dados precisa conter trainevalid responder cantando test subdiretórios, cada um com um _annotations.coco.json Documentação.
2. Como obter a chave da API do Roboflow?
   Faça login em https://app.roboflow.com, localize a chave de API nas configurações de sua conta, copie-a e defina-a na variável de ambiente ROBOFLOW_API_KEY.
3. Quanto tempo dura o treinamento?
   Depende do hardware e do tamanho do conjunto de dados. Em uma GPU T4, 10 épocas podem levar algumas horas. Conjuntos de dados menores podem ser executados em uma CPU, mas é lento.