#! /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