Image Segmentation

Abstract

Image Segmentation is a Python project that uses deep learning to segment images. The application features image processing, model training, and a CLI interface, demonstrating best practices in computer vision and AI.

Prerequisites

• Python 3.8 or above
• A code editor or IDE
• Basic understanding of deep learning and computer vision
• Required libraries: tensorflow, numpy, opencv-python

Before you Start

Install Python and the required libraries:

Install dependencies

pip install tensorflow numpy opencv-python

Getting Started

Create a Project

1. Create a folder named image-segmentation.
2. Open the folder in your code editor or IDE.
3. Create a file named image_segmentation.py.
4. Copy the code below into your file.

Write the Code

Image Segmentation

pch.viewSource

Image Segmentation

"""
Image Segmentation

A full image segmentation pipeline using OpenCV and scikit-image. Includes image loading, segmentation (thresholding, k-means), visualization, and CLI for batch processing.
"""
import cv2
import numpy as np
import argparse
from skimage import filters
from sklearn.cluster import KMeans
import os

def threshold_segmentation(image_path):
    img = cv2.imread(image_path, 0)
    if img is None:
        print(f"Error: Could not load image {image_path}")
        return None
    thresh_val = filters.threshold_otsu(img)
    binary = img > thresh_val
    show_or_save('Threshold Segmentation', binary.astype(np.uint8)*255)
    wait_or_skip(0)
    cv2.destroyAllWindows() if "--show" in __import__("sys").argv else None
    return binary

def kmeans_segmentation(image_path, k=2):
    img = cv2.imread(image_path)
    if img is None:
        print(f"Error: Could not load image {image_path}")
        return None
    Z = img.reshape((-1,3))
    Z = np.float32(Z)
    kmeans = KMeans(n_clusters=k, random_state=42)
    labels = kmeans.fit_predict(Z)
    centers = np.uint8(kmeans.cluster_centers_)
    segmented = centers[labels].reshape(img.shape)
    show_or_save('K-means Segmentation', segmented)
    wait_or_skip(0)
    cv2.destroyAllWindows() if "--show" in __import__("sys").argv else None
    return segmented

def sample_image():
    """Write a small scene with known regions, so there is something to segment.

    Requiring --image means the project cannot be run without finding a
    picture first. Generating one keeps the demo self-contained, and a scene
    with known regions is more useful for judging a segmentation than an
    arbitrary photograph.
    """
    import tempfile
    from pathlib import Path

    import numpy as np
    from PIL import Image

    size = 240
    grid_y, grid_x = np.mgrid[0:size, 0:size] / size
    scene = np.where(grid_y > 0.55, 60, 190).astype("uint8")
    disc = np.hypot(grid_x - 0.35, grid_y - 0.35) < 0.18
    scene[disc] = 125
    rng = np.random.default_rng(0)
    scene = np.clip(scene + rng.normal(0, 6, scene.shape), 0, 255)

    path = Path(tempfile.mkdtemp(prefix="segmentation-demo-")) / "scene.png"
    Image.fromarray(scene.astype("uint8")).save(path)
    print(f"no --image given, so a test scene was written to {path}")
    print("it has three regions: a bright sky, a dark ground and a mid disc\n")
    return str(path)


def main():
    parser = argparse.ArgumentParser(description="Image Segmentation")
    parser.add_argument('--image', type=str, default=None,
                        help='Path to image file (a demo scene is generated '
                             'when this is omitted)')
    parser.add_argument('--mode', type=str, choices=['threshold', 'kmeans'],
                        default='kmeans', help='Segmentation mode')
    parser.add_argument('--k', type=int, default=2, help='Number of clusters for k-means')
    args = parser.parse_args()
    if args.image is None:
        args.image = sample_image()
    if args.mode == 'threshold':
        threshold_segmentation(args.image)
    elif args.mode == 'kmeans':
        kmeans_segmentation(args.image, args.k)

def wait_or_skip(delay=0):
    """cv2.waitKey needs a window; without one it raises. Skip it instead."""
    import sys

    if "--show" in sys.argv:
        return cv2.waitKey(delay)
    return -1


def show_or_save(title, image, _counter=[0]):
    """Display the frame, or write it to a file when no window is available.

    Headless OpenCV has no GUI at all, and even a full build cannot open a
    window over SSH or inside a container. Falling back to a file keeps the
    project runnable everywhere and leaves something to look at afterwards.
    """
    import os
    import re as _re

    if "--show" in __import__("sys").argv:
        cv2.imshow(title, image)
        return None
    _counter[0] += 1
    stem = _re.sub(r"\W+", "_", title).strip("_").lower() or "frame"
    name = f"{stem}.png" if _counter[0] == 1 else f"{stem}_{_counter[0]}.png"
    cv2.imwrite(name, image)
    print(f"saved {name}  ({os.path.getsize(name):,} bytes)")
    return name


if __name__ == "__main__":
    main()

Example Usage

Run image segmentation

python image_segmentation.py

What it produces

Running the file exactly as it ships takes 6.9 s and prints:

python image_segmentation.py

no --image given, so a test scene was written to C:\Users\Zimyo\AppData\Local\Temp\segmentation-demo-hibfvqk8\scene.png
it has three regions: a bright sky, a dark ground and a mid disc
 
saved k_means_segmentation.png  (5,031 bytes)

figure Produced by this project, not drawn for the page matplotlib

Written by the run above. If the project stops producing it, the page's figure asset goes missing and check_docs reports it — which is the point of generating it rather than drawing it.

How it fits together

Read from the top: this is what runs when you execute the file, and which function calls which. It is generated from the code, so it cannot drift from it.

diagram Diagram mermaid

  flowchart TD
  RUN(["python image_segmentation.py"])
  threshold_segmentation("threshold_segmentation")
  kmeans_segmentation("kmeans_segmentation")
  sample_image("sample_image")
  main("main")
  wait_or_skip("wait_or_skip")
  show_or_save("show_or_save")
  RUN --> main
  kmeans_segmentation --> show_or_save
  kmeans_segmentation --> wait_or_skip
  main --> kmeans_segmentation
  main --> sample_image
  main --> threshold_segmentation
  threshold_segmentation --> show_or_save
  threshold_segmentation --> wait_or_skip

Explanation

Key Features

• Image Segmentation: Segments images using deep learning.
• Image Processing: Prepares images for segmentation.
• Error Handling: Validates inputs and manages exceptions.
• CLI Interface: Interactive command-line usage.

Code Breakdown

1. What it imports (lines 6–11)

image_segmentation.py

import cv2
import numpy as np
import argparse
from skimage import filters
from sklearn.cluster import KMeans
import os

2. threshold_segmentation — the function (lines 13–23)

image_segmentation.py

def threshold_segmentation(image_path):
    img = cv2.imread(image_path, 0)
    if img is None:
        print(f"Error: Could not load image {image_path}")
        return None
    thresh_val = filters.threshold_otsu(img)
    binary = img > thresh_val
    cv2.imshow('Threshold Segmentation', binary.astype(np.uint8)*255)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    return binary

3. kmeans_segmentation — the function (lines 25–39)

image_segmentation.py

def kmeans_segmentation(image_path, k=2):
    img = cv2.imread(image_path)
    if img is None:
        print(f"Error: Could not load image {image_path}")
        return None
    Z = img.reshape((-1,3))
    Z = np.float32(Z)
    kmeans = KMeans(n_clusters=k, random_state=42)
    labels = kmeans.fit_predict(Z)
    centers = np.uint8(kmeans.cluster_centers_)
    segmented = centers[labels].reshape(img.shape)
    cv2.imshow('K-means Segmentation', segmented)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    return segmented

4. main — the function (lines 41–50)

image_segmentation.py

def main():
    parser = argparse.ArgumentParser(description="Image Segmentation")
    parser.add_argument('--image', type=str, required=True, help='Path to image file')
    parser.add_argument('--mode', type=str, choices=['threshold', 'kmeans'], required=True, help='Segmentation mode')
    parser.add_argument('--k', type=int, default=2, help='Number of clusters for k-means')
    args = parser.parse_args()
    if args.mode == 'threshold':
        threshold_segmentation(args.image)
    elif args.mode == 'kmeans':
        kmeans_segmentation(args.image, args.k)

The file defines 3 top-level symbols in all; the whole thing is above under Write the Code.

Features

• Image Segmentation: Deep learning and image processing
• Modular Design: Separate functions for each task
• Error Handling: Manages invalid inputs and exceptions
• Production-Ready: Scalable and maintainable code

Next Steps

Enhance the project by:

• Integrating with real image datasets
• Supporting advanced segmentation algorithms
• Creating a GUI for segmentation
• Adding real-time segmentation
• Unit testing for reliability

Educational Value

This project teaches:

• Computer Vision: Image segmentation and deep learning
• Software Design: Modular, maintainable code
• Error Handling: Writing robust Python code

Real-World Applications

• Medical Imaging
• AI Platforms
• Robotics

Conclusion

Image Segmentation demonstrates how to build a scalable and accurate image segmentation tool using Python. With modular design and extensibility, this project can be adapted for real-world applications in AI, healthcare, and more. For more advanced projects, visit Python Central Hub.