working

Author: dbogunowicz

src/sparseml/yolov8/export.py

@@ -12,11 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 
 import click
 from sparseml.yolov8.trainers import SparseYOLO
-from ultralytics.yolo.utils import USER_CONFIG_DIR, get_settings, yaml_save
 
 
 # Options generated from
@@ -88,12 +86,6 @@
     help="Path to override default datasets dir.",
 )
 def main(**kwargs):
-    if kwargs["datasets_dir"] is not None:
-        settings = get_settings()
-        settings["datasets_dir"] = os.path.abspath(
-            os.path.expanduser(kwargs["datasets_dir"])
-        )
-        yaml_save(USER_CONFIG_DIR / "settings.yaml", settings)
 
     model = SparseYOLO(kwargs["model"])
     model.export(**kwargs)

src/sparseml/yolov8/exported/DetectionModel.onnx

@@ -13,11 +13,10 @@
 # limitations under the License.
 
 import logging
-import os
 
 import click
 from sparseml.yolov8.trainers import SparseYOLO
-from ultralytics.yolo.utils import USER_CONFIG_DIR, get_settings, yaml_save
+from sparseml.yolov8.utils import data_from_datasets_dir
 
 
 logger = logging.getLogger()
@@ -219,11 +218,7 @@
 )
 def main(**kwargs):
     if kwargs["datasets_dir"] is not None:
-        settings = get_settings()
-        settings["datasets_dir"] = os.path.abspath(
-            os.path.expanduser(kwargs["datasets_dir"])
-        )
-        yaml_save(USER_CONFIG_DIR / "settings.yaml", settings)
+        kwargs["data"] = data_from_datasets_dir(kwargs["data"], kwargs["datasets_dir"])
 
     model = SparseYOLO(kwargs["model"])
     model.train(**kwargs)

src/sparseml/yolov8/exported/deployment/DetectionModel.onnx

@@ -32,7 +32,11 @@
 from sparseml.pytorch.utils.helpers import download_framework_model_by_recipe_type
 from sparseml.pytorch.utils.logger import LoggerManager, PythonLogger, WANDBLogger
 from sparseml.yolov8.modules import Bottleneck, Conv
-from sparseml.yolov8.utils import check_coco128_segmentation, create_grad_sampler
+from sparseml.yolov8.utils import (
+    check_coco128_segmentation,
+    create_grad_sampler,
+    data_from_datasets_dir,
+)
 from sparseml.yolov8.utils.export_samples import export_sample_inputs_outputs
 from sparseml.yolov8.validators import (
     SparseClassificationValidator,
@@ -655,6 +659,11 @@ def export(self, **kwargs):
             if kwargs["device"] is not None and "cpu" not in kwargs["device"]:
                 overrides["device"] = "cuda:" + kwargs["device"]
             overrides["deterministic"] = kwargs["deterministic"]
+            if kwargs["datasets_dir"] is not None:
+                overrides["data"] = data_from_datasets_dir(
+                    overrides["data"], kwargs["datasets_dir"]
+                )
+
             trainer = self.TrainerClass(overrides=overrides)
             self.model = self.model.to(trainer.device)
 
@@ -710,9 +719,12 @@ def export(self, **kwargs):
         if args["export_samples"]:
             trainer_config = get_cfg(cfg=DEFAULT_SPARSEML_CONFIG_PATH)
 
+            if args["datasets_dir"] is not None:
+                args["data"] = data_from_datasets_dir(
+                    args["data"], args["datasets_dir"]
+                )
             trainer_config.data = args["data"]
             trainer_config.imgsz = args["imgsz"]
-
             trainer = DetectionTrainer(trainer_config)
             # inconsistency in name between
             # validation and test sets

src/sparseml/yolov8/exported/deployment/config.json

@@ -12,17 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import glob
 import os
 import warnings
 from argparse import Namespace
 from typing import Any, Dict
 
+import yaml
+
 from ultralytics.yolo.data.dataloaders.v5loader import create_dataloader
+from ultralytics.yolo.data.utils import ROOT
 from ultralytics.yolo.engine.model import DetectionModel
 from ultralytics.yolo.engine.trainer import BaseTrainer
 
 
-__all__ = ["check_coco128_segmentation", "create_grad_sampler"]
+__all__ = [
+    "check_coco128_segmentation",
+    "create_grad_sampler",
+    "data_from_datasets_dir",
+]
 
 
 def check_coco128_segmentation(args: Namespace) -> Namespace:
@@ -69,3 +77,41 @@ def create_grad_sampler(
         / train_loader.batch_size,
     )
     return grad_sampler
+
+
+def data_from_datasets_dir(data: str, datasets_dir: 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 `datasets_dir`
+    and finally save it to the `datasets_dir` 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 datasets_dir: path to the directory where the dataset is expected to be
+        (and where the new yaml config will be saved)
+     :return: a path to the new yaml config file
+    """
+    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)
+
+        local_dataset_path = os.path.join(
+            datasets_dir, os.path.basename(yaml_config["path"])
+        )
+        yaml_config["path"] = local_dataset_path
+
+        yaml_save_path = os.path.join(datasets_dir, data)
+
+        # 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

src/sparseml/yolov8/train.py

@@ -12,11 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 
 import click
 from sparseml.yolov8.trainers import SparseYOLO
-from ultralytics.yolo.utils import USER_CONFIG_DIR, get_settings, yaml_save
+from sparseml.yolov8.utils import data_from_datasets_dir
 
 
 @click.command(
@@ -74,16 +73,12 @@
 @click.option(
     "--datasets-dir",
     type=str,
-    default=None,
+    default="/home/ubuntu/damian/sparseml/funny_dir",
     help="Path to override default datasets dir.",
 )
 def main(**kwargs):
     if kwargs["datasets_dir"] is not None:
-        settings = get_settings()
-        settings["datasets_dir"] = os.path.abspath(
-            os.path.expanduser(kwargs["datasets_dir"])
-        )
-        yaml_save(USER_CONFIG_DIR / "settings.yaml", settings)
+        kwargs["data"] = data_from_datasets_dir(kwargs["data"], kwargs["datasets_dir"])
 
     model = SparseYOLO(kwargs["model"])
     model.val(**kwargs)