diff --git a/common/research/dataset/twitch/mask_detector.py b/common/research/dataset/twitch/mask_detector.py new file mode 100644 index 0000000000000000000000000000000000000000..f8d9655983bb3bf4306a67cb402c5361f43e4eb1 --- /dev/null +++ b/common/research/dataset/twitch/mask_detector.py @@ -0,0 +1,80 @@ +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) + + + diff --git a/common/research/dataset/twitch/mask.jpg b/common/research/dataset/twitch/mask_robot_view.jpg similarity index 100% rename from common/research/dataset/twitch/mask.jpg rename to common/research/dataset/twitch/mask_robot_view.jpg diff --git a/common/research/dataset/twitch/new_vision.py b/common/research/dataset/twitch/new_vision.py deleted file mode 100644 index 520d509991e075038fbeeb94d5cc1cfa8ca18ff6..0000000000000000000000000000000000000000 --- a/common/research/dataset/twitch/new_vision.py +++ /dev/null @@ -1,92 +0,0 @@ -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')) -