106 lines
2.8 KiB
Python
106 lines
2.8 KiB
Python
#! /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
|