Python script to find the centre of mass of a black-and-white shape

find_center.py

#!/usr/bin/env python3

import numpy as np
from scipy import ndimage
import matplotlib.image as mpimg
import sys
import os


def print_usage_and_exit():
    """Print usage information and exit with error code 1."""
    script_name = os.path.basename(sys.argv[0])
    print(f"\nUsage: {script_name} <image_path> [mark_type]")
    print("\nFinds the center of mass of a black shape on white/transparent background")
    print("Outputs coordinates and saves a new image with the center marked")
    print("\nmark_type options:")
    print("  point      - marks center with a single red dot")
    print("  vertical   - draws a vertical red line (default)")
    print("  horizontal - draws a horizontal red line")
    print("  cross      - draws both vertical and horizontal red lines")
    sys.exit(1)


def convert_to_grayscale(img):
    """
    Convert any image format (RGB, RGBA, grayscale) to grayscale.
    For RGBA images, blend with white background based on alpha.

    Parameters:
    -----------
    img : ndarray
        Input image (RGB, RGBA, or grayscale)

    Returns:
    --------
    ndarray
        Grayscale image
    """
    # Handle images with alpha channel (RGBA)
    if len(img.shape) > 2 and img.shape[2] == 4:
        # Extract alpha channel
        alpha = img[:, :, 3]

        # Convert the RGB part to grayscale
        rgb = img[:, :, :3]
        gray = np.mean(rgb, axis=2)

        # Create a white background (1.0 for matplotlib's 0-1 scale)
        white_bg = np.ones_like(gray)

        # Blend the grayscale image with white background based on alpha
        # alpha=1 means use the image, alpha=0 means use white
        return gray * alpha + white_bg * (1 - alpha)
    elif len(img.shape) > 2:
        # Regular RGB image without alpha
        return np.mean(img, axis=2)
    else:
        # Already grayscale
        return img


def find_center_of_mass(img_gray):
    """
    Find the center of mass of dark areas in a grayscale image.

    Parameters:
    -----------
    img_gray : ndarray
        Grayscale image (2D array)

    Returns:
    --------
    tuple
        (x, y) coordinates of the center of mass
    """
    # Invert the image - center_of_mass works on bright areas
    # If black is 0 and white is 1, we invert to make black areas become bright
    inverted = 1 - img_gray

    # Find center of mass of the dark areas (now bright after inversion)
    center = ndimage.center_of_mass(inverted)

    # Convert to integer coordinates and return (x, y) format
    y, x = int(center[0]), int(center[1])
    return (x, y)


def mark_center(img_gray, center, mark_type="point"):
    """
    Mark the center of mass on a grayscale image.

    Parameters:
    -----------
    img_gray : ndarray
        Grayscale image
    center : tuple
        (x, y) coordinates of the center to mark
    mark_type : str, optional
        How to mark the center: 'point', 'vertical', 'horizontal', or 'cross'
        Default is 'point'

    Returns:
    --------
    ndarray
        RGB image with the center marked in red
    """
    # Create a new RGB image from the grayscale data
    # Stack the grayscale image three times to create an RGB image
    rgb_output = np.stack([img_gray, img_gray, img_gray], axis=2)

    # Extract coordinates
    x, y = center

    # Mark the center based on the mark_type
    if mark_type == "vertical":
        # Draw a vertical red line through the center
        rgb_output[:, x] = [1, 0, 0]
    elif mark_type == "horizontal":
        # Draw a horizontal red line through the center
        rgb_output[y, :] = [1, 0, 0]
    elif mark_type == "cross":
        # Draw both vertical and horizontal red lines through the center
        rgb_output[:, x] = [1, 0, 0]
        rgb_output[y, :] = [1, 0, 0]
    else:
        # Default to a red point (for 'point' or any invalid type)
        rgb_output[y, x] = [1, 0, 0]

    return rgb_output


def process_image(image_path, mark_type="point"):
    """
    Process an image to find center of mass, mark it, and save the result.

    Parameters:
    -----------
    image_path : str
        Path to the input image file
    mark_type : str, optional
        How to mark the center: 'point', 'vertical', 'horizontal', or 'cross'
        Default is 'point'

    Returns:
    --------
    tuple
        (x, y) coordinates of the center of mass
    """
    # Load the image
    img = mpimg.imread(image_path)

    # Convert to grayscale
    img_gray = convert_to_grayscale(img)

    # Find center of mass
    center = find_center_of_mass(img_gray)

    # Mark the center on the image
    marked_image = mark_center(img_gray, center, mark_type)

    # Create output filename based on input filename
    base_name = image_path.rsplit(".", 1)[0]
    output_path = f"{base_name}_center_{mark_type}.png"

    # Save the marked image
    mpimg.imsave(output_path, marked_image, cmap="gray")

    return center


def main():
    """Main function that processes command line arguments and runs the script."""
    # Print usage if no arguments provided
    if len(sys.argv) < 2:
        print_usage_and_exit()

    # Get image path from command line
    image_path = sys.argv[1]

    # Get marking type if provided, default to 'vertical'
    mark_type = "vertical"

    if len(sys.argv) > 2:
        requested_type = sys.argv[2].lower()
        if requested_type in ["point", "vertical", "horizontal", "cross"]:
            mark_type = requested_type
        else:
            # Invalid mark type provided - show usage and exit
            print(f"Error: Unknown mark type '{requested_type}'")
            print("Valid options are: point, vertical, horizontal, cross")
            print_usage_and_exit()

    # Process the image and get the center of mass coordinates directly
    x, y = process_image(image_path, mark_type=mark_type)

    # Print results with information about the marking type
    print(f"Center of mass of black shape: ({x}, {y})")
    print(f"Marked with red {mark_type} in output image")


if __name__ == "__main__":
    main()