[YOLOV8] Add--datasets-dir to yolov8 entrypoints

src/deepsparse/yolov8/utils/validation/helpers.py

@@ -12,16 +12,53 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import argparse
+import glob
 import os
 import warnings
 from typing import List, Optional, Union
 
+import yaml
+
 import torch
 from deepsparse.yolo import YOLOOutput as YOLODetOutput
 from deepsparse.yolov8.schemas import YOLOSegOutput
+from ultralytics.yolo.data.utils import ROOT
+
+
+__all__ = ["data_from_dataset_path", "schema_to_tensor", "check_coco128_segmentation"]
+
+
+def data_from_dataset_path(data: str, dataset_path: str) -> str:
+    """
+    Given a dataset name, fetch the yaml config for the dataset
+    from the Ultralytics dataset repo, overwrite its 'path'
+    attribute (dataset root dir) to point to the `dataset_path`
+    and finally save it to the current working directory.
+    This allows to create load data yaml config files that point
+    to the arbitrary directories on the disk.
+
+    :param data: name of the dataset (e.g. "coco.yaml")
+    :param dataset_path: path to the dataset directory
+    :return: a path to the new yaml config file
+       (saved in the current working directory)
+    """
+    ultralytics_dataset_path = glob.glob(os.path.join(ROOT, "**", data), recursive=True)
+    if len(ultralytics_dataset_path) != 1:
+        raise ValueError(
+            "Expected to find a single path to the "
+            f"dataset yaml file: {data}, but found {ultralytics_dataset_path}"
+        )
+    ultralytics_dataset_path = ultralytics_dataset_path[0]
+    with open(ultralytics_dataset_path, "r") as f:
+        yaml_config = yaml.safe_load(f)
+        yaml_config["path"] = dataset_path
 
+        yaml_save_path = os.path.join(os.getcwd(), data)
 
-__all__ = ["schema_to_tensor", "check_coco128_segmentation"]
+        # save the new dataset yaml file
+        with open(yaml_save_path, "w") as outfile:
+            yaml.dump(yaml_config, outfile, default_flow_style=False)
+        return yaml_save_path
 
 
 def schema_to_tensor(

src/deepsparse/yolov8/validation.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional
+
 import click
 
 from deepsparse import Pipeline
@@ -20,6 +22,7 @@
     DeepSparseDetectionValidator,
     DeepSparseSegmentationValidator,
     check_coco128_segmentation,
+    data_from_dataset_path,
 )
 from ultralytics.yolo.cfg import get_cfg
 from ultralytics.yolo.utils import DEFAULT_CFG
@@ -63,16 +66,6 @@
     show_default=True,
     help="Validation batch size",
 )
-@click.option(
-    "--stride",
-    type=int,
-    default=32,
-    show_default=True,
-    help="YOLOv8 can handle arbitrary sized images as long as "
-    "both sides are a multiple of 32. This is because the "
-    "maximum stride of the backbone is 32 and it is a fully "
-    "convolutional network.",
-)
 @click.option(
     "--engine-type",
     default=DEEPSPARSE_ENGINE,
@@ -95,15 +88,21 @@
     show_default=True,
     help="A subtask of YOLOv8 to run. Default is `detection`.",
 )
+@click.option(
+    "--dataset-path",
+    type=str,
+    default=None,
+    help="Path to override default dataset path.",
+)
 def main(
     dataset_yaml: str,
     model_path: str,
     batch_size: int,
     num_cores: int,
     engine_type: str,
-    stride: int,
     device: str,
     subtask: str,
+    dataset_path: Optional[str],
 ):
 
     pipeline = Pipeline.create(
@@ -124,6 +123,8 @@ def main(
             f"Dataset yaml {dataset_yaml} is not supported. "
             f"Supported dataset configs are {SUPPORTED_DATASET_CONFIGS})"
         )
+    if dataset_path is not None:
+        args.data = data_from_dataset_path(args.data, dataset_path)
     classes = {label: class_ for (label, class_) in enumerate(COCO_CLASSES)}
 
     if subtask == "detection":