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']