drone-rigide/workspace/src/detect_targets/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, threshold_red=120,
        threshold_green=190, return_slices=False):
    """Find three blue targets in the given picture (RGB matrix).

    Args:
        picture: a 2D matrix of RGB values
        threshold_blue: minimal value of the blue channel for a point to be
        considered as blue.
        threshold_red: maximal value of the red channel allowed for a
        target
        threshold_green: maximal value of the green channel allowed for a
        target
        return_slices: Boolean stating if the slices locating the targets
        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(
        (picture[:, :, 2] > threshold_blue)
        & (picture[:, :, 0] < threshold_red)
        & (picture[:, :, 1] < threshold_green),
        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]))

    if return_slices:
        return H, L, R, [(o[0], o[1]) for o in objects]
    else:
        return H, L, R