Merge branch 'armor-lines' into 'origin/armor-colors'

Armor lines

See merge request polystar/robomaster/computer-vision/code!58

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(

robots-at-robots/research/robots_at_robots/dataset/armor_lines_detector.py

+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