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 index d9c9597777c1d7adcba24a70488631872af94dba..7541e8d760b23202bfd43ed8cf2a8cc10057af7f 100644 --- 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 @@ -7,19 +7,24 @@ 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 ..armor_dataset_factory import +#research.robots_at_robots.dataset.armor_dataset_factory import ArmorDatasetFactory ?? Will not find 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? + #line detection + @staticmethod + def detect_lines(image: Image) -> Image: + edges = cv2.Canny(image, 100, 200) lines = cv2.HoughLines(edges, 1, np.pi / 180, 50) - len(lines) + return lineDetector.draw_lines(image, lines) - armor_lines = image.copy() + # 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] @@ -32,14 +37,21 @@ class ArmorLinesDetector: 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') + cv2.line(image_lines, (x1, y1), (x2, y2), (0, 0, 255), 2) + + return image_lines + if __name__ == "__main__": lineDetector = ArmorLinesDetector() - for i, (armor_img, armor, k, p) in enumerate(ArmorDatasetFactory.from_dataset(TwitchROCODataset.TWITCH_470158483)): - lineDetector.detect_lines(armor_img) \ No newline at end of file + #for i, (armor_img, armor, k, p) in enumerate(ArmorDatasetFactory.from_dataset(TwitchROCODataset.TWITCH_470158483)): + # armor_lines = lineDetector.detect_lines(armor_img) + test = cv2.imread('test.png') # 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() \ 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