diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 89b996d..55a1413 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -18,7 +18,7 @@ jobs: fail-fast: false matrix: os: [ubuntu-latest] - python-version: ['3.10'] # requires python<=3.9 + python-version: ['3.10'] model: [yolov8n] steps: - uses: actions/checkout@v3 @@ -51,17 +51,17 @@ jobs: shell: python run: | from ultralytics.yolo.utils.benchmarks import run_benchmarks - run_benchmarks(model='yolov8n.pt', imgsz=160, half=False, hard_fail=False) + run_benchmarks(model='${{ matrix.model }}.pt', imgsz=160, half=False, hard_fail=False) - name: Benchmark SegmentationModel shell: python run: | from ultralytics.yolo.utils.benchmarks import run_benchmarks - run_benchmarks(model='yolov8n-seg.pt', imgsz=160, half=False, hard_fail=False) + run_benchmarks(model='${{ matrix.model }}-seg.pt', imgsz=160, half=False, hard_fail=False) - name: Benchmark ClassificationModel shell: python run: | from ultralytics.yolo.utils.benchmarks import run_benchmarks - run_benchmarks(model='yolov8n-cls.pt', imgsz=160, half=False, hard_fail=False) + run_benchmarks(model='${{ matrix.model }}-cls.pt', imgsz=160, half=False, hard_fail=False) Tests: timeout-minutes: 60 @@ -70,13 +70,10 @@ jobs: fail-fast: false matrix: os: [ubuntu-latest] - python-version: ['3.10'] + python-version: ['3.7', '3.8', '3.9', '3.10'] model: [yolov8n] torch: [latest] include: - - os: ubuntu-latest - python-version: '3.7' # '3.6.8' min - model: yolov8n - os: ubuntu-latest python-version: '3.8' # torch 1.7.0 requires python >=3.6, <=3.8 model: yolov8n @@ -123,9 +120,7 @@ jobs: run: | import os import ultralytics - from ultralytics import hub, yolo key = os.environ['APIKEY'] - print(ultralytics.__version__) ultralytics.checks() # ultralytics.reset_model(key) # reset trained model # ultralytics.start(key) # train model diff --git a/docs/cfg.md b/docs/cfg.md index 1b73abb..5197f0b 100644 --- a/docs/cfg.md +++ b/docs/cfg.md @@ -130,6 +130,7 @@ given task. | half | False | use half precision (FP16) | | device | null | device to run on, i.e. cuda device=0/1/2/3 or device=cpu | | show | False | show results if possible | +| save | False | save images with results | | save_txt | False | save results as .txt file | | save_conf | False | save results with confidence scores | | save_crop | False | save cropped images with results | diff --git a/setup.cfg b/setup.cfg index d7c4cb3..fcc17e1 100644 --- a/setup.cfg +++ b/setup.cfg @@ -3,7 +3,7 @@ # Local usage: pip install pre-commit, pre-commit run --all-files [metadata] -license_file = LICENSE +license_files = LICENSE description_file = README.md [tool:pytest] diff --git a/tests/test_engine.py b/tests/test_engine.py index 52c68c0..e842bdd 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -59,7 +59,7 @@ def test_segment(): # Predictor pred = segment.SegmentationPredictor(overrides={'imgsz': [64, 64]}) result = pred(source=SOURCE, model=f'{MODEL}-seg.pt') - assert len(result) == 2, 'predictor test failed' + assert len(result), 'predictor test failed' # Test resume overrides['resume'] = trainer.last @@ -97,4 +97,4 @@ def test_classify(): # Predictor pred = classify.ClassificationPredictor(overrides={'imgsz': [64, 64]}) result = pred(source=SOURCE, model=trainer.best) - assert len(result) == 2, 'predictor test failed' + assert len(result), 'predictor test failed' diff --git a/tests/test_python.py b/tests/test_python.py index 86425ae..34b00de 100644 --- a/tests/test_python.py +++ b/tests/test_python.py @@ -14,6 +14,13 @@ from ultralytics.yolo.utils import LINUX, ROOT, SETTINGS MODEL = Path(SETTINGS['weights_dir']) / 'yolov8n.pt' CFG = 'yolov8n.yaml' SOURCE = ROOT / 'assets/bus.jpg' +SOURCE_GREYSCALE = Path(f'{SOURCE.parent / SOURCE.stem}_greyscale.jpg') +SOURCE_RGBA = Path(f'{SOURCE.parent / SOURCE.stem}_4ch.png') + +# Convert SOURCE to greyscale and 4-ch +im = Image.open(SOURCE) +im.convert('L').save(SOURCE_GREYSCALE) # greyscale +im.convert('RGBA').save(SOURCE_RGBA) # 4-ch PNG with alpha def test_model_forward(): @@ -42,8 +49,7 @@ def test_predict_dir(): def test_predict_img(): model = YOLO(MODEL) - img = Image.open(str(SOURCE)) - output = model(source=img, save=True, verbose=True) # PIL + output = model(source=Image.open(SOURCE), save=True, verbose=True) # PIL assert len(output) == 1, 'predict test failed' img = cv2.imread(str(SOURCE)) output = model(source=img, save=True, save_txt=True) # ndarray @@ -67,6 +73,13 @@ def test_predict_img(): assert len(output) == 6, 'predict test failed!' +def test_predict_grey_and_4ch(): + model = YOLO(MODEL) + for f in SOURCE_RGBA, SOURCE_GREYSCALE: + for source in Image.open(f), cv2.imread(str(f)), f: + model(source, save=True, verbose=True) + + def test_val(): model = YOLO(MODEL) model.val(data='coco8.yaml', imgsz=32) @@ -151,6 +164,7 @@ def test_predict_callback_and_setup(): # results -> List[batch_size] path, _, im0s, _, _ = predictor.batch # print('on_predict_batch_end', im0s[0].shape) + im0s = im0s if isinstance(im0s, list) else [im0s] bs = [predictor.dataset.bs for _ in range(len(path))] predictor.results = zip(predictor.results, im0s, bs) diff --git a/ultralytics/__init__.py b/ultralytics/__init__.py index 6102be4..cefc173 100644 --- a/ultralytics/__init__.py +++ b/ultralytics/__init__.py @@ -1,6 +1,6 @@ # Ultralytics YOLO 🚀, GPL-3.0 license -__version__ = '8.0.42' +__version__ = '8.0.43' from ultralytics.yolo.engine.model import YOLO from ultralytics.yolo.utils.checks import check_yolo as checks diff --git a/ultralytics/tracker/__init__.py b/ultralytics/tracker/__init__.py index da0f3ab..7f4e335 100644 --- a/ultralytics/tracker/__init__.py +++ b/ultralytics/tracker/__init__.py @@ -1,3 +1,4 @@ # Ultralytics YOLO 🚀, GPL-3.0 license +from .track import register_tracker from .trackers import BOTSORT, BYTETracker diff --git a/ultralytics/tracker/track.py b/ultralytics/tracker/track.py index 9fd8a75..0a83e21 100644 --- a/ultralytics/tracker/track.py +++ b/ultralytics/tracker/track.py @@ -1,13 +1,16 @@ # Ultralytics YOLO 🚀, GPL-3.0 license +from ultralytics.yolo.utils.checks import check_requirements, check_yaml + +check_requirements('lap') # for linear_assignment + import torch -from ultralytics.tracker import BOTSORT, BYTETracker from ultralytics.yolo.utils import IterableSimpleNamespace, yaml_load -from ultralytics.yolo.utils.checks import check_requirements, check_yaml + +from .trackers import BOTSORT, BYTETracker TRACKER_MAP = {'bytetrack': BYTETracker, 'botsort': BOTSORT} -check_requirements('lap') # for linear_assignment def on_predict_start(predictor): diff --git a/ultralytics/yolo/cfg/__init__.py b/ultralytics/yolo/cfg/__init__.py index 2004966..3e6bc60 100644 --- a/ultralytics/yolo/cfg/__init__.py +++ b/ultralytics/yolo/cfg/__init__.py @@ -18,7 +18,7 @@ CLI_HELP_MSG = \ yolo TASK MODE ARGS Where TASK (optional) is one of [detect, segment, classify] - MODE (required) is one of [train, val, predict, export] + MODE (required) is one of [train, val, predict, export, track] ARGS (optional) are any number of custom 'arg=value' pairs like 'imgsz=320' that override defaults. See all ARGS at https://docs.ultralytics.com/cfg or with 'yolo cfg' @@ -197,7 +197,7 @@ def entrypoint(debug=''): # Define tasks and modes tasks = 'detect', 'segment', 'classify' - modes = 'train', 'val', 'predict', 'export', 'track' + modes = 'train', 'val', 'predict', 'export', 'track', 'benchmark' special = { 'help': lambda: LOGGER.info(CLI_HELP_MSG), 'checks': checks.check_yolo, diff --git a/ultralytics/yolo/data/dataloaders/stream_loaders.py b/ultralytics/yolo/data/dataloaders/stream_loaders.py index 106409f..613850a 100644 --- a/ultralytics/yolo/data/dataloaders/stream_loaders.py +++ b/ultralytics/yolo/data/dataloaders/stream_loaders.py @@ -290,13 +290,15 @@ class LoadPilAndNumpy: self.transforms = transforms self.mode = 'image' # generate fake paths - self.paths = [f'image{i}.jpg' for i in range(len(self.im0))] + self.paths = [getattr(im, 'filename', f'image{i}.jpg') for i, im in enumerate(self.im0)] self.bs = len(self.im0) @staticmethod def _single_check(im): assert isinstance(im, (Image.Image, np.ndarray)), f'Expected PIL/np.ndarray image type, but got {type(im)}' if isinstance(im, Image.Image): + if im.mode != 'RGB': + im = im.convert('RGB') im = np.asarray(im)[:, :, ::-1] im = np.ascontiguousarray(im) # contiguous return im diff --git a/ultralytics/yolo/data/dataloaders/v5loader.py b/ultralytics/yolo/data/dataloaders/v5loader.py index b7eb7f4..8ea4582 100644 --- a/ultralytics/yolo/data/dataloaders/v5loader.py +++ b/ultralytics/yolo/data/dataloaders/v5loader.py @@ -1045,7 +1045,7 @@ class HUBDatasetStats(): autodownload: Attempt to download dataset if not found locally Usage - from utils.dataloaders import HUBDatasetStats + from ultralytics.yolo.data.dataloaders.v5loader import HUBDatasetStats stats = HUBDatasetStats('coco128.yaml', autodownload=True) # usage 1 stats = HUBDatasetStats('path/to/coco128.zip') # usage 2 stats.get_json(save=False) @@ -1055,15 +1055,15 @@ class HUBDatasetStats(): def __init__(self, path='coco128.yaml', autodownload=False): # Initialize class zipped, data_dir, yaml_path = self._unzip(Path(path)) - try: - data = yaml_load(check_yaml(yaml_path)) # data dict - if zipped: - data['path'] = data_dir - except Exception as e: - raise Exception('error/HUB/dataset_stats/yaml_load') from e - - check_det_dataset(data, autodownload) # download dataset if missing - self.hub_dir = Path(data['path'] + '-hub') + # try: + # data = yaml_load(check_yaml(yaml_path)) # data dict + # if zipped: + # data['path'] = data_dir + # except Exception as e: + # raise Exception('error/HUB/dataset_stats/yaml_load') from e + + data = check_det_dataset(yaml_path, autodownload) # download dataset if missing + self.hub_dir = Path(str(data['path']) + '-hub') self.im_dir = self.hub_dir / 'images' self.im_dir.mkdir(parents=True, exist_ok=True) # makes /images self.stats = {'nc': data['nc'], 'names': list(data['names'].values())} # statistics dictionary diff --git a/ultralytics/yolo/engine/model.py b/ultralytics/yolo/engine/model.py index e468180..33135c2 100644 --- a/ultralytics/yolo/engine/model.py +++ b/ultralytics/yolo/engine/model.py @@ -9,7 +9,7 @@ from ultralytics.nn.tasks import (ClassificationModel, DetectionModel, Segmentat guess_model_task, nn) from ultralytics.yolo.cfg import get_cfg from ultralytics.yolo.engine.exporter import Exporter -from ultralytics.yolo.utils import DEFAULT_CFG, LOGGER, RANK, callbacks, yaml_load +from ultralytics.yolo.utils import DEFAULT_CFG, DEFAULT_CFG_DICT, LOGGER, RANK, callbacks, yaml_load from ultralytics.yolo.utils.checks import check_file, check_imgsz, check_yaml from ultralytics.yolo.utils.downloads import GITHUB_ASSET_STEMS from ultralytics.yolo.utils.torch_utils import smart_inference_mode @@ -203,7 +203,7 @@ class YOLO: @smart_inference_mode() def track(self, source=None, stream=False, **kwargs): - from ultralytics.tracker.track import register_tracker + from ultralytics.tracker import register_tracker register_tracker(self) # ByteTrack-based method needs low confidence predictions as input conf = kwargs.get('conf') or 0.1 @@ -237,6 +237,20 @@ class YOLO: return validator.metrics + @smart_inference_mode() + def benchmark(self, **kwargs): + """ + Benchmark a model on all export formats. + + Args: + **kwargs : Any other args accepted by the validators. To see all args check 'configuration' section in docs + """ + from ultralytics.yolo.utils.benchmarks import run_benchmarks + overrides = self.model.args.copy() + overrides.update(kwargs) + overrides = {**DEFAULT_CFG_DICT, **overrides} # fill in missing overrides keys with defaults + return run_benchmarks(model=self, imgsz=overrides['imgsz'], half=overrides['half'], device=overrides['device']) + def export(self, **kwargs): """ Export model. diff --git a/ultralytics/yolo/engine/predictor.py b/ultralytics/yolo/engine/predictor.py index 132a504..f318cac 100644 --- a/ultralytics/yolo/engine/predictor.py +++ b/ultralytics/yolo/engine/predictor.py @@ -194,7 +194,7 @@ class BasePredictor: # Print time (inference-only) if self.args.verbose: - LOGGER.info(f"{s}{'' if len(preds) else '(no detections), '}{self.dt[1].dt * 1E3:.1f}ms") + LOGGER.info(f'{s}{self.dt[1].dt * 1E3:.1f}ms') # Release assets if isinstance(self.vid_writer[-1], cv2.VideoWriter): diff --git a/ultralytics/yolo/engine/results.py b/ultralytics/yolo/engine/results.py index 9e7a3fd..f60f2cc 100644 --- a/ultralytics/yolo/engine/results.py +++ b/ultralytics/yolo/engine/results.py @@ -1,3 +1,10 @@ +# Ultralytics YOLO 🚀, GPL-3.0 license +""" +Ultralytics Results, Boxes and Masks classes for handling inference results + +Usage: See https://docs.ultralytics.com/predict/ +""" + from copy import deepcopy from functools import lru_cache @@ -36,7 +43,7 @@ class Results: self.probs = probs if probs is not None else None self.names = names self.path = path - self.comp = ['boxes', 'masks', 'probs'] + self._keys = (k for k in ('boxes', 'masks', 'probs') if getattr(self, k) is not None) def pandas(self): pass @@ -44,10 +51,8 @@ class Results: def __getitem__(self, idx): r = Results(orig_img=self.orig_img, path=self.path, names=self.names) - for item in self.comp: - if getattr(self, item) is None: - continue - setattr(r, item, getattr(self, item)[idx]) + for k in self._keys: + setattr(r, k, getattr(self, k)[idx]) return r def update(self, boxes=None, masks=None, probs=None): @@ -60,57 +65,37 @@ class Results: def cpu(self): r = Results(orig_img=self.orig_img, path=self.path, names=self.names) - for item in self.comp: - if getattr(self, item) is None: - continue - setattr(r, item, getattr(self, item).cpu()) + for k in self._keys: + setattr(r, k, getattr(self, k).cpu()) return r def numpy(self): r = Results(orig_img=self.orig_img, path=self.path, names=self.names) - for item in self.comp: - if getattr(self, item) is None: - continue - setattr(r, item, getattr(self, item).numpy()) + for k in self._keys: + setattr(r, k, getattr(self, k).numpy()) return r def cuda(self): r = Results(orig_img=self.orig_img, path=self.path, names=self.names) - for item in self.comp: - if getattr(self, item) is None: - continue - setattr(r, item, getattr(self, item).cuda()) + for k in self._keys: + setattr(r, k, getattr(self, k).cuda()) return r def to(self, *args, **kwargs): r = Results(orig_img=self.orig_img, path=self.path, names=self.names) - for item in self.comp: - if getattr(self, item) is None: - continue - setattr(r, item, getattr(self, item).to(*args, **kwargs)) + for k in self._keys: + setattr(r, k, getattr(self, k).to(*args, **kwargs)) return r def __len__(self): - for item in self.comp: - if getattr(self, item) is None: - continue - return len(getattr(self, item)) + for k in self._keys: + return len(getattr(self, k)) def __str__(self): - str_out = '' - for item in self.comp: - if getattr(self, item) is None: - continue - str_out = str_out + getattr(self, item).__str__() - return str_out + return ''.join(getattr(self, k).__str__() for k in self._keys) def __repr__(self): - str_out = '' - for item in self.comp: - if getattr(self, item) is None: - continue - str_out = str_out + getattr(self, item).__repr__() - return str_out + return ''.join(getattr(self, k).__repr__() for k in self._keys) def __getattr__(self, attr): name = self.__class__.__name__ @@ -226,20 +211,16 @@ class Boxes: return self.xywh / self.orig_shape[[1, 0, 1, 0]] def cpu(self): - boxes = self.boxes.cpu() - return Boxes(boxes, self.orig_shape) + return Boxes(self.boxes.cpu(), self.orig_shape) def numpy(self): - boxes = self.boxes.numpy() - return Boxes(boxes, self.orig_shape) + return Boxes(self.boxes.numpy(), self.orig_shape) def cuda(self): - boxes = self.boxes.cuda() - return Boxes(boxes, self.orig_shape) + return Boxes(self.boxes.cuda(), self.orig_shape) def to(self, *args, **kwargs): - boxes = self.boxes.to(*args, **kwargs) - return Boxes(boxes, self.orig_shape) + return Boxes(self.boxes.to(*args, **kwargs), self.orig_shape) def pandas(self): LOGGER.info('results.pandas() method not yet implemented') @@ -272,8 +253,7 @@ class Boxes: f'shape: {self.boxes.shape}\n' + f'dtype: {self.boxes.dtype}\n + {self.boxes.__repr__()}') def __getitem__(self, idx): - boxes = self.boxes[idx] - return Boxes(boxes, self.orig_shape) + return Boxes(self.boxes[idx], self.orig_shape) def __getattr__(self, attr): name = self.__class__.__name__ @@ -331,20 +311,16 @@ class Masks: return self.masks def cpu(self): - masks = self.masks.cpu() - return Masks(masks, self.orig_shape) + return Masks(self.masks.cpu(), self.orig_shape) def numpy(self): - masks = self.masks.numpy() - return Masks(masks, self.orig_shape) + return Masks(self.masks.numpy(), self.orig_shape) def cuda(self): - masks = self.masks.cuda() - return Masks(masks, self.orig_shape) + return Masks(self.masks.cuda(), self.orig_shape) def to(self, *args, **kwargs): - masks = self.masks.to(*args, **kwargs) - return Masks(masks, self.orig_shape) + return Masks(self.masks.to(*args, **kwargs), self.orig_shape) def __len__(self): # override len(results) return len(self.masks) @@ -357,8 +333,7 @@ class Masks: f'shape: {self.masks.shape}\n' + f'dtype: {self.masks.dtype}\n + {self.masks.__repr__()}') def __getitem__(self, idx): - masks = self.masks[idx] - return Masks(masks, self.orig_shape) + return Masks(self.masks[idx], self.orig_shape) def __getattr__(self, attr): name = self.__class__.__name__ diff --git a/ultralytics/yolo/engine/trainer.py b/ultralytics/yolo/engine/trainer.py index 4fc09ad..9573471 100644 --- a/ultralytics/yolo/engine/trainer.py +++ b/ultralytics/yolo/engine/trainer.py @@ -243,6 +243,8 @@ class BaseTrainer: metric_keys = self.validator.metrics.keys + self.label_loss_items(prefix='val') self.metrics = dict(zip(metric_keys, [0] * len(metric_keys))) # TODO: init metrics for plot_results()? self.ema = ModelEMA(self.model) + if self.args.plots: + self.plot_training_labels() self.resume_training(ckpt) self.scheduler.last_epoch = self.start_epoch - 1 # do not move self.run_callbacks('on_pretrain_routine_end') @@ -501,6 +503,9 @@ class BaseTrainer: def plot_training_samples(self, batch, ni): pass + def plot_training_labels(self): + pass + def save_metrics(self, metrics): keys, vals = list(metrics.keys()), list(metrics.values()) n = len(metrics) + 1 # number of cols diff --git a/ultralytics/yolo/engine/validator.py b/ultralytics/yolo/engine/validator.py index ae29b04..b2a5394 100644 --- a/ultralytics/yolo/engine/validator.py +++ b/ultralytics/yolo/engine/validator.py @@ -28,7 +28,7 @@ from tqdm import tqdm from ultralytics.nn.autobackend import AutoBackend from ultralytics.yolo.cfg import get_cfg from ultralytics.yolo.data.utils import check_cls_dataset, check_det_dataset -from ultralytics.yolo.utils import DEFAULT_CFG, LOGGER, RANK, SETTINGS, TQDM_BAR_FORMAT, callbacks, emojis +from ultralytics.yolo.utils import DEFAULT_CFG, LOGGER, RANK, SETTINGS, TQDM_BAR_FORMAT, callbacks, colorstr, emojis from ultralytics.yolo.utils.checks import check_imgsz from ultralytics.yolo.utils.files import increment_path from ultralytics.yolo.utils.ops import Profile @@ -194,6 +194,8 @@ class BaseValidator: self.logger.info(f'Saving {f.name}...') json.dump(self.jdict, f) # flatten and save stats = self.eval_json(stats) # update stats + if self.args.plots or self.args.save_json: + LOGGER.info(f"Results saved to {colorstr('bold', self.save_dir)}") return stats def run_callbacks(self, event: str): diff --git a/ultralytics/yolo/utils/benchmarks.py b/ultralytics/yolo/utils/benchmarks.py index bfbee76..eb0c440 100644 --- a/ultralytics/yolo/utils/benchmarks.py +++ b/ultralytics/yolo/utils/benchmarks.py @@ -27,13 +27,14 @@ import time from pathlib import Path import pandas as pd -import torch from ultralytics import YOLO from ultralytics.yolo.engine.exporter import export_formats -from ultralytics.yolo.utils import LOGGER, SETTINGS +from ultralytics.yolo.utils import LOGGER, ROOT, SETTINGS from ultralytics.yolo.utils.checks import check_yolo +from ultralytics.yolo.utils.downloads import download from ultralytics.yolo.utils.files import file_size +from ultralytics.yolo.utils.torch_utils import select_device def run_benchmarks(model=Path(SETTINGS['weights_dir']) / 'yolov8n.pt', @@ -41,8 +42,9 @@ def run_benchmarks(model=Path(SETTINGS['weights_dir']) / 'yolov8n.pt', half=False, device='cpu', hard_fail=False): - device = torch.device(int(device) if device.isnumeric() else device) - model = YOLO(model) + device = select_device(device, verbose=False) + if isinstance(model, (str, Path)): + model = YOLO(model) y = [] t0 = time.time() @@ -65,6 +67,11 @@ def run_benchmarks(model=Path(SETTINGS['weights_dir']) / 'yolov8n.pt', export = YOLO(filename) assert suffix in str(filename), 'export failed' + # Predict + if not (ROOT / 'assets/bus.jpg').exists(): + download(url='https://ultralytics.com/images/bus.jpg', dir=ROOT / 'assets') + export.predict(ROOT / 'assets/bus.jpg', imgsz=imgsz, device=device, half=half) # test + # Validate if model.task == 'detect': data, key = 'coco128.yaml', 'metrics/mAP50-95(B)' @@ -96,6 +103,8 @@ def run_benchmarks(model=Path(SETTINGS['weights_dir']) / 'yolov8n.pt', floor = eval(hard_fail) # minimum metric floor to pass, i.e. = 0.29 mAP for YOLOv5n assert all(x > floor for x in metrics if pd.notna(x)), f'HARD FAIL: metric < floor {floor}' + return df + if __name__ == '__main__': run_benchmarks() diff --git a/ultralytics/yolo/utils/plotting.py b/ultralytics/yolo/utils/plotting.py index 5d5da9f..c3b9c79 100644 --- a/ultralytics/yolo/utils/plotting.py +++ b/ultralytics/yolo/utils/plotting.py @@ -5,19 +5,24 @@ import math from pathlib import Path import cv2 +import matplotlib import matplotlib.pyplot as plt import numpy as np import pandas as pd +import seaborn as sn import torch from PIL import Image, ImageDraw, ImageFont from PIL import __version__ as pil_version -from ultralytics.yolo.utils import LOGGER, threaded +from ultralytics.yolo.utils import LOGGER, TryExcept, threaded from .checks import check_font, check_version, is_ascii from .files import increment_path from .ops import clip_coords, scale_image, xywh2xyxy, xyxy2xywh +matplotlib.rc('font', **{'size': 11}) +matplotlib.use('Agg') # for writing to files only + class Colors: # Ultralytics color palette https://ultralytics.com/ @@ -152,6 +157,52 @@ class Annotator: return np.asarray(self.im) +@TryExcept() # known issue https://github.com/ultralytics/yolov5/issues/5395 +def plot_labels(boxes, cls, names=(), save_dir=Path('')): + # plot dataset labels + LOGGER.info(f"Plotting labels to {save_dir / 'labels.jpg'}... ") + b = boxes.transpose() # classes, boxes + nc = int(cls.max() + 1) # number of classes + x = pd.DataFrame(b.transpose(), columns=['x', 'y', 'width', 'height']) + + # seaborn correlogram + sn.pairplot(x, corner=True, diag_kind='auto', kind='hist', diag_kws=dict(bins=50), plot_kws=dict(pmax=0.9)) + plt.savefig(save_dir / 'labels_correlogram.jpg', dpi=200) + plt.close() + + # matplotlib labels + matplotlib.use('svg') # faster + ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)[1].ravel() + y = ax[0].hist(cls, bins=np.linspace(0, nc, nc + 1) - 0.5, rwidth=0.8) + with contextlib.suppress(Exception): # color histogram bars by class + [y[2].patches[i].set_color([x / 255 for x in colors(i)]) for i in range(nc)] # known issue #3195 + ax[0].set_ylabel('instances') + if 0 < len(names) < 30: + ax[0].set_xticks(range(len(names))) + ax[0].set_xticklabels(list(names.values()), rotation=90, fontsize=10) + else: + ax[0].set_xlabel('classes') + sn.histplot(x, x='x', y='y', ax=ax[2], bins=50, pmax=0.9) + sn.histplot(x, x='width', y='height', ax=ax[3], bins=50, pmax=0.9) + + # rectangles + boxes[:, 0:2] = 0.5 # center + boxes = xywh2xyxy(boxes) * 2000 + img = Image.fromarray(np.ones((2000, 2000, 3), dtype=np.uint8) * 255) + for cls, box in zip(cls[:1000], boxes[:1000]): + ImageDraw.Draw(img).rectangle(box, width=1, outline=colors(cls)) # plot + ax[1].imshow(img) + ax[1].axis('off') + + for a in [0, 1, 2, 3]: + for s in ['top', 'right', 'left', 'bottom']: + ax[a].spines[s].set_visible(False) + + plt.savefig(save_dir / 'labels.jpg', dpi=200) + matplotlib.use('Agg') + plt.close() + + def save_one_box(xyxy, im, file=Path('im.jpg'), gain=1.02, pad=10, square=False, BGR=False, save=True): # Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop xyxy = torch.Tensor(xyxy).view(-1, 4) diff --git a/ultralytics/yolo/utils/torch_utils.py b/ultralytics/yolo/utils/torch_utils.py index 34ddf0d..84cb5ec 100644 --- a/ultralytics/yolo/utils/torch_utils.py +++ b/ultralytics/yolo/utils/torch_utils.py @@ -59,7 +59,7 @@ def DDP_model(model): return DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK) -def select_device(device='', batch=0, newline=False): +def select_device(device='', batch=0, newline=False, verbose=True): # device = None or 'cpu' or 0 or '0' or '0,1,2,3' s = f'Ultralytics YOLOv{__version__} 🚀 Python-{platform.python_version()} torch-{torch.__version__} ' device = str(device).lower() @@ -102,7 +102,7 @@ def select_device(device='', batch=0, newline=False): s += 'CPU\n' arg = 'cpu' - if RANK == -1: + if verbose and RANK == -1: LOGGER.info(s if newline else s.rstrip()) return torch.device(arg) diff --git a/ultralytics/yolo/v8/detect/predict.py b/ultralytics/yolo/v8/detect/predict.py index f26a20a..ecd5169 100644 --- a/ultralytics/yolo/v8/detect/predict.py +++ b/ultralytics/yolo/v8/detect/predict.py @@ -56,7 +56,7 @@ class DetectionPredictor(BasePredictor): det = results[idx].boxes # TODO: make boxes inherit from tensors if len(det) == 0: - return log_string + return f'{log_string}(no detections), ' for c in det.cls.unique(): n = (det.cls == c).sum() # detections per class log_string += f"{n} {self.model.names[int(c)]}{'s' * (n > 1)}, " diff --git a/ultralytics/yolo/v8/detect/train.py b/ultralytics/yolo/v8/detect/train.py index 348502e..874d489 100644 --- a/ultralytics/yolo/v8/detect/train.py +++ b/ultralytics/yolo/v8/detect/train.py @@ -1,6 +1,7 @@ # Ultralytics YOLO 🚀, GPL-3.0 license from copy import copy +import numpy as np import torch import torch.nn as nn @@ -12,7 +13,7 @@ from ultralytics.yolo.engine.trainer import BaseTrainer from ultralytics.yolo.utils import DEFAULT_CFG, RANK, colorstr from ultralytics.yolo.utils.loss import BboxLoss from ultralytics.yolo.utils.ops import xywh2xyxy -from ultralytics.yolo.utils.plotting import plot_images, plot_results +from ultralytics.yolo.utils.plotting import plot_images, plot_labels, plot_results from ultralytics.yolo.utils.tal import TaskAlignedAssigner, dist2bbox, make_anchors from ultralytics.yolo.utils.torch_utils import de_parallel @@ -102,6 +103,11 @@ class DetectionTrainer(BaseTrainer): def plot_metrics(self): plot_results(file=self.csv) # save results.png + def plot_training_labels(self): + boxes = np.concatenate([lb['bboxes'] for lb in self.train_loader.dataset.labels], 0) + cls = np.concatenate([lb['cls'] for lb in self.train_loader.dataset.labels], 0) + plot_labels(boxes, cls.squeeze(), names=self.data['names'], save_dir=self.save_dir) + # Criterion class for computing training losses class Loss: diff --git a/ultralytics/yolo/v8/segment/predict.py b/ultralytics/yolo/v8/segment/predict.py index 24bc9d6..84adde7 100644 --- a/ultralytics/yolo/v8/segment/predict.py +++ b/ultralytics/yolo/v8/segment/predict.py @@ -59,7 +59,7 @@ class SegmentationPredictor(DetectionPredictor): result = results[idx] if len(result) == 0: - return log_string + return f'{log_string}(no detections), ' det, mask = result.boxes, result.masks # getting tensors TODO: mask mask,box inherit for tensor # Print results