SegFormer (b4-sized) encoder pre-entrenado

nvidia
Clasificación de imagen

SegFormer consiste en un codificador Transformer jerárquico y una cabeza de decodificación completamente de MLP ligera para lograr excelentes resultados en los benchmarks de segmentación semántica como ADE20K y Cityscapes. El Transformer jerárquico se pre-entrenó primero en ImageNet-1k, después de lo cual se agregó una cabeza de decodificación y se afinó junto con un conjunto de datos descendente. Este repositorio solo contiene el Transformer jerárquico pre-entrenado, por lo tanto, se puede usar con fines de ajuste fino.

Como usar

Aquí está 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-b4")
model = SegformerForImageClassification.from_pretrained("nvidia/mit-b4")

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

Para más ejemplos de código, consultar la documentación.

Funcionalidades

Codificador Transformer jerárquico
Cabeza de decodificación ligera completamente de MLP
Capaz de lograr excelentes resultados en benchmarks de segmentación semántica como ADE20K y Cityscapes

Casos de uso

Ajuste fino de segmentación semántica