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 bf50395d934c49e78b4e6bde2aed37e05836d064..d9c9597777c1d7adcba24a70488631872af94dba 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,7 +7,7 @@ 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 ArmorDatasetFactory
 #from ..code/robots-at-robots/research/robots_at_robots/dataset/armor_dataset_factory.py import ArmorDatasetFactory
 
 
@@ -40,8 +40,6 @@ class ArmorLinesDetector:
 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
+    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