KieDani/SegformerPlusPlus

1import torch.hub
2from PIL import Image
3
4# --- IMPORTANT: TorchHub Dependencies ---
5# Install the dependencies via:
6# pip install tomesd omegaconf numpy rich yapf addict tqdm packaging torchvision
7
8# Load the SegFormer++ model with predefined parameters.
9print("Loading SegFormer++ Model...")
10# Replace 'your_username/your_repo' with the actual path to your repository
11model = torch.hub.load(
12    'KieDani/SegformerPlusPlus',
13    'segformer_plusplus',
14    pretrained=True,
15    backbone='b5',
16    tome_strategy='bsm_hq',
17    checkpoint_url='https://mediastore.rz.uni-augsburg.de/get/yzE65lzm6N/',
18    out_channels=19,
19)
20model.eval()
21print("Model loaded successfully.")
22
23# Load the data transformations via the 'data_transforms' entry point.
24print("Loading data transformations...")
25transform = torch.hub.load(
26    'KieDani/SegformerPlusPlus',
27    'data_transforms',
28)
29print("Transformations loaded successfully.")
30
31# --- Example for Image Preparation and Inference ---
32# Create a dummy image, as we don't need a real image file.
33# In a real scenario, you would load an image from the hard drive, e.g.
34# from PIL import Image
35# image = Image.open('path_to_your_image.jpg').convert('RGB')
36print("Creating a dummy image for demonstration...")
37dummy_image = Image.new('RGB', (1300, 1300), color='red')
38print("Original image size:", dummy_image.size)
39
40# Apply the transformations loaded from the Hub to the image.
41print("Applying transformations to the image...")
42input_tensor = transform(dummy_image).unsqueeze(0)  # Adds a batch dimension
43print("Transformed image tensor size:", input_tensor.shape)
44
45# Run inference.
46print("Running inference...")
47with torch.no_grad():
48    output = model(input_tensor)
49
50# The output tensor has the shape [1, num_classes, height, width]
51# We remove the batch dimension (1)
52output_tensor = output.squeeze(0)
53
54print(f"\nInference completed. Output tensor size: {output_tensor.shape}")
55
56# To get the final segmentation map, you would use argmax.
57segmentation_map = torch.argmax(output_tensor, dim=0)
58print(f"Size of the generated segmentation map: {segmentation_map.shape}")
59