From e89b2b8bf061d8732bc84b3765a86b29102750dd Mon Sep 17 00:00:00 2001
From: Yasmine Moumou <yasmine.moumou@polymtl.ca>
Date: Tue, 11 Aug 2020 20:55:00 -0400
Subject: [PATCH] [robots](armor-lines) Add line detection function structure
---
.../dataset/armor_lines_detector.py | 47 +++++++++++++++++++
1 file changed, 47 insertions(+)
create mode 100644 robots-at-robots/research/robots_at_robots/dataset/armor_lines_detector.py
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 0000000..bf50395
--- /dev/null
+++ b/robots-at-robots/research/robots_at_robots/dataset/armor_lines_detector.py
@@ -0,0 +1,47 @@
+from typing import Tuple
+
+import cv2
+import numpy as np
+import matplotlib.pyplot as plt
+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 .armor_dataset_factory import ArmorDatasetFactory
+#from ..code/robots-at-robots/research/robots_at_robots/dataset/armor_dataset_factory.py import ArmorDatasetFactory
+
+
+class ArmorLinesDetector:
+ @staticmethod
+ def detect_lines(image: Image):
+
+ edges = cv2.Canny(image, image.size, image.size) #size?
+ lines = cv2.HoughLines(edges, 1, np.pi / 180, 50)
+ len(lines)
+
+ armor_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(armor_lines, (x1, y1), (x2, y2), (0, 0, 255), 2)
+ plt.imshow(armor_lines, cmap='gray')
+ plt.title('50 lignes principales')
+
+
+if __name__ == "__main__":
+
+ lineDetector = ArmorLinesDetector()
+ armorFactory = ArmorDatasetFactory()
+
+ armorData = armorFactory.from_dataset(TwitchROCODataset.TWITCH_470158483)
+ for element in armorData:
+ lineDetector.detect_lines(element[0]) # get first image to test
\ No newline at end of file
--
GitLab