File size: 1,221 Bytes
883856e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | import torch
import torch.nn as nn
from torchvision.models.segmentation import deeplabv3_resnet50, DeepLabV3_ResNet50_Weights
class ResNetSegmentation(nn.Module):
def __init__(self):
super().__init__()
# Load Pretrained DeepLabV3 with ResNet50 backbone
# weights="DEFAULT" loads the best available weights (COCO)
self.model = deeplabv3_resnet50(weights=DeepLabV3_ResNet50_Weights.DEFAULT)
# Replace the classifier head for binary segmentation (1 class + background? No, usually binary is 1 channel)
# DeepLabV3 classifier is DeepLabHead sequentially:
# 0: ASPP...
# 1: Conv2d(256, 256, 3, padding=1)
# 2: BatchNorm
# 3: ReLU
# 4: Conv2d(256, num_classes, 1)
# We need final output channel = 1 (Binary mask)
self.model.classifier[4] = nn.Conv2d(256, 1, kernel_size=1)
# Also auxiliary classifier needs to be changed if present (it is by default)
self.model.aux_classifier[4] = nn.Conv2d(256, 1, kernel_size=1)
def forward(self, x):
# DeepLabV3 returns an OrderedDict with keys 'out' and 'aux'
return self.model(x)['out']
|