Views
No views yet
nvidia/mit-b5)1from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
2from PIL import Image
3from datasets import load_dataset
4
5# Load an image from the coralscapes dataset or load your own image
6dataset = load_dataset("EPFL-ECEO/coralscapes")
7image = dataset["test"][42]["image"]
8
9preprocessor = SegformerImageProcessor.from_pretrained("EPFL-ECEO/segformer-b5-finetuned-coralscapes-1024-1024")
10model = SegformerForSemanticSegmentation.from_pretrained("EPFL-ECEO/segformer-b5-finetuned-coralscapes-1024-1024")
11
12inputs = preprocessor(image, return_tensors = "pt")
13outputs = model(**inputs)
14outputs = preprocessor.post_process_semantic_segmentation(outputs, target_sizes=[(image.size[1], image.size[0])])
15label_pred = outputs[0].numpy()1import torch
2import torch.nn.functional as F
3from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
4from PIL import Image
5import numpy as np
6from datasets import load_dataset
7device = 'cuda' if torch.cuda.is_available() else 'cpu'
8
9def resize_image(image, target_size=1024):
10 """
11 Used to resize the image such that the smaller side equals 1024
12 """
13 h_img, w_img = image.size
14 if h_img < w_img:
15 new_h, new_w = target_size, int(w_img * (target_size / h_img))
16 else:
17 new_h, new_w = int(h_img * (target_size / w_img)), target_size
18 resized_img = image.resize((new_h, new_w))
19 return resized_img
20
21def segment_image(image, preprocessor, model, crop_size = (1024, 1024), num_classes = 40, transform=None):
22 """
23 Finds an optimal stride based on the image size and aspect ratio to create
24 overlapping sliding windows of size 1024x1024 which are then fed into the model.
25 """
26 h_crop, w_crop = crop_size
27
28 img = torch.Tensor(np.array(resize_image(image, target_size=1024)).transpose(2, 0, 1)).unsqueeze(0)
29 batch_size, _, h_img, w_img = img.size()
30
31 if transform:
32 img = torch.Tensor(transform(image = img.numpy())["image"]).to(device)
33
34 h_grids = int(np.round(3/2*h_img/h_crop)) if h_img > h_crop else 1
35 w_grids = int(np.round(3/2*w_img/w_crop)) if w_img > w_crop else 1
36
37 h_stride = int((h_img - h_crop + h_grids -1)/(h_grids -1)) if h_grids > 1 else h_crop
38 w_stride = int((w_img - w_crop + w_grids -1)/(w_grids -1)) if w_grids > 1 else w_crop
39
40 preds = img.new_zeros((batch_size, num_classes, h_img, w_img))
41 count_mat = img.new_zeros((batch_size, 1, h_img, w_img))
42
43 for h_idx in range(h_grids):
44 for w_idx in range(w_grids):
45 y1 = h_idx * h_stride
46 x1 = w_idx * w_stride
47 y2 = min(y1 + h_crop, h_img)
48 x2 = min(x1 + w_crop, w_img)
49 y1 = max(y2 - h_crop, 0)
50 x1 = max(x2 - w_crop, 0)
51 crop_img = img[:, :, y1:y2, x1:x2]
52 with torch.no_grad():
53 if(preprocessor):
54 inputs = preprocessor(crop_img, return_tensors = "pt")
55 inputs["pixel_values"] = inputs["pixel_values"].to(device)
56 else:
57 inputs = crop_img.to(device)
58 outputs = model(**inputs)
59
60 resized_logits = F.interpolate(
61 outputs.logits[0].unsqueeze(dim=0), size=crop_img.shape[-2:], mode="bilinear", align_corners=False
62 )
63 preds += F.pad(resized_logits,
64 (int(x1), int(preds.shape[3] - x2), int(y1),
65 int(preds.shape[2] - y2))).cpu()
66 count_mat[:, :, y1:y2, x1:x2] += 1
67
68 assert (count_mat == 0).sum() == 0
69 preds = preds / count_mat
70 preds = preds.argmax(dim=1)
71 preds = F.interpolate(preds.unsqueeze(0).type(torch.uint8), size=image.size[::-1], mode='nearest')
72 label_pred = preds.squeeze().cpu().numpy()
73 return label_pred
74
75# Load an image from the coralscapes dataset or load your own image
76dataset = load_dataset("EPFL-ECEO/coralscapes")
77image = dataset["test"][42]["image"]
78
79preprocessor = SegformerImageProcessor.from_pretrained("EPFL-ECEO/segformer-b5-finetuned-coralscapes-1024-1024")
80model = SegformerForSemanticSegmentation.from_pretrained("EPFL-ECEO/segformer-b5-finetuned-coralscapes-1024-1024")
81
82label_pred = segment_image(image, preprocessor, model)1@misc{sauder2025coralscapesdatasetsemanticscene,
2 title={The Coralscapes Dataset: Semantic Scene Understanding in Coral Reefs},
3 author={Jonathan Sauder and Viktor Domazetoski and Guilhem Banc-Prandi and Gabriela Perna and Anders Meibom and Devis Tuia},
4 year={2025},
5 eprint={2503.20000},
6 archivePrefix={arXiv},
7 primaryClass={cs.CV},
8 url={https://arxiv.org/abs/2503.20000},
9 }