SegFormer (encoder solamente preentrenado de tamaño b1)

nvidia

Clasificación de imagen

SegFormer (encoder solamente preentrenado de tamaño b1) es un modelo de encoder con un Transformer jerárquico preentrenado en ImageNet-1k. SegFormer consiste en un encoder basado en Transformer y un cabeza de decodificación todo-MLP ligero que logra excelentes resultados en benchmarks de segmentación semántica como ADE20K y Cityscapes. Este repositorio solo contiene el Transformer jerárquico preentrenado, por lo que puede usarse para propósitos de ajuste fino.

Como usar

Aquí se explica cómo usar este modelo para clasificar una imagen del conjunto de datos COCO 2017 en una de las 1,000 clases de ImageNet:
from transformers import SegformerFeatureExtractor, SegformerForImageClassification
from PIL import Image
import requests

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

feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/mit-b1")
model = SegformerForImageClassification.from_pretrained("nvidia/mit-b1")

inputs = feature_extractor(images=image, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits
# el modelo predice una de las 1000 clases de ImageNet
predicted_class_idx = logits.argmax(-1).item()
print("Clase predicha:", model.config.id2label[predicted_class_idx])

Funcionalidades

Encoder basado en Transformer jerárquico
Cabeza de decodificación ligera todo-MLP
Preentrenado en ImageNet-1k

Casos de uso

Ajuste fino para segmentación semántica