تجزئة الصور: UNet وDeepLab

import torch
import torch.nn as nn
import torchvision.models as models
from torchvision.models.segmentation import DeepLabV3_ResNet101_Weights

class="tok-comment"># ── class="tok-num">1. Pretrained semantic segmentation (DeepLabV3) ───────────────────────────
model = models.segmentation.deeplabv3_resnet101(
    weights=DeepLabV3_ResNet101_Weights.DEFAULT
)
model.eval()

class="tok-comment"># Inference
from PIL import Image
import torchvision.transforms.functional as F
import numpy as np

img = Image.open(class="tok-str">"street.jpg").convert(class="tok-str">"RGB")
img_t = F.to_tensor(img).unsqueeze(class="tok-num">0)   class="tok-comment"># (class="tok-num">1, class="tok-num">3, H, W)
img_t = F.normalize(img_t, mean=[class="tok-num">0.485,class="tok-num">0.456,class="tok-num">0.406], std=[class="tok-num">0.229,class="tok-num">0.224,class="tok-num">0.225])

with torch.no_grad():
    output = model(img_t)[class="tok-str">"out"]         class="tok-comment"># (class="tok-num">1, class="tok-num">21, H, W) — class="tok-num">21 PASCAL VOC classes
pred_mask = output.argmax(dim=class="tok-num">1)[class="tok-num">0]     class="tok-comment"># (H, W) class labels

class="tok-comment"># ── class="tok-num">2. Minimal UNet ────────────────────────────────────────────────────────────
class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, class="tok-num">3, padding=class="tok-num">1, bias=False),
            nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, class="tok-num">3, padding=class="tok-num">1, bias=False),
            nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True),
        )
    def forward(self, x): return self.conv(x)

class UNet(nn.Module):
    def __init__(self, in_channels=class="tok-num">3, n_classes=class="tok-num">2, base_channels=class="tok-num">64):
        super().__init__()
        bc = base_channels
        class="tok-comment"># Encoder
        self.enc1 = DoubleConv(in_channels, bc)
        self.enc2 = DoubleConv(bc, bc*class="tok-num">2)
        self.enc3 = DoubleConv(bc*class="tok-num">2, bc*class="tok-num">4)
        self.pool = nn.MaxPool2d(class="tok-num">2)
        class="tok-comment"># Bottleneck
        self.bottleneck = DoubleConv(bc*class="tok-num">4, bc*class="tok-num">8)
        class="tok-comment"># Decoder
        self.up3  = nn.ConvTranspose2d(bc*class="tok-num">8, bc*class="tok-num">4, class="tok-num">2, stride=class="tok-num">2)
        self.dec3 = DoubleConv(bc*class="tok-num">8, bc*class="tok-num">4)   class="tok-comment"># bc*class="tok-num">8 because of skip connection
        self.up2  = nn.ConvTranspose2d(bc*class="tok-num">4, bc*class="tok-num">2, class="tok-num">2, stride=class="tok-num">2)
        self.dec2 = DoubleConv(bc*class="tok-num">4, bc*class="tok-num">2)
        self.up1  = nn.ConvTranspose2d(bc*class="tok-num">2, bc, class="tok-num">2, stride=class="tok-num">2)
        self.dec1 = DoubleConv(bc*class="tok-num">2, bc)
        class="tok-comment"># Output
        self.out  = nn.Conv2d(bc, n_classes, class="tok-num">1)

    def forward(self, x):
        e1 = self.enc1(x)
        e2 = self.enc2(self.pool(e1))
        e3 = self.enc3(self.pool(e2))
        b  = self.bottleneck(self.pool(e3))
        d3 = self.dec3(torch.cat([self.up3(b),  e3], dim=class="tok-num">1))
        d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=class="tok-num">1))
        d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=class="tok-num">1))
        return self.out(d1)

unet = UNet(n_classes=class="tok-num">2)
x = torch.randn(class="tok-num">2, class="tok-num">3, class="tok-num">256, class="tok-num">256)
out = unet(x)
print(fclass="tok-str">"UNet output: {out.shape}")   class="tok-comment"># (class="tok-num">2, class="tok-num">2, class="tok-num">256, class="tok-num">256)

class="tok-comment"># ── class="tok-num">3. Dice loss ──────────────────────────────────────────────────────────────
def dice_loss(pred, target, eps=class="tok-num">1e-6):
    class="tok-str">"""pred: (B, C, H, W) softmax, target: (B, H, W) long"""
    pred_soft = torch.softmax(pred, dim=class="tok-num">1)
    target_oh = torch.zeros_like(pred_soft)
    target_oh.scatter_(class="tok-num">1, target.unsqueeze(class="tok-num">1), class="tok-num">1)
    inter = (pred_soft * target_oh).sum(dim=(class="tok-num">2,class="tok-num">3))
    union = pred_soft.sum(dim=(class="tok-num">2,class="tok-num">3)) + target_oh.sum(dim=(class="tok-num">2,class="tok-num">3))
    return class="tok-num">1 - (class="tok-num">2*inter + eps) / (union + eps)

pred   = torch.randn(class="tok-num">2, class="tok-num">2, class="tok-num">64, class="tok-num">64)
target = torch.randint(class="tok-num">0, class="tok-num">2, (class="tok-num">2, class="tok-num">64, class="tok-num">64))
loss   = dice_loss(pred, target).mean()
print(fclass="tok-str">"Dice loss: {loss.item():.4f}")