78 lines
2 KiB
Python
78 lines
2 KiB
Python
#! /usr/bin/python
|
|
|
|
import numpy as np
|
|
from scipy.ndimage import label, find_objects, center_of_mass
|
|
|
|
|
|
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[1]
|
|
)
|
|
H = high[0]
|
|
|
|
sides = sorted(
|
|
high[1:],
|
|
key=lambda x: x[0]
|
|
)
|
|
# Leftmost point
|
|
L = sides[0]
|
|
# Rightmost point
|
|
R = sides[-1]
|
|
|
|
if return_slices:
|
|
return H, L, R, [(o[0], o[1]) for o in objects]
|
|
else:
|
|
return H, L, R
|