Projects / 2024

Bird Image Semantic Segmentation

Python · PyTorch · Matplotlib

A bird image segmentation model using U-Net, trained on the CUB-200-2011 dataset.

0.1193 validation BCE loss

11,788 images CUB-200-2011

200 bird species

The task

Semantic segmentation is per-pixel classification: rather than asking “is there a bird,” the model labels every pixel as bird or background, producing a precise mask of the animal. I trained a U-Net to do exactly that on bird photographs.

Dataset

The CUB-200-2011 dataset, also known as the Caltech-UCSD Birds-200-2011, is a popular dataset, often used for bird-related computer vision projects. It consists of 11,788 images of 200 different bird species. This dataset also includes detailed annotations, such as species labels, bounding boxes, and segmentation masks, making it a perfect choice for this project.

Model

The U-Net architecture is a well-regarded model in the field of image semantic segmentation. Its encoder-decoder structure, paired with skip connections, is very effective for generating high precision segmentation masks. These skip connections bridge the encoder and decoder paths, helping retain spatial details while also capturing broader semantic features.

The encoder portion of U-Net uses convolutional layers along with max-pooling to progressively reduce the spatial dimensions of the image while extracting high-level features. The decoder then uses transposed convolutional layers to reconstruct the image, merging the high-level features from the encoder with the spatial information through skip connections.

Results

On in-domain images, North American birds like those in CUB-200-2011, the model segments cleanly.

In-domain segmentation result on a North American bird (1 of 3): original image, predicted mask, and segmented overlay

In-domain segmentation result on a North American bird (2 of 3): original image, predicted mask, and segmented overlay

In-domain segmentation result on a North American bird (3 of 3): original image, predicted mask, and segmented overlay

But, surprisingly, for images not belonging to that domain space (ood data), the model seemed to struggle to find proper segmentation masks. Given below are the results given by the model on images of bird species of the Indian subcontinent.

Out-of-domain segmentation result on a bird from the Indian subcontinent (1 of 3): original image, predicted mask, and segmented overlay

Out-of-domain segmentation result on a bird from the Indian subcontinent (2 of 3): original image, predicted mask, and segmented overlay

Out-of-domain segmentation result on a bird from the Indian subcontinent (3 of 3): original image, predicted mask, and segmented overlay

Conclusion

Since this was a very quick project, I only used a simple U-Net architecture. However, in future, I would like to explore some improvements on the U-Net architecture, particularly by adding Attention Gates (AG). These gates can help the model focus on relevant regions in the image, potentially improving its ability to generate accurate segmentation masks, especially for bird species outside the training domain space.

Next project →

DnCNN Image Gaussian Denoising