diff --git a/python/find_targets.py b/python/find_targets.py new file mode 100644 index 0000000..f4204a9 --- /dev/null +++ b/python/find_targets.py @@ -0,0 +1,78 @@ +#! /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 diff --git a/python/image.jpeg b/python/image.jpeg new file mode 100644 index 0000000..eb00eff Binary files /dev/null and b/python/image.jpeg differ diff --git a/python/image_1.jpeg b/python/image_1.jpeg new file mode 100644 index 0000000..1282163 Binary files /dev/null and b/python/image_1.jpeg differ diff --git a/python/result.png b/python/result.png new file mode 100644 index 0000000..9e1aef9 Binary files /dev/null and b/python/result.png differ diff --git a/python/result_1.png b/python/result_1.png new file mode 100644 index 0000000..1c96ff0 Binary files /dev/null and b/python/result_1.png differ diff --git a/python/test_find_targets.py b/python/test_find_targets.py new file mode 100644 index 0000000..2a6b554 --- /dev/null +++ b/python/test_find_targets.py @@ -0,0 +1,19 @@ +import matplotlib.pyplot as pl +from matplotlib.patches import Rectangle + +from find_targets import find_targets + +fig, ax = pl.subplots(1) +img = pl.imread('image.jpeg') + + +H, L, R, objects = find_targets(img, return_slices=True) + +for o in objects: + x, y = o + r = Rectangle((y.start, x.start), y.stop-y.start, x.stop-x.start, linewidth=1,edgecolor='r',facecolor='none') + ax.add_patch(r) +ax.imshow(img) +ax.plot([H[1], L[1], R[1]], [H[0], L[0], R[0]], 'o', color='red') + +pl.show()