[common] (twitch-dset) clean code

common/research/dataset/twitch/mask_detector.py

+import math
+from pathlib import Path
+from typing import List, Tuple
+
+import cv2
+import numpy as np
+
+
+class Zone:
+    def __init__(self, y_min, y_max, threshold, active_pixels, image_mask):
+
+        self.pixels = [
+            (pix[0], pix[1])
+            for pix in active_pixels
+            if y_min < pix[1] < y_max
+        ]
+
+        self.mean_r = self.get_mean(0, image_mask)
+        self.mean_g = self.get_mean(1, image_mask)
+        self.mean_b = self.get_mean(2, image_mask)
+        self.threshold = threshold
+
+    def get_mean(self, color, img):
+        return sum([img[pix[0], pix[1]][color] for pix in self.pixels]) / len(self.pixels)
+
+    def get_means(self, img):
+        mr, mg, mb = 0, 0, 0
+        for pix in self.pixels:
+            p = img[pix[0], pix[1]]
+            mr += p[0]
+            mg += p[1]
+            mb += p[2]
+
+        n_pixels = len(self.pixels)
+        return mr/n_pixels, mg/n_pixels, mb/n_pixels
+
+    def is_matching(self, frame: np.ndarray):
+        mean_r, mean_g, mean_b = self.get_means(frame)
+        return math.sqrt(pow(mean_r - self.mean_r, 2) + pow(mean_g - self.mean_g, 2) + pow(mean_b - self.mean_b, 2)) < self.threshold
+
+
+class MaskDetector:
+    def __init__(self, image_path: Path, zones_params: List[Tuple[int, int, int]]):
+        image_mask = cv2.imread(str(image_path))
+        active_px = [
+             (a, b)
+             for a in range(0, 720) for b in range(0, 1280)
+             if (
+                image_mask[a, b].any() and
+                int(image_mask[a, b][0]) + int(image_mask[a, b][1]) + int(image_mask[a, b][2]) > 50
+                )
+        ]
+
+        self.zones = [
+            Zone(*zone_params, active_px, image_mask)
+            for zone_params in zones_params
+        ]
+
+    def is_matching(self, frame: np.ndarray):
+        return all(
+            zone.is_matching(frame)
+            for zone in self.zones
+        )
+
+
+robot_view_detector = MaskDetector(
+    Path(__file__).parent / 'mask_robot_view.jpg',
+    [
+        (20, 70, 20),
+        (270, 370, 20),
+        (510, 770, 20),
+    ]
+)
+
+
+def is_image_from_robot_view(frame):
+    return robot_view_detector.is_matching(frame)
+
+
+

common/research/dataset/twitch/new_vision.py

-import math
-from pathlib import Path
-
-import cv2
-
-Y_MIN_Z1 = 20
-Y_MAX_Z1 = 70
-THRESHOLD_ZONE_1 = 45
-
-Y_MIN_Z2 = 270
-Y_MAX_Z2 = 370
-THRESHOLD_ZONE_2 = 45
-
-Y_MIN_Z3 = 510
-Y_MAX_Z3 = 770
-THRESHOLD_ZONE_3 = 45
-
-R = 0
-G = 1
-B = 2
-
-
-class Zone:
-    def __init__(self, y_min, y_max, threshold, active_pixels, image_mask):
-
-        self.y_min = y_min
-        self.y_max = y_max
-        self.pixels = [
-            (pix[0], pix[1])
-            for pix in active_pixels
-            if self.y_min < pix[1] < self.y_max
-        ]
-
-        self.moyenne_r = self.get_moyenne(R, image_mask)
-        self.moyenne_g = self.get_moyenne(G, image_mask)
-        self.moyenne_b = self.get_moyenne(B, image_mask)
-        self.threshold = threshold
-
-    def get_moyenne(self, color, img):
-        return sum([img[pix[0], pix[1]][color] for pix in self.pixels]) / len(self.pixels)
-
-    def get_moyennes(self, img):
-        mr, mg, mb = 0, 0, 0
-        for pix in self.pixels:
-            p = img[pix[0], pix[1]]
-            mr += p[0]
-            mg += p[1]
-            mb += p[2]
-
-        leng = len(self.pixels)
-        return mr/leng, mg/leng, mb/leng
-
-
-class Mask:
-    def __init__(self, image):
-        self.image = cv2.imread(image)     # Passer l'url et l'ouvrir avec cv2
-        self.active_px = [                  # Iterer sur tous les pixels de l'image pour trouver ceux qui sont !=  0 0 0
-                         (a, b)
-                         for a in range(0, 720) for b in range(0, 1280)
-                         if (self.image[a, b].any() and
-                             (int(self.image[a, b][0]) + int(self.image[a, b][1]) + int(self.image[a, b][2]) > 50))
-                        ]
-
-        self.zone1 = Zone(Y_MIN_Z1, Y_MAX_Z1, 20, self.active_px, self.image)
-        self.zone2 = Zone(Y_MIN_Z2, Y_MAX_Z2, 20, self.active_px, self.image)
-        self.zone3 = Zone(Y_MIN_Z3, Y_MAX_Z3, 20, self.active_px, self.image)
-
-
-def process_frame_moyennes(frame, mask):
-
-    if(process_zone_moyennes(frame, mask.zone1) and    # Passer à travers les zones de la classe mask, faire la moyenne des pixels
-       process_zone_moyennes(frame, mask.zone2) and
-       process_zone_moyennes(frame, mask.zone3)):
-        return 1
-    else:
-        return 0
-
-
-def process_zone_moyennes(frame, zone):                # calculer la diff avec la moyenne des bonnes immages
-    moy_r, moy_g, moy_b = zone.get_moyennes(frame)
-    if math.sqrt(pow(moy_r - zone.moyenne_r, 2) + pow(moy_g - zone.moyenne_g, 2) + pow(moy_b - zone.moyenne_b, 2)) < zone.threshold:
-        return 1
-    else:
-        return 0
-
-
-def is_image_from_robot_view(frame):
-    return process_frame_moyennes(frame, mask)
-
-
-mask = Mask(str(Path(__file__).parent / 'mask.jpg'))
-