xbe.at

Computer vision

Image Segmentation using K-means Clustering in Python

7 min read, 15 sections

Contents
  1. Introduction to Image Segmentation with K-means Clustering
  2. K-means Clustering Algorithm Overview
  3. Preprocessing the Image
  4. Applying K-means to Image Segmentation
  5. Visualizing Cluster Centers
  6. Effect of Varying the Number of Clusters
  7. Handling Different Color Spaces
  8. Real-Life Example: Satellite Image Segmentation
  9. Real-Life Example: Medical Image Segmentation
  10. Challenges and Limitations
  11. Improving K-means Segmentation
  12. Post-processing Segmentation Results
  13. Evaluating Segmentation Quality
  14. Conclusion and Future Directions
  15. Additional Resources

Introduction to Image Segmentation with K-means Clustering

Image segmentation is a crucial task in computer vision that involves partitioning an image into multiple segments or regions. K-means clustering is a popular unsupervised learning algorithm that can be applied to image segmentation. This technique groups pixels with similar characteristics into clusters, effectively separating different objects or regions in an image.

import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from skimage import io

# Load and display an example image
image = io.imread('example_image.jpg')
plt.imshow(image)
plt.title('Original Image')
plt.show()

K-means Clustering Algorithm Overview

K-means clustering aims to partition n observations into k clusters, where each observation belongs to the cluster with the nearest mean. In the context of image segmentation, pixels are treated as observations, and their color values (typically in RGB space) serve as features for clustering.

def kmeans_clustering(image, n_clusters):
    # Reshape the image to a 2D array of pixels
    pixels = image.reshape((-1, 3))
    
    # Perform K-means clustering
    kmeans = KMeans(n_clusters=n_clusters, random_state=42)
    kmeans.fit(pixels)
    
    # Get the labels and cluster centers
    labels = kmeans.labels_
    centers = kmeans.cluster_centers_
    
    return labels, centers

Preprocessing the Image

Before applying K-means clustering, we need to preprocess the image. This involves reading the image, converting it to a suitable format, and normalizing the pixel values if necessary.

def preprocess_image(image_path):
    # Read the image
    image = io.imread(image_path)
    
    # Convert to float32 and normalize
    image = image.astype(np.float32) / 255.0
    
    return image

# Example usage
image = preprocess_image('example_image.jpg')
plt.imshow(image)
plt.title('Preprocessed Image')
plt.show()

Applying K-means to Image Segmentation

We'll now apply the K-means algorithm to segment our preprocessed image. The number of clusters (k) determines the number of segments in the output image.

# Apply K-means clustering
n_clusters = 5
labels, centers = kmeans_clustering(image, n_clusters)

# Reshape labels to the image shape
segmented = centers[labels].reshape(image.shape)

# Display the segmented image
plt.imshow(segmented)
plt.title(f'Segmented Image (k={n_clusters})')
plt.show()

Visualizing Cluster Centers

The cluster centers represent the average color of each segment. Visualizing these centers can provide insights into the dominant colors in the segmented image.

def plot_color_centers(centers):
    # Create a bar plot of cluster centers
    plt.figure(figsize=(10, 2))
    plt.imshow([centers], aspect='auto')
    plt.title('Cluster Color Centers')
    plt.xticks([])
    plt.yticks([])
    plt.show()

# Visualize cluster centers
plot_color_centers(centers)

Effect of Varying the Number of Clusters

The choice of k (number of clusters) significantly impacts the segmentation results. Let's explore how different values of k affect the output.

def segment_and_plot(image, k):
    labels, centers = kmeans_clustering(image, k)
    segmented = centers[labels].reshape(image.shape)
    plt.subplot(1, 2, 1)
    plt.imshow(image)
    plt.title('Original Image')
    plt.subplot(1, 2, 2)
    plt.imshow(segmented)
    plt.title(f'Segmented (k={k})')
    plt.show()

# Experiment with different k values
for k in [3, 5, 7, 10]:
    segment_and_plot(image, k)

Handling Different Color Spaces

K-means clustering can be applied to various color spaces. While RGB is common, other spaces like LAB or HSV might yield better results for certain images.

from skimage import color

def segment_in_color_space(image, color_space, n_clusters):
    if color_space == 'RGB':
        transformed = image
    elif color_space == 'LAB':
        transformed = color.rgb2lab(image)
    elif color_space == 'HSV':
        transformed = color.rgb2hsv(image)
    
    labels, _ = kmeans_clustering(transformed, n_clusters)
    segmented = labels.reshape(image.shape[:2])
    
    plt.imshow(segmented, cmap='viridis')
    plt.title(f'Segmented in {color_space} space')
    plt.show()

# Segment in different color spaces
for space in ['RGB', 'LAB', 'HSV']:
    segment_in_color_space(image, space, n_clusters=5)

Real-Life Example: Satellite Image Segmentation

K-means clustering can be applied to satellite imagery for land cover classification. This technique helps identify different types of terrain, such as water bodies, forests, urban areas, and agricultural land.

# Load a satellite image
satellite_image = io.imread('satellite_image.jpg')

# Segment the image
n_clusters = 6
labels, centers = kmeans_clustering(satellite_image, n_clusters)
segmented = centers[labels].reshape(satellite_image.shape)

# Display results
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.imshow(satellite_image)
plt.title('Original Satellite Image')
plt.subplot(1, 2, 2)
plt.imshow(segmented)
plt.title('Segmented Satellite Image')
plt.show()

Real-Life Example: Medical Image Segmentation

K-means clustering is also useful in medical image analysis, such as segmenting MRI brain scans to identify different tissues or potential abnormalities.

# Load an MRI brain scan
mri_image = io.imread('brain_mri.jpg', as_gray=True)

# Segment the image
n_clusters = 4
labels, centers = kmeans_clustering(mri_image.reshape(-1, 1), n_clusters)
segmented = centers[labels].reshape(mri_image.shape)

# Display results
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.imshow(mri_image, cmap='gray')
plt.title('Original MRI Scan')
plt.subplot(1, 2, 2)
plt.imshow(segmented, cmap='viridis')
plt.title('Segmented MRI Scan')
plt.show()

Challenges and Limitations

While K-means clustering is powerful, it has some limitations for image segmentation:

  1. Sensitivity to initialization: Results can vary based on initial centroid positions.
  2. Need to specify k: The optimal number of clusters isn't always known beforehand.
  3. Assumption of spherical clusters: K-means assumes clusters are spherical and equally sized, which may not always hold true for image data.
def demonstrate_initialization_sensitivity(image, k, n_runs=5):
    plt.figure(figsize=(15, 3))
    for i in range(n_runs):
        labels, _ = kmeans_clustering(image, k)
        segmented = labels.reshape(image.shape[:2])
        plt.subplot(1, n_runs, i+1)
        plt.imshow(segmented, cmap='viridis')
        plt.title(f'Run {i+1}')
    plt.show()

# Demonstrate sensitivity to initialization
demonstrate_initialization_sensitivity(image, k=5)

Improving K-means Segmentation

Several techniques can enhance K-means segmentation results:

  1. Multiple initializations: Run K-means multiple times and select the best result.
  2. Elbow method: Determine the optimal k by plotting the within-cluster sum of squares against k.
  3. Spatial information: Incorporate pixel coordinates as features to consider spatial relationships.
def kmeans_with_spatial_info(image, n_clusters):
    h, w = image.shape[:2]
    x, y = np.meshgrid(np.arange(w), np.arange(h))
    spatial_features = np.dstack((image, x/w, y/h))
    
    labels, _ = kmeans_clustering(spatial_features.reshape((-1, 5)), n_clusters)
    return labels.reshape((h, w))

# Segment with spatial information
spatial_segmentation = kmeans_with_spatial_info(image, n_clusters=5)
plt.imshow(spatial_segmentation, cmap='viridis')
plt.title('Segmentation with Spatial Information')
plt.show()

Post-processing Segmentation Results

After applying K-means, post-processing can refine the segmentation:

  1. Morphological operations: Remove small regions or fill holes.
  2. Connected component analysis: Identify and label distinct regions.
  3. Boundary smoothing: Refine segment boundaries for a more natural appearance.
from scipy import ndimage

def post_process_segmentation(segmentation):
    # Apply morphological opening to remove small regions
    opened = ndimage.binary_opening(segmentation)
    
    # Fill holes in the segments
    filled = ndimage.binary_fill_holes(opened)
    
    # Label connected components
    labeled, _ = ndimage.label(filled)
    
    return labeled

# Post-process the segmentation
processed = post_process_segmentation(spatial_segmentation)
plt.imshow(processed, cmap='viridis')
plt.title('Post-processed Segmentation')
plt.show()

Evaluating Segmentation Quality

Assessing the quality of image segmentation is crucial. While ground truth is ideal, unsupervised metrics can be used when it's unavailable:

  1. Inertia: Sum of squared distances of samples to their closest cluster center.
  2. Silhouette score: Measure of how similar an object is to its own cluster compared to other clusters.
  3. Calinski-Harabasz index: Ratio of between-cluster dispersion to within-cluster dispersion.
from sklearn.metrics import silhouette_score, calinski_harabasz_score

def evaluate_segmentation(image, labels):
    pixels = image.reshape((-1, 3))
    inertia = KMeans(n_clusters=len(np.unique(labels))).fit(pixels).inertia_
    silhouette = silhouette_score(pixels, labels)
    calinski = calinski_harabasz_score(pixels, labels)
    
    print(f"Inertia: {inertia:.2f}")
    print(f"Silhouette Score: {silhouette:.2f}")
    print(f"Calinski-Harabasz Index: {calinski:.2f}")

# Evaluate the segmentation
evaluate_segmentation(image, labels)

Conclusion and Future Directions

K-means clustering provides a simple yet effective approach to image segmentation. While it has limitations, it serves as a foundation for more advanced techniques. Future directions include:

  1. Exploring other clustering algorithms (e.g., DBSCAN, mean-shift)
  2. Incorporating deep learning for feature extraction before clustering
  3. Developing adaptive methods to automatically select the optimal number of clusters
# Placeholder for future improvements
def advanced_segmentation(image):
    # TODO: Implement more sophisticated segmentation techniques
    pass

# Placeholder for automatic cluster number selection
def optimal_cluster_number(image):
    # TODO: Implement method to determine optimal k
    pass

Additional Resources

For further exploration of image segmentation and K-means clustering, consider the following resources:

  1. ArXiv paper: "A Survey of Recent Advances in CNN-based Single Image Crowd Counting and Density Estimation" (arXiv:1707.01202)
  2. ArXiv paper: "Image Segmentation Using Deep Learning: A Survey" (arXiv:2001.05566)
  3. ArXiv paper: "A Review of Modern Deep Learning Techniques for Image Classification" (arXiv:2101.01169)

These papers provide comprehensive overviews of advanced techniques in image analysis and segmentation, building upon the foundations of classic algorithms like K-means clustering.

More in Computer vision

  1. Implementing CNN Image Classification with PyTorch8 min

    Convolutional Neural Networks (CNNs) are a powerful class of deep learning models particularly effective for image classification tasks. PyTorch, a popular...

  2. Explaining the Purpose of MaxPooling in Convolutional Neural Networks12 min

    MaxPooling is a crucial operation in Convolutional Neural Networks (CNNs) that helps reduce the spatial dimensions of feature maps while retaining the most...

  3. Introduction to Convolutional Neural Networks in Python10 min

    Convolutional Neural Networks (CNNs) are a class of deep learning models primarily used for processing grid-like data, such as images. They are designed to...

  4. Evaluating GANs with Inception Score and FID7 min

    The evaluation of Generative Adversarial Networks requires quantitative metrics to assess the quality and diversity of generated samples. Two fundamental...

  5. Leveraging Deep Depth Prediction for Monocular Direct Sparse Odometry:12 min

    Monocular Direct Sparse Odometry (DSO) is a visual odometry method that estimates camera motion and scene structure from a single camera. It's a key...

  6. Deep Residual Learning for Image Recognition in Python8 min

    Deep Residual Learning is a revolutionary approach in image recognition that addresses the problem of vanishing gradients in deep neural networks. It...