drone-rigide/scripts/find_targets.py
2019-06-02 18:02:15 +02:00

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