Fix yolov6 parser (#337)

Co-authored-by: klemen1999 <[email protected]>

Author: JSabadin

luxonis_ml/data/parsers/luxonis_parser.py

@@ -127,6 +127,9 @@ class name will be taken from the dictionary.
             self.parser_type = (
                 ParserType.DIR
                 if Path(self.dataset_dir / "train").exists()
+                or Path(
+                    self.dataset_dir / "images" / "train"
+                ).exists()  # temporary fix for yolov6
                 else ParserType.SPLIT
             )
         else:

tests/test_data/test_parsers.py

@@ -103,3 +103,116 @@ def test_dir_parser(
     task_types = {get_task_type(task) for task in ann}
     assert task_types == expected_task_types
     dataset.delete_dataset()
+
+
+@pytest.mark.parametrize(
+    ("url", "dataset_type", "expected_task_types"),
+    [
+        (
+            "COCO_people_subset.zip",
+            "coco",
+            {
+                "boundingbox",
+                "keypoints",
+                "segmentation",
+                "classification",
+                "instance_segmentation",
+            },
+        ),
+        (
+            "Thermal_Dogs_and_People.v1-resize-416x416.coco.zip",
+            "coco",
+            {"boundingbox", "classification"},
+        ),
+        (
+            "roboflow://team-roboflow/coco-128/2/coco",
+            "coco",
+            {"boundingbox", "classification"},
+        ),
+        (
+            "Thermal_Dogs_and_People.v1-resize-416x416.voc.zip",
+            "voc",
+            {"boundingbox", "classification"},
+        ),
+        (
+            "Thermal_Dogs_and_People.v1-resize-416x416.darknet.zip",
+            "darknet",
+            {"boundingbox", "classification"},
+        ),
+        (
+            "Thermal_Dogs_and_People.v1-resize-416x416.yolov4pytorch.zip",
+            "yolov4",
+            {"boundingbox", "classification"},
+        ),
+        (
+            "Thermal_Dogs_and_People.v1-resize-416x416.mt-yolov6.zip",
+            "yolov6",
+            {"boundingbox", "classification"},
+        ),
+        (
+            "Thermal_Dogs_and_People.v1-resize-416x416.createml.zip",
+            "createml",
+            {"boundingbox", "classification"},
+        ),
+        (
+            "Thermal_Dogs_and_People.v1-resize-416x416.tensorflow.zip",
+            "tfcsv",
+            {"boundingbox", "classification"},
+        ),
+        (
+            "D2_Tile.png-mask-semantic.zip",
+            "segmask",
+            {"segmentation", "classification"},
+        ),
+        (
+            "Flowers_Classification.v2-raw.folder.zip",
+            "clsdir",
+            {"classification"},
+        ),
+        (
+            "D2_ParkingLot.zip",
+            "solo",
+            {"boundingbox", "segmentation", "classification", "keypoints"},
+        ),
+        (
+            "D2_ParkingLot_Native.zip",
+            "native",
+            {
+                "boundingbox",
+                "instance_segmentation",
+                "classification",
+                "keypoints",
+                "metadata/color",
+                "metadata/brand",
+            },
+        ),
+    ],
+)
+def test_dir_parser_explicit_type(
+    dataset_name: str,
+    url: str,
+    dataset_type: str,
+    expected_task_types: set[str],
+    storage_url: str,
+    tempdir: Path,
+):
+    if not url.startswith("roboflow://"):
+        url = f"{storage_url}/{url}"
+
+    elif environ.ROBOFLOW_API_KEY is None:
+        pytest.skip("Roboflow API key is not set")
+
+    parser = LuxonisParser(
+        url,
+        dataset_name=dataset_name,
+        dataset_type=dataset_type,  # type: ignore
+        delete_local=True,
+        save_dir=tempdir,
+    )
+    dataset = parser.parse()
+    assert len(dataset) > 0
+    loader = LuxonisLoader(dataset)
+    _, ann = next(iter(loader))
+    task_types = {get_task_type(task) for task in ann}
+    assert task_types == expected_task_types
+    dataset.delete_dataset()