deeplabv3

import torch
from torchvision import models, transforms
from PIL import Image
from network.utils import IntermediateLayerGetter

# Load the image
img_path = ''
input_image = Image.open(img_path).convert("RGB") 

# input image = preprocessed stage image

preprocess = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0)  # create a mini-batch as expected by the model

# Check if GPU is available and use it
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
input_batch = input_batch.to(device)

# Load model
model = models.segmentation.deeplabv3_resnet101(pretrained=True).to(device)
model.eval()

# Define the layers to extract
return_layers = {'layer1': 'feat1', 'layer4': 'feat2'} # optional names for the layers
intermediate_layer_getter = IntermediateLayerGetter(model.backbone, return_layers=return_layers)


with torch.no_grad():
    features = intermediate_layer_getter(input_batch)


for name, feature in features.items():
    print(f"{name}: {feature.shape}")

# feat1: torch.Size([1, 256, 64, 64])
# feat2: torch.Size([1, 2048, 32, 32])