Image Segmentation

import torch
import torch.nn as nn
import torch.nn.functional as F

class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True),
        )
    def forward(self, x): return self.net(x)

class Down(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.net = nn.Sequential(nn.MaxPool2d(2), DoubleConv(in_ch, out_ch))
    def forward(self, x): return self.net(x)

class Up(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.up   = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.conv = DoubleConv(in_ch, out_ch)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        diffH = x2.size(2) - x1.size(2)
        diffW = x2.size(3) - x1.size(3)
        x1 = F.pad(x1, [diffW//2, diffW-diffW//2, diffH//2, diffH-diffH//2])
        return self.conv(torch.cat([x2, x1], dim=1))

class UNet(nn.Module):
    def __init__(self, in_channels=3, num_classes=2, base_ch=64):
        super().__init__()
        self.inc   = DoubleConv(in_channels, base_ch)
        self.down1 = Down(base_ch,     base_ch * 2)
        self.down2 = Down(base_ch * 2, base_ch * 4)
        self.down3 = Down(base_ch * 4, base_ch * 8)
        self.down4 = Down(base_ch * 8, base_ch * 8)
        self.up1   = Up(base_ch * 16, base_ch * 4)
        self.up2   = Up(base_ch * 8,  base_ch * 2)
        self.up3   = Up(base_ch * 4,  base_ch)
        self.up4   = Up(base_ch * 2,  base_ch)
        self.outc  = nn.Conv2d(base_ch, num_classes, 1)

    def forward(self, x):
        x1 = self.inc(x)
        x2, x3, x4, x5 = self.down1(x1), self.down2(self.down1(x1)), self.down3(self.down2(self.down1(x1))), self.down4(self.down3(self.down2(self.down1(x1))))
        # Cleaner version:
        x1=self.inc(x); x2=self.down1(x1); x3=self.down2(x2); x4=self.down3(x3); x5=self.down4(x4)
        x=self.up1(x5,x4); x=self.up2(x,x3); x=self.up3(x,x2); x=self.up4(x,x1)
        return self.outc(x)

# Test
model = UNet(in_channels=3, num_classes=21)
out = model(torch.randn(2, 3, 512, 512))
print(out.shape)  # (2, 21, 512, 512)

# Dice + CE combined loss
class SegmentationLoss(nn.Module):
    def __init__(self, dice_w=0.5, ce_w=0.5):
        super().__init__()
        self.ce = nn.CrossEntropyLoss()
        self.dw, self.cw = dice_w, ce_w

    def dice_loss(self, logits, targets):
        probs = F.softmax(logits, dim=1)
        C = probs.shape[1]
        targets_oh = F.one_hot(targets, C).permute(0, 3, 1, 2).float()
        inter = (probs * targets_oh).sum(dim=(2, 3))
        union = probs.sum(dim=(2, 3)) + targets_oh.sum(dim=(2, 3))
        return (1.0 - ((2*inter + 1) / (union + 1)).mean())

    def forward(self, logits, targets):
        return self.cw * self.ce(logits, targets) + self.dw * self.dice_loss(logits, targets)

# pip install segment-anything
from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator
import numpy as np
from PIL import Image

sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth")
sam = sam.to("cuda")

# Interactive prediction (point prompt)
predictor = SamPredictor(sam)
image = np.array(Image.open("photo.jpg").convert("RGB"))
predictor.set_image(image)   # compute embeddings once

masks, scores, logits = predictor.predict(
    point_coords=np.array([[300, 200]]),
    point_labels=np.array([1]),   # 1=foreground
    multimask_output=True,        # 3 masks at different granularities
)
best_mask = masks[scores.argmax()]
print(f"Mask coverage: {best_mask.mean():.2%}")

# Box prompt
masks, scores, _ = predictor.predict(
    box=np.array([[x1, y1, x2, y2]]),
    multimask_output=False,
)

# Automatic mask generation (everything mode)
mask_generator = SamAutomaticMaskGenerator(
    model=sam,
    points_per_side=32,
    pred_iou_thresh=0.88,
    stability_score_thresh=0.95,
    min_mask_region_area=100,
)
auto_masks = mask_generator.generate(image)
print(f"Found {len(auto_masks)} segments")