Semantic segmentation of pets using a U-Net implemented from scratch in PyTorch.
Trained on the Oxford-IIIT Pet Dataset, achieving Validation Dice 0.80 on CPU.
This is my first deep learning segmentation project. I built a U-Net from scratch in PyTorch and trained it to segment cats and dogs from background pixels in the Oxford-IIIT Pet Dataset.
The project includes the full ML pipeline - data loading, augmentation, training, validation, prediction, and several model interpretability visualizations (feature maps, Grad-CAM, U-shape architecture diagram).
| Metric | Train | Validation |
|---|---|---|
| Dice Score | 0.8342 | 0.7993 |
| IoU Score | 0.7175 | 0.6681 |
| Loss | 0.3741 | 0.4399 (best) |
Trained for 50 epochs on CPU (~12 hours).
- U-Net from scratch: 3-level encoder-decoder with skip connections
- Combined loss: Dice Loss + 0.5 × BCE Loss for stable training under class imbalance
- Synced augmentation: Image and mask transformed together
- Metrics: Dice Score and IoU computed per epoch
- TensorBoard logging: Real-time training monitoring
- Model interpretability: Feature maps, Grad-CAM, U-shape architecture diagram
- Best model checkpointing: Saves the model with lowest validation loss
Input → Ground Truth → Model Prediction:
Raw model output - bright areas = high cat-probability:
8 of the 32 filters in enc1 - each detects different features (edges, textures, contrast):
Watching features get more abstract from shallow (top) to deep (bottom). Notice how the bottleneck compresses the cat into a small bright blob:
The classic "U-shape" - encoder shrinks (left), bottleneck at the bottom, decoder grows (right):
Shows where the model "looks" when making predictions - concentrated on the cat's body:
Validation Dice climbing steadily to ~0.80 over 50 epochs:
A real U-Net with 3 encoder levels, bottleneck, and 3 decoder levels with skip connections.
Encoder:
enc1: ConvBlock(3 → 32) → 128×128
enc2: ConvBlock(32 → 64) → 64×64
enc3: ConvBlock(64 → 128) → 32×32
Bottleneck:
ConvBlock(128 → 256) → 16×16
Decoder (with skip connections):
upconv3 + dec3(256 → 128) → 32×32
upconv2 + dec2(128 → 64) → 64×64
upconv1 + dec1(64 → 32) → 128×128
Output:
Conv2d(32 → 1, 1×1) → segmentation mask
Each ConvBlock is Conv → BatchNorm → ReLU repeated twice.
Oxford-IIIT Pet Dataset - 7,390 images of 37 cat and dog breeds with pixel-level segmentation masks.
The original masks have three classes:
- 1 = pet border/outline (~9.6% of pixels)
- 2 = background (~82.8% of pixels)
- 3 = pet body (~7.6% of pixels)
For binary segmentation, I treat border + body as foreground:
mask = ((mask == 1) | (mask == 3)).astype("float32")This gives ~17% foreground coverage - much easier to learn than body-only (7.6%) and produces clean filled cat silhouettes.
git clone https://github.com/saeed-moo/pet-image-segmentation.git
cd pet-image-segmentation
python -m venv .venv
# Windows
.venv\Scripts\activate
# macOS/Linux
source .venv/bin/activate
pip install -r requirements.txtpython download_dataset.pypython check_masks.pyThis prints unique mask values and shows pet/background coverage - useful for verifying mask binarization.
python train.pyTrains for 50 epochs, saves the best model as model_best.pth, and logs metrics to TensorBoard.
Monitor in another terminal:
tensorboard --logdir=runspython compare_predictions.py # Input + GT + Prediction (3 panels)
python predict.py # Input + Prediction (2 panels)python visualise_features.py # Feature maps from enc1
python gradcam_visualise.py # Grad-CAM attention heatmapThis project taught me a lot about real-world ML debugging. Three big issues I had to track down:
Initially I used mask <= 2 thinking it kept the pet. But values 1 and 2 in this dataset are border + background, not pet. This inversion gave a misleading fake Dice 0.937 because the model just predicted "everything is foreground" and got 92% pixel accuracy.
I diagnosed it by inspecting unique mask values and visualizing what each binarization option actually selected:
Value 1: 9.56% (border)
Value 2: 82.83% (background — biggest!)
Value 3: 7.61% (pet body)
The visualization made it obvious - mask <= 2 was selecting border + background (everything except the cat body), not the pet. This is why class imbalance was so extreme and the model collapsed to predicting "all foreground."
My RandomHorizontalFlip only flipped the image, not the mask. So the model was being trained on image-mask pairs where the cat was in one position and the label said it was in another - essentially noise. Fixed by using torchvision.transforms.functional to apply the same transform to both.
With pet body at only 7.6% of pixels, plain BCE loss converged to "predict all background." Solved by combining Dice Loss (handles imbalance natively) with weighted BCE Loss for boundary detail.
Final result: Real Val Dice 0.80 (up from a broken fake 0.937 → real but mediocre 0.42 → genuinely good 0.80).
pet-image-segmentation/
├── train.py # Training script
├── predict.py # Single-image prediction
├── compare_predictions.py # 3-panel comparison
├── visualise_features.py # Feature map visualization
├── gradcam_visualise.py # Grad-CAM heatmap
├── check_masks.py # Mask inspection utility
├── dataset.py # PyTorch Dataset
├── metrics.py # Dice & IoU + DiceLoss class
├── download_dataset.py # Dataset downloader
├── requirements.txt # Python dependencies
├── models/
│ └── unet.py # U-Net model
├── data/ # Images and masks (not in repo)
├── results/ # Visualizations (in repo)
└── runs/ # TensorBoard logs (not in repo)
PyTorch • Torchvision • NumPy • PIL • OpenCV • Matplotlib • TensorBoard
- U-Net paper (Ronneberger et al., 2015)
- Grad-CAM paper (Selvaraju et al., 2016)
- Oxford-IIIT Pet Dataset
MIT License
⭐ Star this repo if you found it useful!