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