Object Detection Semantic Segmentation
# Object Detection & Semantic Segmentation
## Introduction & Motivation
Object detection and semantic segmentation extend basic image classification, which assigns a single label to an entire image, into the more structured and practically demanding task of localizing and delineating individual objects or regions within a scene. Object detection asks a model to predict a set of bounding boxes, each paired with a class label and confidence score, identifying where every instance of a relevant object appears in an image. Semantic segmentation asks a model to assign a class label to every individual pixel in an image, producing a dense, per-pixel understanding of scene content rather than a sparse set of boxes; a closely related task, instance segmentation, combines both, producing a pixel-precise mask for each individual object instance. These capabilities underlie a vast range of real-world computer vision applications: autonomous vehicles must detect and segment pedestrians, vehicles, and lane markings in real time; medical imaging systems must segment tumors or anatomical structures from scans; retail and manufacturing systems use detection for inventory tracking, quality control, and automated checkout; and satellite and aerial imagery analysis relies on segmentation for land-use classification and disaster response mapping. The motivation for studying detection and segmentation as distinct problems from classification is that most real-world visual understanding tasks require knowing not just what is present in an image, but precisely where, a distinction with major implications for both model architecture and evaluation methodology.
## Core Concepts & Theory
Modern object detectors fall into two broad architectural families. Two-stage detectors, exemplified by the R-CNN family (R-CNN, Fast R-CNN, Faster R-CNN), first generate a set of candidate object regions (region proposals), then classify and refine each proposal in a second stage; Faster R-CNN's key innovation was the Region Proposal Network (RPN), which generates proposals using a learned neural network rather than a separate, non-learned algorithm like Selective Search, making the entire pipeline trainable end-to-end. One-stage detectors, exemplified by YOLO (You Only Look Once) and SSD (Single Shot Detector), skip the separate proposal-generation stage entirely, directly predicting bounding boxes and class labels densely across a grid overlaid on the image in a single forward pass, trading some accuracy for substantially faster inference, which made YOLO the architecture of choice for real-time detection applications. Both families historically relied on anchor boxes: a predefined set of reference boxes at various scales and aspect ratios placed at each spatial location, against which the network predicts offsets and class probabilities rather than raw box coordinates, simplifying the learning problem but introducing a set of hyperparameters (anchor scales and ratios) that must be tuned to the target dataset. More recent anchor-free detectors (e.g., FCOS, CenterNet) and transformer-based detectors (DETR) eliminate anchor boxes entirely, either predicting object centers and extents directly or, in DETR's case, framing detection as a direct set-prediction problem solved via a Transformer encoder-decoder and bipartite matching, removing the need for hand-designed anchor configurations and non-maximum suppression post-processing altogether.
## Mathematical Formulation
Intersection over Union (IoU) is the fundamental metric for comparing a predicted bounding box to a ground-truth box, defined as the area of overlap divided by the area of union:
$$ ext{IoU}(B_{pred}, B_{gt}) = \frac{ ext{Area}(B_{pred} \cap B_{gt})}{ ext{Area}(B_{pred} \cup B_{gt})} $$
A predicted detection is typically counted as a true positive if its IoU with a ground-truth box exceeds a threshold (commonly 0.5), and as a false positive otherwise. The overall detection loss combines a classification term and a bounding-box regression term across all predicted boxes:
$$ \mathcal{L}_{det} = \mathcal{L}_{cls} + \lambda_{box} \, \mathcal{L}_{box} $$
where the classification loss is typically cross-entropy (or focal loss, discussed below) over object classes plus a background class, and the box regression loss is commonly a smooth L1 loss applied to the four box coordinate offsets (center x, center y, width, height) relative to an anchor or reference point. Non-Maximum Suppression (NMS), applied as post-processing in anchor-based detectors, greedily selects the highest-confidence box, removes all other boxes with IoU above a threshold against it, and repeats, eliminating duplicate detections of the same object. For semantic segmentation, the standard per-pixel loss is cross-entropy summed or averaged over all pixels:
$$ \mathcal{L}_{seg} = -\frac{1}{HW}\sum_{i=1}^{H}\sum_{j=1}^{W} \sum_{c=1}^{C} y_{ijc} \log(\hat{p}_{ijc}) $$
where H and W are image height and width, C is the number of classes, y is the one-hot ground-truth label at each pixel, and p-hat is the predicted class probability distribution at each pixel. The Dice loss, an alternative to cross-entropy commonly used in medical image segmentation where foreground classes are often severely underrepresented relative to background pixels, directly optimizes the overlap-based Dice coefficient between predicted and ground-truth masks rather than per-pixel classification accuracy.
## Advanced Theory & Extensions
Feature Pyramid Networks (FPN) address the challenge that objects appear at vastly different scales within a single image by building a multi-scale feature hierarchy with strong semantics at every level, combining coarse, semantically rich features from deep layers with fine-grained, spatially precise features from shallow layers via a top-down pathway with lateral connections, and have become a near-universal component of modern detection and segmentation architectures regardless of the underlying backbone. Focal loss, introduced alongside the RetinaNet one-stage detector, addresses the severe foreground-background class imbalance inherent to dense one-stage detection (where the vast majority of anchor locations correspond to background) by down-weighting the loss contribution of well-classified, easy examples, allowing the training signal to focus on hard, informative examples. DETR (Detection Transformer) reframes object detection as a direct set-prediction problem: a fixed number of learned object queries attend over image features via a Transformer decoder, and a bipartite matching loss (via the Hungarian algorithm) assigns each ground-truth object to exactly one prediction during training, eliminating both anchor boxes and non-maximum suppression as separate, hand-designed pipeline components. Mask R-CNN extends Faster R-CNN to instance segmentation by adding a parallel branch that predicts a binary segmentation mask for each detected object region, alongside the existing classification and box-regression branches, using RoIAlign (a precise, non-quantized feature extraction operation) to preserve pixel-level spatial alignment that the original RoIPool operation lacked. U-Net, originally developed for biomedical image segmentation, established the now-standard encoder-decoder architecture with skip connections directly linking corresponding encoder and decoder resolution levels, allowing the decoder to recover fine spatial detail lost during downsampling while still benefiting from the encoder's increasingly abstract, semantically rich features at deeper levels.
## Computational Considerations
Multi-scale feature processing, essential for handling objects of widely varying size within a single image, multiplies the computational and memory cost of detection and segmentation models relative to single-scale classification networks, since feature maps must be maintained, combined, and processed at several spatial resolutions simultaneously. Two-stage detectors incur additional computational overhead from the region proposal and per-region classification stages, generally making them slower than comparable one-stage detectors, though modern efficient implementations have substantially narrowed this historical accuracy-speed gap. Non-maximum suppression, while computationally cheap for a small number of candidate boxes, can become a bottleneck at very high object density (e.g., crowded scenes with many overlapping objects) and is inherently non-differentiable and sequential, motivating the set-prediction reformulation used by DETR-style architectures that eliminates this post-processing step entirely. Semantic segmentation models must produce output at or near the original input resolution, requiring either computationally expensive upsampling/decoder pathways (as in U-Net and Feature Pyramid Network-based segmentation heads) or dilated/atrous convolutions (as used in the DeepLab family) that expand the receptive field without reducing spatial resolution, both of which carry non-trivial memory costs for high-resolution inputs. Real-time deployment requirements (autonomous driving, robotics, video analytics) place hard latency and often on-device compute constraints on detection and segmentation models, motivating a substantial line of research into efficient architectures (MobileNet-based backbones, quantization, and pruning, connecting directly to the model compression literature) specifically tuned for these dense prediction tasks rather than only for classification.
## Practical Implementation Strategies
Transfer learning from a backbone pretrained on a large classification dataset (ImageNet, or increasingly self-supervised or vision-language pretrained backbones) is standard practice for both detection and segmentation, since detection and segmentation datasets are typically far smaller than classification datasets, and pretrained low- and mid-level visual features transfer effectively to these denser prediction tasks. Data augmentation strategies specific to detection and segmentation must correctly transform both the image and its associated annotations (bounding boxes must be recomputed under geometric transforms, segmentation masks must be spatially warped identically to the image), and libraries like Albumentations provide annotation-aware augmentation pipelines that handle this coupling correctly. Anchor box configuration (scales and aspect ratios) should be tuned to match the statistical distribution of object sizes and shapes in the target dataset, commonly via k-means clustering of ground-truth box dimensions, since poorly matched anchors can substantially degrade detection recall regardless of downstream architecture quality; anchor-free and DETR-style detectors sidestep this tuning burden entirely. Class imbalance, both in the background-versus-foreground sense (addressed by focal loss) and in the sense of rare object categories being underrepresented in training data, requires explicit handling through loss reweighting, resampling strategies, or synthetic data generation, since naive training tends to systematically underperform on rare classes. Evaluation-driven iteration should examine per-class and per-scale performance breakdowns (small versus medium versus large objects, as standardized in the COCO evaluation protocol) rather than relying solely on an aggregate metric, since aggregate scores can mask severe underperformance on specific object categories or size ranges that matter for the target application.
## Benchmark Datasets & Evaluation
The COCO (Common Objects in Context) dataset, with over 200,000 labeled images spanning 80 object categories and rich annotations including bounding boxes, instance segmentation masks, and keypoints, is the dominant benchmark for general-purpose object detection and instance segmentation, and its evaluation protocol (mean Average Precision averaged across IoU thresholds from 0.5 to 0.95) has become the standard reporting metric across the field. Pascal VOC, an earlier and smaller benchmark with 20 object categories, remains in use for historical comparison and rapid prototyping, using a simpler mean Average Precision computed at a single IoU threshold of 0.5. Cityscapes provides densely annotated urban street scenes specifically designed for autonomous driving research, with fine-grained semantic and instance segmentation labels for road-relevant categories (vehicles, pedestrians, road surface, traffic signs). ADE20K offers broad-coverage scene parsing annotations spanning indoor and outdoor scenes across a very large number of semantic categories, testing generalization across a much wider vocabulary of object and stuff (non-object background regions like sky, grass, and road) categories than COCO. Mean Average Precision (mAP) for detection integrates precision-recall behavior across confidence thresholds, computed per class and then averaged, while mean Intersection over Union (mIoU), averaged across classes, is the standard aggregate metric for semantic segmentation, measuring the overlap between predicted and ground-truth pixel masks for each class independent of instance-level distinctions.
## Key Challenges & Limitations
Small object detection remains substantially harder than detecting large, prominent objects, since small objects are represented by very few pixels (and correspondingly few activations in deep, heavily downsampled feature maps), a challenge that motivated feature pyramid architectures but which persists as a source of systematic performance gaps, particularly visible in the small-object breakdown of COCO evaluation results. Occlusion and crowded scenes, where multiple objects of the same class heavily overlap, challenge both the detection stage (non-maximum suppression can incorrectly discard a correct detection of a heavily occluded object because it overlaps substantially with a neighboring correct detection) and the annotation process itself, since even human annotators struggle to precisely delineate heavily occluded object boundaries. Domain shift between training and deployment conditions (different camera characteristics, lighting, weather, geographic region) can cause substantial performance degradation, a particularly acute concern for safety-critical applications like autonomous driving, motivating research into domain adaptation and robustness-focused training and evaluation protocols. Annotation cost and quality present a persistent practical bottleneck: pixel-precise segmentation masks are dramatically more expensive to produce than image-level classification labels or even bounding boxes, motivating weakly supervised and semi-supervised segmentation approaches that learn from cheaper, coarser annotation signals (e.g., image-level labels or bounding boxes alone). Long-tailed class distributions, where a small number of common categories dominate the training data while many categories have only a handful of examples, systematically bias models toward strong performance on common classes at the expense of rare but potentially important categories, requiring specialized loss functions and resampling strategies beyond standard training recipes.
## Hyperparameter Tuning
The IoU threshold used to define positive versus negative training examples (commonly 0.5 for standard detection, though some architectures use different thresholds for different network stages, as in Cascade R-CNN's progressively stricter multi-stage thresholds) directly affects the localization precision-recall trade-off the trained model learns to optimize for. The non-maximum suppression IoU threshold controls how aggressively overlapping detections are pruned, with a lower threshold removing more duplicate detections at the risk of incorrectly suppressing genuinely separate, overlapping objects, particularly relevant in crowded-scene applications. Focal loss's focusing parameter (commonly denoted gamma, typically set around 2) controls how strongly the loss down-weights easy, well-classified examples relative to hard examples, with higher values placing progressively more training emphasis on the hardest examples at the risk of instability if set too aggressively. Anchor box scale and aspect ratio configurations must be matched to the expected object size and shape distribution of the target dataset and deployment domain, with poor matches degrading recall regardless of how well-tuned other hyperparameters are. Backbone architecture and depth selection trades off accuracy against inference latency and memory footprint, a decision that should be driven by the deployment constraints of the target application (real-time embedded deployment favoring lightweight backbones like MobileNet, versus offline batch processing affording much larger, more accurate backbones like ResNet-101 or Swin Transformer variants).
## Real-World Applications & Case Studies
Autonomous vehicles rely on real-time object detection and semantic segmentation as core perception components, jointly identifying and localizing pedestrians, vehicles, cyclists, traffic signs, and drivable road surface from camera (and often LiDAR-fused) sensor data, with extremely stringent accuracy and latency requirements given the safety-critical nature of the application. Medical imaging applications use segmentation extensively, including tumor and lesion segmentation in MRI and CT scans, organ segmentation for radiotherapy planning, and cell and nucleus segmentation in microscopy images, frequently using U-Net or its many derivatives as the standard architectural starting point given the relatively small, specialized datasets typical of medical imaging. Retail and manufacturing applications use object detection for automated checkout systems (identifying products as they move through a scanning area), inventory and shelf-monitoring systems, and visual quality control on production lines, detecting defects or misassembled components in real time. Satellite and aerial imagery analysis applies detection and segmentation to land-use classification, crop health monitoring in precision agriculture, building footprint extraction for urban planning, and disaster response mapping (rapidly identifying damaged structures after natural disasters from aerial imagery). Video surveillance and security systems combine detection with tracking algorithms to monitor object and person movement across camera feeds over time, raising both valuable safety applications and significant privacy considerations that have prompted ongoing regulatory and ethical discussion.
## Integration with Other Methods
Detection and segmentation architectures share substantial architectural DNA with the broader convolutional and vision transformer literature, with backbone networks (ResNet, Swin Transformer, ConvNeXt) typically pretrained via standard image classification or self-supervised objectives before being adapted for dense prediction tasks via task-specific heads. Vision-language models increasingly incorporate detection and segmentation capabilities through open-vocabulary detection and segmentation, where a model can localize and segment arbitrary object categories described in natural language at inference time (e.g., Grounding DINO, Segment Anything Model) rather than being restricted to a fixed, closed set of training categories, directly connecting detection research to the multimodal vision-language literature. Video object detection and segmentation extend the single-image problem to the temporal domain, incorporating motion information and temporal consistency constraints, and increasingly leverage the same self-attention and Transformer mechanisms used in sequence modeling to reason across frames. Detection and segmentation models are frequently combined with tracking algorithms (e.g., DeepSORT, ByteTrack) in downstream multi-object tracking pipelines, associating per-frame detections into consistent object trajectories over time for applications like video surveillance and autonomous driving perception stacks. Model compression techniques (quantization, pruning, knowledge distillation) are heavily applied to detection and segmentation models destined for edge or real-time deployment, since the dense, multi-scale computation inherent to these architectures is often more expensive than comparable classification models, making compression a practically essential rather than merely optional step for many deployment scenarios.
## Future Research Directions
Open-vocabulary and zero-shot detection and segmentation, enabling models to localize and segment object categories not seen during training by leveraging natural-language category descriptions, represents one of the most active current research directions, promising to substantially reduce the closed-vocabulary limitation of traditional detection and segmentation systems. Improving small-object detection and handling of extreme scale variation within a single image remains an open challenge despite substantial progress from feature pyramid architectures, with continued research into higher-resolution processing strategies that remain computationally tractable. Reducing annotation cost through weakly supervised, semi-supervised, and self-supervised approaches to detection and segmentation is an important practical direction, given the substantial expense of producing pixel-precise or even box-level annotations at the scale required for training high-performing models. 3D object detection and segmentation, extending these tasks to point clouds and volumetric data (critical for autonomous driving LiDAR perception and medical volumetric imaging), presents distinct architectural challenges from the 2D image case and remains an active area of specialized research. Finally, improving robustness to distribution shift, adversarial perturbations, and out-of-distribution inputs is increasingly recognized as essential for safety-critical deployment of detection and segmentation systems, motivating research into uncertainty quantification and robust training methods specifically tailored to dense prediction tasks rather than borrowed directly from the classification robustness literature.
## Summary & Key Takeaways
Object detection and semantic segmentation extend image classification into spatially structured prediction tasks, with detection localizing objects via bounding boxes and segmentation providing dense, per-pixel class labels, and instance segmentation combining both. Two-stage detectors (Faster R-CNN) trade speed for accuracy relative to one-stage detectors (YOLO, SSD), while anchor-free and Transformer-based approaches (DETR) increasingly eliminate the hand-designed anchor boxes and non-maximum suppression that characterized earlier architectures. Feature Pyramid Networks address multi-scale object handling, focal loss addresses foreground-background class imbalance in dense detection, and U-Net's encoder-decoder-with-skip-connections design remains the dominant architectural pattern for segmentation, particularly in medical imaging. Persistent challenges include small object detection, occlusion in crowded scenes, domain shift robustness, and the substantial cost of producing pixel-precise annotations, while open-vocabulary detection and segmentation, driven by vision-language model integration, represents one of the most active current research frontiers extending these capabilities beyond fixed, closed category vocabularies.
Keywords: object detection, semantic segmentation, instance segmentation, Faster R-CNN, YOLO, SSD, anchor boxes, Feature Pyramid Network, focal loss, non-maximum suppression, DETR, Mask R-CNN, U-Net, intersection over union, mean average precision, mIoU, RoIAlign, open-vocabulary detection, Segment Anything Model, panoptic segmentation
---
## Appendix: Practical Labs
### Lab 1: Intersection over Union and Non-Maximum Suppression
import numpy as np
def compute_iou(box_a, box_b):
"""box format: [x_min, y_min, x_max, y_max]"""
x_min = max(box_a[0], box_b[0])
y_min = max(box_a[1], box_b[1])
x_max = min(box_a[2], box_b[2])
y_max = min(box_a[3], box_b[3])
inter_area = max(0, x_max - x_min) * max(0, y_max - y_min)
area_a = (box_a[2] - box_a[0]) * (box_a[3] - box_a[1])
area_b = (box_b[2] - box_b[0]) * (box_b[3] - box_b[1])
union_area = area_a + area_b - inter_area
return inter_area / union_area if union_area > 0 else 0.0
def non_max_suppression(boxes, scores, iou_threshold=0.5):
"""Greedy NMS: repeatedly select the highest-scoring remaining box and
suppress all other boxes overlapping it above the IoU threshold."""
order = np.argsort(scores)[::-1]
keep = []
while len(order) > 0:
current = order[0]
keep.append(current)
remaining = order[1:]
suppressed = []
for idx in remaining:
iou = compute_iou(boxes[current], boxes[idx])
if iou <= iou_threshold:
suppressed.append(idx)
order = np.array(suppressed)
return keep
def test_iou_and_nms():
boxes = np.array([
[10, 10, 50, 50], # box 0
[12, 12, 52, 52], # box 1, heavily overlaps box 0
[100, 100, 140, 140], # box 2, separate object
[15, 15, 55, 55], # box 3, overlaps box 0 and 1
])
scores = np.array([0.9, 0.75, 0.95, 0.6])
iou_01 = compute_iou(boxes[0], boxes[1])
iou_02 = compute_iou(boxes[0], boxes[2])
print(f"IoU(box0, box1) = {iou_01:.3f} (should be high, overlapping)")
print(f"IoU(box0, box2) = {iou_02:.3f} (should be 0, non-overlapping)")
assert iou_01 > 0.5, "Heavily overlapping boxes should have high IoU"
assert iou_02 == 0.0, "Non-overlapping boxes should have zero IoU"
kept = non_max_suppression(boxes, scores, iou_threshold=0.5)
kept_boxes = [i for i in kept]
print(f"Boxes kept after NMS: {kept_boxes}")
# box 2 (highest score, separate object) and box 0's cluster's top scorer should survive
assert 2 in kept, "The high-confidence, non-overlapping box 2 must survive NMS"
assert len(kept) < len(boxes), "NMS should suppress at least one duplicate detection"
print("IoU and NMS test passed.")
if __name__ == "__main__":
test_iou_and_nms()### Lab 2: Focal Loss for Dense Detection Class Imbalance
import torch
import torch.nn.functional as F
def focal_loss(logits, targets, alpha=0.25, gamma=2.0):
"""Binary focal loss, as used in RetinaNet, down-weighting easy examples
(high predicted probability for the correct class) relative to hard
examples, addressing the severe foreground/background imbalance in
dense one-stage detection."""
probs = torch.sigmoid(logits)
ce_loss = F.binary_cross_entropy_with_logits(logits, targets, reduction="none")
p_t = probs * targets + (1 - probs) * (1 - targets)
focal_weight = (1 - p_t) ** gamma
alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
loss = alpha_t * focal_weight * ce_loss
return loss.mean()
def test_focal_loss_downweights_easy_examples():
torch.manual_seed(0)
# "Easy" examples: model is already confident and correct.
easy_logits = torch.tensor([5.0, 5.0, 5.0, -5.0, -5.0])
easy_targets = torch.tensor([1.0, 1.0, 1.0, 0.0, 0.0])
# "Hard" examples: model is uncertain or wrong.
hard_logits = torch.tensor([0.1, -0.2, 0.3, 0.0, -0.1])
hard_targets = torch.tensor([1.0, 1.0, 1.0, 0.0, 0.0])
easy_ce = F.binary_cross_entropy_with_logits(easy_logits, easy_targets).item()
hard_ce = F.binary_cross_entropy_with_logits(hard_logits, hard_targets).item()
easy_focal = focal_loss(easy_logits, easy_targets).item()
hard_focal = focal_loss(hard_logits, hard_targets).item()
print(f"Easy examples: CE={easy_ce:.4f}, Focal={easy_focal:.4f}")
print(f"Hard examples: CE={hard_ce:.4f}, Focal={hard_focal:.4f}")
# Focal loss should shrink the easy-example loss much more aggressively
# than the hard-example loss, relative to standard cross-entropy.
easy_ratio = easy_focal / (easy_ce + 1e-8)
hard_ratio = hard_focal / (hard_ce + 1e-8)
print(f"Focal/CE ratio -- easy: {easy_ratio:.4f}, hard: {hard_ratio:.4f}")
assert easy_ratio < hard_ratio, "Focal loss should down-weight easy examples more than hard examples"
print("Focal loss test passed.")
if __name__ == "__main__":
test_focal_loss_downweights_easy_examples()### Lab 3: U-Net-Style Encoder-Decoder with Skip Connections
import torch
import torch.nn as nn
class ConvBlock(nn.Module):
def __init__(self, in_channels, out_channels):
super().__init__()
self.block = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
)
def forward(self, x):
return self.block(x)
class SimpleUNet(nn.Module):
"""A small U-Net for binary segmentation, with an encoder that
downsamples via max pooling, a decoder that upsamples via transposed
convolution, and skip connections linking corresponding resolutions."""
def __init__(self, in_channels=3, n_classes=1, base_channels=16):
super().__init__()
c = base_channels
self.enc1 = ConvBlock(in_channels, c)
self.enc2 = ConvBlock(c, c * 2)
self.pool = nn.MaxPool2d(2)
self.bottleneck = ConvBlock(c * 2, c * 4)
self.up2 = nn.ConvTranspose2d(c * 4, c * 2, kernel_size=2, stride=2)
self.dec2 = ConvBlock(c * 4, c * 2) # doubled due to skip concat
self.up1 = nn.ConvTranspose2d(c * 2, c, kernel_size=2, stride=2)
self.dec1 = ConvBlock(c * 2, c) # doubled due to skip concat
self.out_conv = nn.Conv2d(c, n_classes, kernel_size=1)
def forward(self, x):
e1 = self.enc1(x) # skip connection 1
e2 = self.enc2(self.pool(e1)) # skip connection 2
b = self.bottleneck(self.pool(e2))
d2 = self.up2(b)
d2 = self.dec2(torch.cat([d2, e2], dim=1)) # skip connection used here
d1 = self.up1(d2)
d1 = self.dec1(torch.cat([d1, e1], dim=1)) # skip connection used here
return self.out_conv(d1)
def test_unet_forward_and_skip_connections():
torch.manual_seed(0)
model = SimpleUNet(in_channels=3, n_classes=1, base_channels=8)
x = torch.randn(2, 3, 64, 64)
output = model(x)
print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")
assert output.shape == (2, 1, 64, 64), "U-Net output should match input spatial resolution"
# Verify gradients flow through the skip connections by checking the
# earliest encoder layer receives a nonzero gradient.
loss = output.sum()
loss.backward()
enc1_grad_norm = sum(p.grad.norm().item() for p in model.enc1.parameters() if p.grad is not None)
assert enc1_grad_norm > 0, "Gradients should flow back to the first encoder block via skip connections"
print(f"Encoder block 1 gradient norm: {enc1_grad_norm:.4f}")
print("U-Net skip connection test passed.")
if __name__ == "__main__":
test_unet_forward_and_skip_connections()### Lab 4: Dice Loss and Mean IoU for Segmentation Evaluation
import torch
import numpy as np
def dice_loss(pred_probs, target_mask, epsilon=1e-6):
"""Soft Dice loss: 1 - Dice coefficient, computed on predicted
probabilities (not thresholded), suitable as a differentiable training
objective for segmentation, especially under class imbalance."""
pred_flat = pred_probs.flatten()
target_flat = target_mask.flatten()
intersection = (pred_flat * target_flat).sum()
dice_coeff = (2.0 * intersection + epsilon) / (pred_flat.sum() + target_flat.sum() + epsilon)
return 1.0 - dice_coeff
def compute_mean_iou(pred_mask, target_mask, n_classes):
"""Computes per-class IoU and averages across classes present in
either the prediction or the ground truth (mean IoU)."""
ious = []
for c in range(n_classes):
pred_c = (pred_mask == c)
target_c = (target_mask == c)
intersection = np.logical_and(pred_c, target_c).sum()
union = np.logical_or(pred_c, target_c).sum()
if union == 0:
continue # class not present in either prediction or ground truth
ious.append(intersection / union)
return np.mean(ious) if ious else 0.0
def test_dice_loss_and_miou():
torch.manual_seed(0)
# Case 1: near-perfect prediction.
target = torch.zeros(1, 1, 20, 20)
target[:, :, 5:15, 5:15] = 1.0
good_pred = target.clone() * 0.95 + 0.025 # confident, mostly correct
# Case 2: poor prediction (mostly misses the foreground region).
bad_pred = torch.full_like(target, 0.05)
good_loss = dice_loss(good_pred, target).item()
bad_loss = dice_loss(bad_pred, target).item()
print(f"Dice loss (good prediction): {good_loss:.4f}")
print(f"Dice loss (poor prediction): {bad_loss:.4f}")
assert good_loss < bad_loss, "A more accurate prediction should have lower Dice loss"
# Mean IoU sanity check with discrete class masks.
pred_mask = np.zeros((20, 20), dtype=int)
pred_mask[5:15, 5:15] = 1
target_mask = np.zeros((20, 20), dtype=int)
target_mask[6:16, 6:16] = 1 # slightly shifted ground truth
miou = compute_mean_iou(pred_mask, target_mask, n_classes=2)
print(f"Mean IoU (slightly shifted prediction): {miou:.4f}")
assert 0.0 < miou < 1.0, "Mean IoU should be a valid value between 0 and 1 for imperfect overlap"
print("Dice loss and mean IoU test passed.")
if __name__ == "__main__":
test_dice_loss_and_miou()