Views
No views yet
nvidia/segformer-b0-finetuned-ade-512-512 on a custom dataset with 3 semantic classes:nvidia/segformer-b0-finetuned-ade-512-5121from transformers import AutoModelForSemanticSegmentation, AutoImageProcessor
2from PIL import Image
3import torch
4
5# Load model and processor
6model = AutoModelForSemanticSegmentation.from_pretrained("Wilbur1240/segformer-b0-finetuned-ade-512-512-finetune-mastr1325")
7processor = AutoImageProcessor.from_pretrained("Wilbur1240/segformer-b0-finetuned-ade-512-512-finetune-mastr1325")
8
9# Load and preprocess an image
10image = Image.open("example.jpg").convert("RGB")
11inputs = processor(images=image, return_tensors="pt")
12
13# Inference
14with torch.no_grad():
15 outputs = model(**inputs)
16 logits = outputs.logits # [1, num_classes, H, W]
17 pred_seg = logits.argmax(dim=1) # [1, H, W]