Détection des cibles

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

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()