Semantic Segmentation

1import torch
2import torch.nn as nn
3import torchvision.transforms as transforms
4from PIL import Image
5import numpy as np
6
7# Semantic segmentation output format
8# Shape: [batch_size, num_classes, height, width]
9# Each pixel has probabilities for each class
10
11class SegmentationHead(nn.Module):
12    """Simple segmentation head for pixel classification."""
13
14    def __init__(self, in_channels, num_classes):
15        super().__init__()
16        self.conv = nn.Conv2d(in_channels, num_classes, kernel_size=1)
17
18    def forward(self, x):
19        return self.conv(x)
20
21# Example: Convert logits to segmentation mask
22def logits_to_mask(logits):
23    """Convert model output to segmentation mask."""
24    # logits shape: [B, C, H, W]
25    # Apply softmax and argmax
26    probs = torch.softmax(logits, dim=1)
27    mask = torch.argmax(probs, dim=1)  # [B, H, W]
28    return mask
29
30# Color mapping for visualization
31def create_color_map(num_classes):
32    """Create color palette for segmentation visualization."""
33    colors = np.zeros((num_classes, 3), dtype=np.uint8)
34    for i in range(num_classes):
35        colors[i] = [
36            (i * 100) % 256,
37            (i * 50 + 100) % 256,
38            (i * 150 + 50) % 256
39        ]
40    return colors
41
42def visualize_segmentation(image, mask, color_map):
43    """Overlay segmentation mask on image."""
44    # Convert mask to colored image
45    h, w = mask.shape
46    colored_mask = np.zeros((h, w, 3), dtype=np.uint8)
47
48    for class_idx in range(len(color_map)):
49        colored_mask[mask == class_idx] = color_map[class_idx]
50
51    # Blend with original image
52    alpha = 0.5
53    blended = (alpha * image + (1 - alpha) * colored_mask).astype(np.uint8)
54    return blended

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4import torchvision.models as models
5
6class FCN8s(nn.Module):
7    """
8    FCN-8s architecture for semantic segmentation.
9    Uses VGG16 backbone with skip connections.
10    """
11
12    def __init__(self, num_classes=21, pretrained_backbone=True):
13        super().__init__()
14
15        # Load pretrained VGG16
16        vgg16 = models.vgg16(pretrained=pretrained_backbone)
17        features = list(vgg16.features.children())
18
19        # Encoder stages (pool indices: 4, 9, 16, 23, 30)
20        self.stage1 = nn.Sequential(*features[:5])   # -> pool1
21        self.stage2 = nn.Sequential(*features[5:10])  # -> pool2
22        self.stage3 = nn.Sequential(*features[10:17]) # -> pool3
23        self.stage4 = nn.Sequential(*features[17:24]) # -> pool4
24        self.stage5 = nn.Sequential(*features[24:])   # -> pool5
25
26        # Replace FC layers with conv layers
27        self.fc6 = nn.Sequential(
28            nn.Conv2d(512, 4096, kernel_size=7, padding=3),
29            nn.ReLU(inplace=True),
30            nn.Dropout2d()
31        )
32        self.fc7 = nn.Sequential(
33            nn.Conv2d(4096, 4096, kernel_size=1),
34            nn.ReLU(inplace=True),
35            nn.Dropout2d()
36        )
37
38        # Scoring layers
39        self.score_fc7 = nn.Conv2d(4096, num_classes, kernel_size=1)
40        self.score_pool4 = nn.Conv2d(512, num_classes, kernel_size=1)
41        self.score_pool3 = nn.Conv2d(256, num_classes, kernel_size=1)
42
43        # Upsampling layers (transposed convolution)
44        self.upsample_2x_fc7 = nn.ConvTranspose2d(
45            num_classes, num_classes, kernel_size=4, stride=2, padding=1
46        )
47        self.upsample_2x_pool4 = nn.ConvTranspose2d(
48            num_classes, num_classes, kernel_size=4, stride=2, padding=1
49        )
50        self.upsample_8x = nn.ConvTranspose2d(
51            num_classes, num_classes, kernel_size=16, stride=8, padding=4
52        )
53
54        self._initialize_weights()
55
56    def _initialize_weights(self):
57        for m in [self.score_fc7, self.score_pool4, self.score_pool3]:
58            nn.init.zeros_(m.weight)
59            nn.init.zeros_(m.bias)
60
61    def forward(self, x):
62        # Encoder
63        x1 = self.stage1(x)   # 1/2
64        x2 = self.stage2(x1)  # 1/4
65        x3 = self.stage3(x2)  # 1/8  (pool3)
66        x4 = self.stage4(x3)  # 1/16 (pool4)
67        x5 = self.stage5(x4)  # 1/32 (pool5)
68
69        # FC layers (convolutionalized)
70        x = self.fc6(x5)
71        x = self.fc7(x)
72
73        # Score from fc7 (1/32 resolution)
74        score_fc7 = self.score_fc7(x)
75
76        # Upsample fc7 score to 1/16 and fuse with pool4
77        upscore_fc7 = self.upsample_2x_fc7(score_fc7)
78        score_pool4 = self.score_pool4(x4)
79        fuse_pool4 = upscore_fc7 + score_pool4
80
81        # Upsample fused to 1/8 and fuse with pool3
82        upscore_pool4 = self.upsample_2x_pool4(fuse_pool4)
83        score_pool3 = self.score_pool3(x3)
84        fuse_pool3 = upscore_pool4 + score_pool3
85
86        # Final 8x upsampling to original resolution
87        output = self.upsample_8x(fuse_pool3)
88
89        return output
90
91# Usage example
92model = FCN8s(num_classes=21)
93x = torch.randn(1, 3, 256, 256)
94output = model(x)
95print(f"Input: {x.shape} -> Output: {output.shape}")
96# Output: torch.Size([1, 21, 256, 256])

1import torch
2import torch.nn as nn
3
4class DoubleConv(nn.Module):
5    """Two consecutive conv-bn-relu blocks."""
6
7    def __init__(self, in_channels, out_channels, mid_channels=None):
8        super().__init__()
9        if mid_channels is None:
10            mid_channels = out_channels
11
12        self.double_conv = nn.Sequential(
13            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
14            nn.BatchNorm2d(mid_channels),
15            nn.ReLU(inplace=True),
16            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
17            nn.BatchNorm2d(out_channels),
18            nn.ReLU(inplace=True)
19        )
20
21    def forward(self, x):
22        return self.double_conv(x)
23
24class Down(nn.Module):
25    """Downsampling block: maxpool + double conv."""
26
27    def __init__(self, in_channels, out_channels):
28        super().__init__()
29        self.maxpool_conv = nn.Sequential(
30            nn.MaxPool2d(2),
31            DoubleConv(in_channels, out_channels)
32        )
33
34    def forward(self, x):
35        return self.maxpool_conv(x)
36
37class Up(nn.Module):
38    """Upsampling block: upsample + concatenate + double conv."""
39
40    def __init__(self, in_channels, out_channels, bilinear=True):
41        super().__init__()
42
43        if bilinear:
44            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
45            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
46        else:
47            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
48            self.conv = DoubleConv(in_channels, out_channels)
49
50    def forward(self, x1, x2):
51        x1 = self.up(x1)
52
53        # Handle size mismatch due to odd dimensions
54        diff_y = x2.size()[2] - x1.size()[2]
55        diff_x = x2.size()[3] - x1.size()[3]
56        x1 = F.pad(x1, [diff_x // 2, diff_x - diff_x // 2,
57                       diff_y // 2, diff_y - diff_y // 2])
58
59        # Concatenate along channel dimension
60        x = torch.cat([x2, x1], dim=1)
61        return self.conv(x)
62
63class UNet(nn.Module):
64    """
65    U-Net architecture for semantic segmentation.
66
67    Architecture:
68        Input -> Encoder (contracting) -> Bottleneck -> Decoder (expanding) -> Output
69
70    Skip connections concatenate encoder features to decoder.
71    """
72
73    def __init__(self, n_channels=3, n_classes=2, bilinear=True):
74        super().__init__()
75        self.n_channels = n_channels
76        self.n_classes = n_classes
77        self.bilinear = bilinear
78
79        factor = 2 if bilinear else 1
80
81        # Encoder (contracting path)
82        self.inc = DoubleConv(n_channels, 64)
83        self.down1 = Down(64, 128)
84        self.down2 = Down(128, 256)
85        self.down3 = Down(256, 512)
86        self.down4 = Down(512, 1024 // factor)
87
88        # Decoder (expanding path)
89        self.up1 = Up(1024, 512 // factor, bilinear)
90        self.up2 = Up(512, 256 // factor, bilinear)
91        self.up3 = Up(256, 128 // factor, bilinear)
92        self.up4 = Up(128, 64, bilinear)
93
94        # Output layer
95        self.outc = nn.Conv2d(64, n_classes, kernel_size=1)
96
97    def forward(self, x):
98        # Encoder
99        x1 = self.inc(x)    # 64 channels
100        x2 = self.down1(x1) # 128 channels, 1/2 resolution
101        x3 = self.down2(x2) # 256 channels, 1/4 resolution
102        x4 = self.down3(x3) # 512 channels, 1/8 resolution
103        x5 = self.down4(x4) # 512/1024 channels, 1/16 resolution
104
105        # Decoder with skip connections
106        x = self.up1(x5, x4) # Concatenate with x4
107        x = self.up2(x, x3)  # Concatenate with x3
108        x = self.up3(x, x2)  # Concatenate with x2
109        x = self.up4(x, x1)  # Concatenate with x1
110
111        # Output logits
112        logits = self.outc(x)
113        return logits
114
115# Example usage
116model = UNet(n_channels=3, n_classes=21)
117x = torch.randn(2, 3, 256, 256)
118output = model(x)
119print(f"U-Net: {x.shape} -> {output.shape}")
120
121# Count parameters
122params = sum(p.numel() for p in model.parameters())
123print(f"Total parameters: {params:,}")

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5class AtrousConv(nn.Module):
6    """
7    Atrous (Dilated) Convolution.
8    Increases receptive field without increasing parameters.
9    """
10
11    def __init__(self, in_channels, out_channels, kernel_size=3, dilation=1):
12        super().__init__()
13        padding = dilation * (kernel_size - 1) // 2
14
15        self.conv = nn.Sequential(
16            nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size,
17                     padding=padding, dilation=dilation, bias=False),
18            nn.BatchNorm2d(out_channels),
19            nn.ReLU(inplace=True)
20        )
21
22    def forward(self, x):
23        return self.conv(x)
24
25class ASPPModule(nn.Module):
26    """
27    Atrous Spatial Pyramid Pooling (ASPP).
28    Captures multi-scale information with parallel atrous convolutions.
29    """
30
31    def __init__(self, in_channels, out_channels, dilations=[1, 6, 12, 18]):
32        super().__init__()
33
34        modules = []
35
36        # 1x1 convolution
37        modules.append(nn.Sequential(
38            nn.Conv2d(in_channels, out_channels, 1, bias=False),
39            nn.BatchNorm2d(out_channels),
40            nn.ReLU(inplace=True)
41        ))
42
43        # Atrous convolutions at different rates
44        for dilation in dilations[1:]:
45            modules.append(AtrousConv(in_channels, out_channels, dilation=dilation))
46
47        # Global average pooling branch (image-level features)
48        modules.append(nn.Sequential(
49            nn.AdaptiveAvgPool2d(1),
50            nn.Conv2d(in_channels, out_channels, 1, bias=False),
51            nn.BatchNorm2d(out_channels),
52            nn.ReLU(inplace=True)
53        ))
54
55        self.convs = nn.ModuleList(modules)
56
57        # Projection layer
58        self.project = nn.Sequential(
59            nn.Conv2d(out_channels * (len(dilations) + 1), out_channels, 1, bias=False),
60            nn.BatchNorm2d(out_channels),
61            nn.ReLU(inplace=True),
62            nn.Dropout(0.5)
63        )
64
65    def forward(self, x):
66        res = []
67        for conv in self.convs[:-1]:
68            res.append(conv(x))
69
70        # Image-level features need upsampling
71        global_feat = self.convs[-1](x)
72        global_feat = F.interpolate(global_feat, size=x.shape[2:],
73                                    mode='bilinear', align_corners=False)
74        res.append(global_feat)
75
76        # Concatenate all branches
77        x = torch.cat(res, dim=1)
78        return self.project(x)
79
80class DeepLabV3Plus(nn.Module):
81    """
82    DeepLabV3+ architecture with encoder-decoder structure.
83    Combines ASPP with low-level features for sharp boundaries.
84    """
85
86    def __init__(self, num_classes=21, backbone='resnet50'):
87        super().__init__()
88
89        # Use ResNet backbone (simplified for illustration)
90        import torchvision.models as models
91        resnet = models.resnet50(pretrained=True)
92
93        # Encoder layers
94        self.layer0 = nn.Sequential(
95            resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool
96        )
97        self.layer1 = resnet.layer1  # Low-level features (1/4)
98        self.layer2 = resnet.layer2  # (1/8)
99        self.layer3 = resnet.layer3  # (1/16)
100        self.layer4 = resnet.layer4  # High-level features (1/16 with dilation)
101
102        # Modify layer4 for output stride 16
103        for m in self.layer4.modules():
104            if isinstance(m, nn.Conv2d) and m.stride == (2, 2):
105                m.stride = (1, 1)
106                m.dilation = (2, 2)
107                m.padding = (2, 2)
108
109        # ASPP module
110        self.aspp = ASPPModule(2048, 256)
111
112        # Decoder
113        self.low_level_conv = nn.Sequential(
114            nn.Conv2d(256, 48, 1, bias=False),
115            nn.BatchNorm2d(48),
116            nn.ReLU(inplace=True)
117        )
118
119        self.decoder = nn.Sequential(
120            nn.Conv2d(256 + 48, 256, 3, padding=1, bias=False),
121            nn.BatchNorm2d(256),
122            nn.ReLU(inplace=True),
123            nn.Dropout(0.5),
124            nn.Conv2d(256, 256, 3, padding=1, bias=False),
125            nn.BatchNorm2d(256),
126            nn.ReLU(inplace=True),
127            nn.Dropout(0.1),
128            nn.Conv2d(256, num_classes, 1)
129        )
130
131    def forward(self, x):
132        input_shape = x.shape[-2:]
133
134        # Encoder
135        x = self.layer0(x)
136        low_level = self.layer1(x)  # 1/4 resolution
137        x = self.layer2(low_level)
138        x = self.layer3(x)
139        x = self.layer4(x)  # 1/16 resolution
140
141        # ASPP
142        x = self.aspp(x)
143
144        # Decoder
145        x = F.interpolate(x, size=low_level.shape[-2:],
146                         mode='bilinear', align_corners=False)
147        low_level = self.low_level_conv(low_level)
148        x = torch.cat([x, low_level], dim=1)
149        x = self.decoder(x)
150
151        # Final upsampling
152        x = F.interpolate(x, size=input_shape,
153                         mode='bilinear', align_corners=False)
154        return x
155
156# Example
157model = DeepLabV3Plus(num_classes=21)
158x = torch.randn(1, 3, 512, 512)
159output = model(x)
160print(f"DeepLabV3+: {x.shape} -> {output.shape}")

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5class DiceLoss(nn.Module):
6    """
7    Dice Loss for semantic segmentation.
8    Handles class imbalance well.
9
10    Dice = 2 * |A intersection B| / (|A| + |B|)
11    """
12
13    def __init__(self, smooth=1e-6):
14        super().__init__()
15        self.smooth = smooth
16
17    def forward(self, pred, target):
18        # pred: [B, C, H, W] (logits)
19        # target: [B, H, W] (class indices)
20
21        num_classes = pred.shape[1]
22
23        # Convert to probabilities
24        pred = F.softmax(pred, dim=1)
25
26        # One-hot encode target
27        target_one_hot = F.one_hot(target, num_classes)  # [B, H, W, C]
28        target_one_hot = target_one_hot.permute(0, 3, 1, 2).float()  # [B, C, H, W]
29
30        # Calculate Dice for each class
31        dims = (0, 2, 3)  # Batch, Height, Width
32        intersection = torch.sum(pred * target_one_hot, dim=dims)
33        cardinality = torch.sum(pred + target_one_hot, dim=dims)
34
35        dice_score = (2. * intersection + self.smooth) / (cardinality + self.smooth)
36
37        # Average over classes
38        return 1 - dice_score.mean()
39
40class FocalLoss(nn.Module):
41    """
42    Focal Loss for handling class imbalance.
43    Focuses training on hard examples.
44
45    FL(p_t) = -alpha_t * (1 - p_t)^gamma * log(p_t)
46    """
47
48    def __init__(self, alpha=None, gamma=2.0, ignore_index=-100):
49        super().__init__()
50        self.alpha = alpha  # Class weights
51        self.gamma = gamma  # Focusing parameter
52        self.ignore_index = ignore_index
53
54    def forward(self, pred, target):
55        # pred: [B, C, H, W]
56        # target: [B, H, W]
57
58        ce_loss = F.cross_entropy(pred, target, reduction='none',
59                                  ignore_index=self.ignore_index)
60
61        # Get probabilities
62        pt = torch.exp(-ce_loss)
63
64        # Apply focal weight
65        focal_weight = (1 - pt) ** self.gamma
66        focal_loss = focal_weight * ce_loss
67
68        # Apply class weights if provided
69        if self.alpha is not None:
70            alpha_t = self.alpha[target]
71            focal_loss = alpha_t * focal_loss
72
73        return focal_loss.mean()
74
75class CombinedLoss(nn.Module):
76    """
77    Combined loss: Cross-Entropy + Dice Loss.
78    Balances pixel-wise accuracy and region overlap.
79    """
80
81    def __init__(self, ce_weight=0.5, dice_weight=0.5, class_weights=None):
82        super().__init__()
83        self.ce_weight = ce_weight
84        self.dice_weight = dice_weight
85        self.ce_loss = nn.CrossEntropyLoss(weight=class_weights)
86        self.dice_loss = DiceLoss()
87
88    def forward(self, pred, target):
89        ce = self.ce_loss(pred, target)
90        dice = self.dice_loss(pred, target)
91        return self.ce_weight * ce + self.dice_weight * dice
92
93class BoundaryLoss(nn.Module):
94    """
95    Boundary Loss penalizes errors at object boundaries.
96    Useful for precise boundary delineation.
97    """
98
99    def __init__(self):
100        super().__init__()
101        # Sobel filters for edge detection
102        sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32)
103        sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32)
104
105        self.register_buffer('sobel_x', sobel_x.view(1, 1, 3, 3))
106        self.register_buffer('sobel_y', sobel_y.view(1, 1, 3, 3))
107
108    def get_boundary(self, mask):
109        """Extract boundaries using Sobel operator."""
110        mask = mask.float().unsqueeze(1)  # [B, 1, H, W]
111
112        grad_x = F.conv2d(mask, self.sobel_x, padding=1)
113        grad_y = F.conv2d(mask, self.sobel_y, padding=1)
114
115        boundary = torch.sqrt(grad_x ** 2 + grad_y ** 2)
116        boundary = (boundary > 0.5).float()
117
118        return boundary.squeeze(1)
119
120    def forward(self, pred, target):
121        pred_mask = torch.argmax(pred, dim=1)
122
123        pred_boundary = self.get_boundary(pred_mask)
124        target_boundary = self.get_boundary(target)
125
126        # Binary cross-entropy on boundaries
127        loss = F.binary_cross_entropy_with_logits(
128            pred_boundary, target_boundary
129        )
130        return loss
131
132# Example usage
133criterion = CombinedLoss(ce_weight=0.5, dice_weight=0.5)
134pred = torch.randn(4, 21, 256, 256)  # Model output
135target = torch.randint(0, 21, (4, 256, 256))  # Ground truth
136
137loss = criterion(pred, target)
138print(f"Combined Loss: {loss.item():.4f}")

1import torch
2import numpy as np
3from collections import defaultdict
4
5class SegmentationMetrics:
6    """
7    Comprehensive metrics for semantic segmentation evaluation.
8    """
9
10    def __init__(self, num_classes, ignore_index=255):
11        self.num_classes = num_classes
12        self.ignore_index = ignore_index
13        self.reset()
14
15    def reset(self):
16        """Reset all counters."""
17        self.confusion_matrix = np.zeros((self.num_classes, self.num_classes))
18
19    def update(self, pred, target):
20        """
21        Update confusion matrix with batch predictions.
22
23        Args:
24            pred: [B, H, W] predicted class indices
25            target: [B, H, W] ground truth class indices
26        """
27        pred = pred.cpu().numpy().flatten()
28        target = target.cpu().numpy().flatten()
29
30        # Remove ignored pixels
31        mask = (target != self.ignore_index)
32        pred = pred[mask]
33        target = target[mask]
34
35        # Update confusion matrix
36        for p, t in zip(pred, target):
37            if 0 <= p < self.num_classes and 0 <= t < self.num_classes:
38                self.confusion_matrix[t, p] += 1
39
40    def pixel_accuracy(self):
41        """Overall pixel accuracy."""
42        correct = np.diag(self.confusion_matrix).sum()
43        total = self.confusion_matrix.sum()
44        return correct / (total + 1e-10)
45
46    def class_pixel_accuracy(self):
47        """Per-class pixel accuracy."""
48        return np.diag(self.confusion_matrix) / (
49            self.confusion_matrix.sum(axis=1) + 1e-10
50        )
51
52    def mean_pixel_accuracy(self):
53        """Mean class pixel accuracy."""
54        class_acc = self.class_pixel_accuracy()
55        return np.nanmean(class_acc)
56
57    def intersection_over_union(self):
58        """Per-class IoU (Jaccard Index)."""
59        intersection = np.diag(self.confusion_matrix)
60        union = (
61            self.confusion_matrix.sum(axis=0) +
62            self.confusion_matrix.sum(axis=1) -
63            intersection
64        )
65        return intersection / (union + 1e-10)
66
67    def mean_iou(self):
68        """Mean Intersection over Union (mIoU)."""
69        iou = self.intersection_over_union()
70        return np.nanmean(iou)
71
72    def frequency_weighted_iou(self):
73        """Frequency-weighted IoU."""
74        freq = self.confusion_matrix.sum(axis=1) / (
75            self.confusion_matrix.sum() + 1e-10
76        )
77        iou = self.intersection_over_union()
78        return (freq * iou).sum()
79
80    def dice_coefficient(self):
81        """Per-class Dice coefficient (F1 score)."""
82        intersection = np.diag(self.confusion_matrix)
83        pred_sum = self.confusion_matrix.sum(axis=0)
84        target_sum = self.confusion_matrix.sum(axis=1)
85        return 2 * intersection / (pred_sum + target_sum + 1e-10)
86
87    def mean_dice(self):
88        """Mean Dice coefficient."""
89        dice = self.dice_coefficient()
90        return np.nanmean(dice)
91
92    def get_results(self):
93        """Get all metrics as dictionary."""
94        return {
95            'pixel_accuracy': self.pixel_accuracy(),
96            'mean_pixel_accuracy': self.mean_pixel_accuracy(),
97            'mean_iou': self.mean_iou(),
98            'frequency_weighted_iou': self.frequency_weighted_iou(),
99            'mean_dice': self.mean_dice(),
100            'per_class_iou': self.intersection_over_union(),
101            'per_class_dice': self.dice_coefficient()
102        }
103
104def compute_iou_batch(pred, target, num_classes):
105    """
106    Fast IoU computation for a batch.
107
108    Args:
109        pred: [B, H, W] predicted class indices
110        target: [B, H, W] ground truth
111        num_classes: number of classes
112
113    Returns:
114        IoU for each class
115    """
116    ious = []
117
118    for cls in range(num_classes):
119        pred_cls = (pred == cls)
120        target_cls = (target == cls)
121
122        intersection = (pred_cls & target_cls).sum().float()
123        union = (pred_cls | target_cls).sum().float()
124
125        if union > 0:
126            ious.append((intersection / union).item())
127        else:
128            ious.append(float('nan'))
129
130    return ious
131
132# Example evaluation
133metrics = SegmentationMetrics(num_classes=21)
134
135# Simulate predictions and targets
136for _ in range(10):
137    pred = torch.randint(0, 21, (4, 256, 256))
138    target = torch.randint(0, 21, (4, 256, 256))
139    metrics.update(pred, target)
140
141results = metrics.get_results()
142print(f"Pixel Accuracy: {results['pixel_accuracy']:.4f}")
143print(f"Mean IoU: {results['mean_iou']:.4f}")
144print(f"Mean Dice: {results['mean_dice']:.4f}")
145print(f"Frequency Weighted IoU: {results['frequency_weighted_iou']:.4f}")
146
147# Print per-class IoU
148print("\nPer-class IoU:")
149for i, iou in enumerate(results['per_class_iou'][:5]):
150    print(f"  Class {i}: {iou:.4f}")

1import torch
2import torch.nn as nn
3import torch.optim as optim
4from torch.utils.data import Dataset, DataLoader
5from torch.cuda.amp import autocast, GradScaler
6import albumentations as A
7from albumentations.pytorch import ToTensorV2
8import numpy as np
9from tqdm import tqdm
10
11class SegmentationDataset(Dataset):
12    """Custom dataset for semantic segmentation."""
13
14    def __init__(self, images, masks, transform=None):
15        self.images = images  # List of image paths or arrays
16        self.masks = masks    # List of mask paths or arrays
17        self.transform = transform
18
19    def __len__(self):
20        return len(self.images)
21
22    def __getitem__(self, idx):
23        image = self.images[idx]  # Load image
24        mask = self.masks[idx]    # Load mask
25
26        if self.transform:
27            augmented = self.transform(image=image, mask=mask)
28            image = augmented['image']
29            mask = augmented['mask']
30
31        return image, mask.long()
32
33def get_training_augmentation():
34    """Training-time augmentations."""
35    return A.Compose([
36        A.Resize(512, 512),
37        A.HorizontalFlip(p=0.5),
38        A.VerticalFlip(p=0.5),
39        A.RandomRotate90(p=0.5),
40        A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=30, p=0.5),
41        A.OneOf([
42            A.RandomBrightnessContrast(p=1),
43            A.RandomGamma(p=1),
44            A.HueSaturationValue(p=1),
45        ], p=0.5),
46        A.OneOf([
47            A.GaussNoise(p=1),
48            A.GaussianBlur(p=1),
49            A.MotionBlur(p=1),
50        ], p=0.3),
51        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
52        ToTensorV2()
53    ])
54
55def get_validation_augmentation():
56    """Validation-time augmentations (no randomness)."""
57    return A.Compose([
58        A.Resize(512, 512),
59        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
60        ToTensorV2()
61    ])
62
63class SegmentationTrainer:
64    """Training handler for segmentation models."""
65
66    def __init__(self, model, criterion, optimizer, scheduler, device, num_classes):
67        self.model = model.to(device)
68        self.criterion = criterion
69        self.optimizer = optimizer
70        self.scheduler = scheduler
71        self.device = device
72        self.scaler = GradScaler()  # For mixed precision
73        self.metrics = SegmentationMetrics(num_classes)
74
75    def train_epoch(self, dataloader):
76        """Train for one epoch."""
77        self.model.train()
78        total_loss = 0
79
80        pbar = tqdm(dataloader, desc='Training')
81        for images, masks in pbar:
82            images = images.to(self.device)
83            masks = masks.to(self.device)
84
85            self.optimizer.zero_grad()
86
87            # Mixed precision training
88            with autocast():
89                outputs = self.model(images)
90                loss = self.criterion(outputs, masks)
91
92            self.scaler.scale(loss).backward()
93            self.scaler.step(self.optimizer)
94            self.scaler.update()
95
96            total_loss += loss.item()
97            pbar.set_postfix({'loss': loss.item()})
98
99        return total_loss / len(dataloader)
100
101    @torch.no_grad()
102    def validate(self, dataloader):
103        """Validate the model."""
104        self.model.eval()
105        self.metrics.reset()
106        total_loss = 0
107
108        for images, masks in tqdm(dataloader, desc='Validation'):
109            images = images.to(self.device)
110            masks = masks.to(self.device)
111
112            with autocast():
113                outputs = self.model(images)
114                loss = self.criterion(outputs, masks)
115
116            total_loss += loss.item()
117
118            # Update metrics
119            pred = torch.argmax(outputs, dim=1)
120            self.metrics.update(pred, masks)
121
122        results = self.metrics.get_results()
123        results['loss'] = total_loss / len(dataloader)
124        return results
125
126    def fit(self, train_loader, val_loader, epochs, save_path='best_model.pth'):
127        """Full training loop."""
128        best_miou = 0
129
130        for epoch in range(epochs):
131            print(f"\nEpoch {epoch + 1}/{epochs}")
132
133            train_loss = self.train_epoch(train_loader)
134            val_results = self.validate(val_loader)
135
136            if self.scheduler:
137                self.scheduler.step(val_results['mean_iou'])
138
139            print(f"Train Loss: {train_loss:.4f}")
140            print(f"Val Loss: {val_results['loss']:.4f}")
141            print(f"Val mIoU: {val_results['mean_iou']:.4f}")
142            print(f"Val Pixel Acc: {val_results['pixel_accuracy']:.4f}")
143
144            # Save best model
145            if val_results['mean_iou'] > best_miou:
146                best_miou = val_results['mean_iou']
147                torch.save({
148                    'epoch': epoch,
149                    'model_state_dict': self.model.state_dict(),
150                    'optimizer_state_dict': self.optimizer.state_dict(),
151                    'best_miou': best_miou,
152                }, save_path)
153                print(f"Saved best model with mIoU: {best_miou:.4f}")
154
155# Example training setup
156def setup_training():
157    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
158    num_classes = 21  # Pascal VOC
159
160    # Model
161    model = UNet(n_channels=3, n_classes=num_classes)
162
163    # Loss and optimizer
164    criterion = CombinedLoss(ce_weight=0.5, dice_weight=0.5)
165    optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)
166    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
167        optimizer, mode='max', factor=0.5, patience=5
168    )
169
170    # Create trainer
171    trainer = SegmentationTrainer(
172        model, criterion, optimizer, scheduler, device, num_classes
173    )
174
175    return trainer
176
177print("Training pipeline setup complete!")