ultralytics 8.0.73 minor fixes (#1929)

Co-authored-by: Yonghye Kwon <developer.0hye@gmail.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: joseliraGB <122470533+joseliraGB@users.noreply.github.com>

ultralytics/yolo/engine/model.py

@@ -356,7 +356,6 @@ class YOLO:
             raise AttributeError("Dataset required but missing, i.e. pass 'data=coco128.yaml'")
         if overrides.get('resume'):
             overrides['resume'] = self.ckpt_path
-
         self.task = overrides.get('task') or self.task
         self.trainer = TASK_MAP[self.task][1](overrides=overrides, _callbacks=self.callbacks)
         if not overrides.get('resume'):  # manually set model only if not resuming

ultralytics/yolo/engine/predictor.py

@@ -109,8 +109,35 @@ class BasePredictor:
     def preprocess(self, img):
         pass
 
-    def write_results(self, results, batch, print_string):
-        raise NotImplementedError('print_results function needs to be implemented')
+    def write_results(self, idx, results, batch):
+        p, im, _ = batch
+        log_string = ''
+        if len(im.shape) == 3:
+            im = im[None]  # expand for batch dim
+        self.seen += 1
+        if self.source_type.webcam or self.source_type.from_img:  # batch_size >= 1
+            log_string += f'{idx}: '
+            frame = self.dataset.count
+        else:
+            frame = getattr(self.dataset, 'frame', 0)
+        self.data_path = p
+        self.txt_path = str(self.save_dir / 'labels' / p.stem) + ('' if self.dataset.mode == 'image' else f'_{frame}')
+        log_string += '%gx%g ' % im.shape[2:]  # print string
+        result = results[idx]
+        log_string += result.verbose()
+
+        if self.args.save or self.args.show:  # Add bbox to image
+            plot_args = dict(line_width=self.args.line_thickness, boxes=self.args.boxes)
+            if not self.args.retina_masks:
+                plot_args['im_gpu'] = im[idx]
+            self.plotted_img = result.plot(**plot_args)
+        # write
+        if self.args.save_txt:
+            result.save_txt(f'{self.txt_path}.txt', save_conf=self.args.save_conf)
+        if self.args.save_crop:
+            result.save_crop(save_dir=self.save_dir / 'crops', file_name=self.data_path.stem)
+
+        return log_string
 
     def postprocess(self, preds, img, orig_img):
         return preds

ultralytics/yolo/engine/results.py

@@ -7,13 +7,14 @@ Usage: See https://docs.ultralytics.com/modes/predict/
 
 from copy import deepcopy
 from functools import lru_cache
+from pathlib import Path
 
 import numpy as np
 import torch
 
 from ultralytics.yolo.data.augment import LetterBox
 from ultralytics.yolo.utils import LOGGER, SimpleClass, deprecation_warn, ops
-from ultralytics.yolo.utils.plotting import Annotator, colors
+from ultralytics.yolo.utils.plotting import Annotator, colors, save_one_box
 
 
 class BaseTensor(SimpleClass):
@@ -233,6 +234,80 @@ class Results(SimpleClass):
 
         return annotator.result()
 
+    def verbose(self):
+        """
+        Return log string for each tasks.
+        """
+        log_string = ''
+        probs = self.probs
+        boxes = self.boxes
+        if len(self) == 0:
+            return log_string if probs is not None else log_string + '(no detections), '
+        if probs is not None:
+            n5 = min(len(self.names), 5)
+            top5i = probs.argsort(0, descending=True)[:n5].tolist()  # top 5 indices
+            log_string += f"{', '.join(f'{self.names[j]} {probs[j]:.2f}' for j in top5i)}, "
+        if boxes:
+            for c in boxes.cls.unique():
+                n = (boxes.cls == c).sum()  # detections per class
+                log_string += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, "
+        return log_string
+
+    def save_txt(self, txt_file, save_conf=False):
+        """Save predictions into txt file.
+
+        Args:
+            txt_file (str): txt file path.
+            save_conf (bool): save confidence score or not.
+        """
+        boxes = self.boxes
+        masks = self.masks
+        probs = self.probs
+        kpts = self.keypoints
+        texts = []
+        if probs is not None:
+            # classify
+            n5 = min(len(self.names), 5)
+            top5i = probs.argsort(0, descending=True)[:n5].tolist()  # top 5 indices
+            [texts.append(f'{probs[j]:.2f} {self.names[j]}') for j in top5i]
+        elif boxes:
+            # detect/segment/pose
+            for j, d in enumerate(boxes):
+                c, conf, id = int(d.cls), float(d.conf), None if d.id is None else int(d.id.item())
+                line = (c, *d.xywhn.view(-1))
+                if masks:
+                    seg = masks[j].xyn[0].copy().reshape(-1)  # reversed mask.xyn, (n,2) to (n*2)
+                    line = (c, *seg)
+                if kpts is not None:
+                    kpt = (kpts[j][:, :2] / d.orig_shape[[1, 0]]).reshape(-1).tolist()
+                    line += (*kpt, )
+                line += (conf, ) * save_conf + (() if id is None else (id, ))
+                texts.append(('%g ' * len(line)).rstrip() % line)
+
+        with open(txt_file, 'a') as f:
+            for text in texts:
+                f.write(text + '\n')
+
+    def save_crop(self, save_dir, file_name=Path('im.jpg')):
+        """Save cropped predictions to `save_dir/cls/file_name.jpg`.
+
+        Args:
+            save_dir (str | pathlib.Path): Save path.
+            file_name (str | pathlib.Path): File name.
+        """
+        if self.probs is not None:
+            LOGGER.warning('Warning: Classify task do not support `save_crop`.')
+            return
+        if isinstance(save_dir, str):
+            save_dir = Path(save_dir)
+        if isinstance(file_name, str):
+            file_name = Path(file_name)
+        for d in self.boxes:
+            save_one_box(d.xyxy,
+                         self.orig_img.copy(),
+                         file=save_dir / self.names[int(d.cls)] / f'{file_name.stem}.jpg',
+                         BGR=True)
+
 
 class Boxes(BaseTensor):
     """
@@ -339,6 +414,8 @@ class Masks(BaseTensor):
     """
 
     def __init__(self, masks, orig_shape) -> None:
+        if masks.ndim == 2:
+            masks = masks[None, :]
         self.masks = masks  # N, h, w
         self.orig_shape = orig_shape
 

ultralytics/yolo/engine/trainer.py

@@ -552,7 +552,7 @@ class BaseTrainer:
         if self.resume:
             assert start_epoch > 0, \
                 f'{self.args.model} training to {self.epochs} epochs is finished, nothing to resume.\n' \
-                f"Start a new training without --resume, i.e. 'yolo task=... mode=train model={self.args.model}'"
+                f"Start a new training without resuming, i.e. 'yolo train model={self.args.model}'"
             LOGGER.info(
                 f'Resuming training from {self.args.model} from epoch {start_epoch + 1} to {self.epochs} total epochs')
         if self.epochs < start_epoch: