YOLOS (pequeño tamaño) modelo (300 épocas preentrenamiento)

hustvl
Detección de objetos

YOLOS es un Transformer de Visión (ViT) entrenado usando la pérdida DETR. A pesar de su simplicidad, un modelo YOLOS de tamaño base es capaz de alcanzar 42 AP en la validación de COCO 2017 (similar a DETR y marcos más complejos como Faster R-CNN). El modelo está entrenado usando una 'pérdida de emparejamiento bipartito': se comparan las clases predichas + cajas delimitadoras de cada una de las N = 100 consultas de objeto con las anotaciones de verdad en el terreno, acolchadas hasta la misma longitud N. El algoritmo de emparejamiento húngaro se utiliza para crear una correspondencia óptima uno-a-uno entre cada una de las consultas N y cada una de las anotaciones N. A continuación, se utilizan la entropía cruzada estándar (para las clases) y una combinación lineal de la pérdida L1 y IoU generalizada (para las cajas delimitadoras) para optimizar los parámetros del modelo.

Como usar

Aquí está cómo usar este modelo:

from transformers import YolosImageProcessor, YolosForObjectDetection
from PIL import Image
import requests

url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
image = Image.open(requests.get(url, stream=True).raw)

image_processor = YolosImageProcessor.from_pretrained('hustvl/yolos-small-300')
model = YolosForObjectDetection.from_pretrained('hustvl/yolos-small-300')

inputs = image_processor(images=image, return_tensors="pt")
outputs = model(**inputs)

# el modelo predice cajas delimitadoras y clases COCO correspondientes
logits = outputs.logits
bboxes = outputs.pred_boxes

Actualmente, tanto el procesador de imágenes como el modelo soportan PyTorch.

Funcionalidades

Transformers para detección de objetos
Entrenado usando pérdida DETR
Alcanza 42 AP en la validación de COCO 2017
Modelo compatible con PyTorch

Casos de uso

Detección de objetos en imágenes