drone-rigide/scripts/find_targets.py

#! /usr/bin/python

import numpy as np
from scipy.ndimage import label, find_objects, center_of_mass


def normalize_coordinates(p, w, h):
    """Normalize the coordinates of a point so it does not depends on the
    size of the picture.

    Args:
        p: The point (tuple-like)
        w: Width of the picture
        h: height of the picture
    Returns:
        (x,y): The normalized coordinates.
    """
    j = float(max(w, h))  # Python2 is shitty

    A = np.array([
        [1.0/j, 0.0],
        [0.0, -1.0/j]
    ])
    b = np.array([-float(w)/(2.0*j), float(h)/(2.0*j)])

    return np.array(p).dot(A) + b


def find_targets(
        picture, threshold_blue=(140, 255), threshold_red=(0, 120),
        threshold_green=(0, 190), return_slices=False, return_binary=False):
    """Find three blue targets in the given picture (RGB matrix).

    Args:
        picture: a 2D matrix of RGB values
        threshold_blue: interval of allowed values for blue channel.
        threshold_red: interval of allowed values for red channel
        threshold_green: interval of allowed values for green channel
        return_slices: Boolean stating if the slices locating the targets
        should be returned.
        return_binary: Boolean stating if the binary map should be returned.

    Returns:
        (H,L,R,[objects]) the positions of the targets in the picture (center of mass). objects is the list of slices controlled by the return_slices parameter.

    Raises:
        ValueError when less than three targets are found.
    """

    blue_points = np.where(
        (threshold_red[0] <= picture[:, :, 0])
        & (picture[:, :, 0] <= threshold_red[1])
        & (threshold_green[0] <= picture[:, :, 1])
        & (picture[:, :, 1] <= threshold_green[1])
        & (threshold_blue[0] <= picture[:, :, 2])
        & (picture[:, :, 2] <= threshold_blue[1]),
        1,
        0
    )

    structure = [
        [0, 1, 0],
        [1, 1, 1],
        [0, 1, 0]
    ]
    labels, n = label(blue_points, structure)

    if n < 3:
        raise ValueError("Less than three potential targets were found")

    objects = [(a[0], a[1], i+1) for i, a in enumerate(find_objects(labels))]

    objects = sorted(
        objects,
        key=lambda x: (x[0].stop - x[0].start) * (x[1].stop - x[1].start)
    )[-3:]

    coordinates = center_of_mass(
        blue_points,
        labels,
        index=[o[2] for o in objects]
    )

    # Highest point
    high = sorted(
        coordinates,
        key=lambda x: x[0]
    )
    H = tuple(reversed(high[0]))

    sides = sorted(
        high[1:],
        key=lambda x: x[1]
    )
    # Leftmost point
    L = tuple(reversed(sides[0]))
    # Rightmost point
    R = tuple(reversed(sides[-1]))

    ret = (H, L, R)
    if return_slices:
        ret = ret + (((o[0], o[1]) for o in objects),)
    if return_binary:
        ret = ret + (blue_points,)

    return ret