diff --git a/robots-at-robots/research/robots_at_robots/armor_color/baseline_experiments.py b/robots-at-robots/research/robots_at_robots/armor_color/baseline_experiments.py index c0ab8deab01a3df1c18be2d2c7b24fb26d5474c5..d723c6810c9966dbe41ff157542184ae5bd98403 100644 --- a/robots-at-robots/research/robots_at_robots/armor_color/baseline_experiments.py +++ b/robots-at-robots/research/robots_at_robots/armor_color/baseline_experiments.py @@ -13,8 +13,8 @@ if __name__ == "__main__": logging.getLogger().setLevel("INFO") reporter = ArmorColorPipelineReporterFactory.from_roco_datasets( - train_roco_datasets=[ROCODatasetsZoo.TWITCH.T470151286, ROCODatasetsZoo.TWITCH.T470150052], - test_roco_datasets=[ROCODatasetsZoo.TWITCH.T470152289], + train_roco_datasets=[ROCODatasetsZoo.TWITCH.T470149568], + test_roco_datasets=[ROCODatasetsZoo.TWITCH.T470149568], ) red_blue_comparison_pipeline = ClassifierImagePipeline( diff --git a/robots-at-robots/research/robots_at_robots/dataset/armor_lines_detector.py b/robots-at-robots/research/robots_at_robots/dataset/armor_lines_detector.py new file mode 100644 index 0000000000000000000000000000000000000000..d472b704df6f084c21489961d2d95c53ffd47ebb --- /dev/null +++ b/robots-at-robots/research/robots_at_robots/dataset/armor_lines_detector.py @@ -0,0 +1,68 @@ +from itertools import islice +from typing import Tuple + +import cv2 +import numpy as np +import matplotlib.pyplot as plt + +from dataset.armor_dataset_factory import ArmorDataset +from polystar.common.models.image import Image + +from polystar.common.models.image_annotation import ImageAnnotation +#from research.common.dataset.twitch.twitch_roco_datasets import TwitchROCODataset +#from research.robots_at_robots.dataset.armor_dataset_factory import ArmorDatasetFactory +from research.common.datasets.roco.zoo.roco_dataset_zoo import ROCODatasetsZoo + + +class ArmorLinesDetector: + + #line detection + @staticmethod + def detect_lines(image: Image) -> Image: + + edges = cv2.Canny(image, 50, 200) + lines = cv2.HoughLines(edges, 1, np.pi / 180, threshold=25) + return lineDetector.draw_lines(image, lines) + + # Drawing lines on image + @staticmethod + def draw_lines(image: Image, lines) -> Image: + image_lines = image.copy() + + for objet in lines[0:50]: + rho, theta = objet[0] + a = np.cos(theta) + b = np.sin(theta) + x0 = a * rho + y0 = b * rho + x1 = int(x0 + 1000 * (-b)) + y1 = int(y0 + 1000 * (a)) + x2 = int(x0 - 1000 * (-b)) + y2 = int(y0 - 1000 * (a)) + + cv2.line(image_lines, (x1, y1), (x2, y2), (0, 0, 255), 1) + + return image_lines + + + + +if __name__ == "__main__": + + lineDetector = ArmorLinesDetector() + armor_lines = [] + for _armor_img, _armor, _name in islice(ArmorDataset(ROCODatasetsZoo.TWITCH.T470149568.to_images()), 20, 30): + # armor_lines = lineDetector.detect_lines(_armor_img) + #print(_name, repr(_armor)) + #plt.imshow(armor_lines) + #plt.show() + #plt.clf() + + + test = cv2.imread(r'C:\Users\moumo\Documents\polySTAR\code\dataset\twitch\v1\470149568\colors\470149568-frame-001120-0-red.jpg') # ouverture + test = cv2.cvtColor(test, cv2.COLOR_BGR2RGB) # conversion des couleurs + armor_lines = lineDetector.detect_lines(test) + plt.imshow(armor_lines, cmap='gray') + plt.title('50 lignes principales') + plt.show() + plt.clf() \ No newline at end of file diff --git a/robots-at-robots/research/robots_at_robots/dataset/test.png b/robots-at-robots/research/robots_at_robots/dataset/test.png new file mode 100644 index 0000000000000000000000000000000000000000..51d6eb8103b216c65d3a060562153e342b6507fb Binary files /dev/null and b/robots-at-robots/research/robots_at_robots/dataset/test.png differ