CoreML NMS and half fixes (#143)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

ultralytics/yolo/engine/exporter.py

@@ -67,7 +67,7 @@ import torch
 
 import ultralytics
 from ultralytics.nn.modules import Detect, Segment
-from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel, attempt_load_weights
+from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel
 from ultralytics.yolo.configs import get_config
 from ultralytics.yolo.data.dataloaders.stream_loaders import LoadImages
 from ultralytics.yolo.data.utils import check_dataset
@@ -154,7 +154,7 @@ class Exporter:
         # Load PyTorch model
         self.device = select_device(self.args.device or 'cpu')
         if self.args.half:
-            if self.device.type == 'cpu' or not coreml:
+            if self.device.type == 'cpu' and not coreml:
                 LOGGER.info('half=True only compatible with GPU or CoreML export, i.e. use device=0 or format=coreml')
                 self.args.half = False
             assert not self.args.dynamic, '--half not compatible with --dynamic, i.e. use either --half or --dynamic'
@@ -769,17 +769,22 @@ class Exporter:
 def export(cfg):
     cfg.model = cfg.model or "yolov8n.yaml"
     cfg.format = cfg.format or "torchscript"
-    exporter = Exporter(cfg)
 
-    model = None
-    if isinstance(cfg.model, (str, Path)):
-        if Path(cfg.model).suffix == '.yaml':
-            model = DetectionModel(cfg.model)
-        elif Path(cfg.model).suffix == '.pt':
-            model = attempt_load_weights(cfg.model, fuse=True)
-        else:
-            TypeError(f'Unsupported model type {cfg.model}')
-    exporter(model=model)
+    # exporter = Exporter(cfg)
+    #
+    # model = None
+    # if isinstance(cfg.model, (str, Path)):
+    #     if Path(cfg.model).suffix == '.yaml':
+    #         model = DetectionModel(cfg.model)
+    #     elif Path(cfg.model).suffix == '.pt':
+    #         model = attempt_load_weights(cfg.model, fuse=True)
+    #     else:
+    #         TypeError(f'Unsupported model type {cfg.model}')
+    # exporter(model=model)
+
+    from ultralytics import YOLO
+    model = YOLO(cfg.model)
+    model.export(**cfg)
 
 
 if __name__ == "__main__":

ultralytics/yolo/engine/model.py

@@ -64,7 +64,7 @@ class YOLO:
             verbose (bool): display model info on load
         """
         cfg = check_yaml(cfg)  # check YAML
-        cfg_dict = yaml_load(cfg)  # model dict
+        cfg_dict = yaml_load(cfg, append_filename=True)  # model dict
         self.task = guess_task_from_head(cfg_dict["head"][-1][-2])
         self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = \
             self._guess_ops_from_task(self.task)
@@ -183,7 +183,7 @@ class YOLO:
         overrides.update(kwargs)
         if kwargs.get("cfg"):
             LOGGER.info(f"cfg file passed. Overriding default params with {kwargs['cfg']}.")
-            overrides = yaml_load(check_yaml(kwargs["cfg"]))
+            overrides = yaml_load(check_yaml(kwargs["cfg"]), append_filename=True)
         overrides["task"] = self.task
         overrides["mode"] = "train"
         if not overrides.get("data"):

ultralytics/yolo/engine/validator.py

@@ -157,7 +157,8 @@ class BaseValidator:
         self.run_callbacks('on_val_end')
         if self.training:
             model.float()
-            return {**stats, **trainer.label_loss_items(self.loss.cpu() / len(self.dataloader), prefix="val")}
+            results = {**stats, **trainer.label_loss_items(self.loss.cpu() / len(self.dataloader), prefix="val")}
+            return {k: round(float(v), 5) for k, v in results.items()}  # return results as 5 decimal place floats
         else:
             self.logger.info('Speed: %.1fms pre-process, %.1fms inference, %.1fms loss, %.1fms post-process per image' %
                              self.speed)