Views
No views yet
1from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
2import torch
3import cv2
4import numpy as np
5
6# Load model and processor
7model = SegformerForSemanticSegmentation.from_pretrained("simhaq-trmb/segformer-parker")
8processor = SegformerImageProcessor.from_pretrained("simhaq-trmb/segformer-parker")
9
10# Load image
11image = cv2.imread("your_image.jpg")
12image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
13image_resized = cv2.resize(image_rgb, (512, 512))
14
15# Process the image
16inputs = processor(images=image_resized, return_tensors="pt")
17
18# Perform inference
19with torch.no_grad():
20 outputs = model(pixel_values=inputs["pixel_values"])
21
22 # Interpolate the logits to match the input resolution
23 upsampled_logits = torch.nn.functional.interpolate(
24 outputs.logits,
25 size=(512, 512),
26 mode="bilinear",
27 align_corners=False
28 )
29
30 # Get the segmentation mask
31 segmentation_mask = upsampled_logits.argmax(dim=1).squeeze().cpu().numpy()| Class ID | Color | Description |
|---|---|---|
| 0 | Blue | Sky |
| 1 | Green | Dense Foliage |
| 2 | Black | Obstruction |
| 3 | Yellow | Sparse Foliage |