mirror of
https://github.com/THU-MIG/yolov10.git
synced 2025-05-23 21:44:22 +08:00
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>
This commit is contained in:
parent
95f96dc5bc
commit
5629ed0bb7
4
.github/workflows/ci.yaml
vendored
4
.github/workflows/ci.yaml
vendored
@ -76,9 +76,9 @@ jobs:
|
|||||||
run: |
|
run: |
|
||||||
python -m pip install --upgrade pip wheel
|
python -m pip install --upgrade pip wheel
|
||||||
if [ "${{ matrix.os }}" == "macos-latest" ]; then
|
if [ "${{ matrix.os }}" == "macos-latest" ]; then
|
||||||
pip install -e . coremltools openvino-dev tensorflow-macos --extra-index-url https://download.pytorch.org/whl/cpu
|
pip install -e . coremltools openvino-dev tensorflow-macos tensorflowjs --extra-index-url https://download.pytorch.org/whl/cpu
|
||||||
else
|
else
|
||||||
pip install -e . coremltools openvino-dev tensorflow-cpu --extra-index-url https://download.pytorch.org/whl/cpu
|
pip install -e . coremltools openvino-dev tensorflow-cpu tensorflowjs --extra-index-url https://download.pytorch.org/whl/cpu
|
||||||
fi
|
fi
|
||||||
yolo export format=tflite
|
yolo export format=tflite
|
||||||
- name: Check environment
|
- name: Check environment
|
||||||
|
103
examples/hub.ipynb
Normal file
103
examples/hub.ipynb
Normal file
@ -0,0 +1,103 @@
|
|||||||
|
{
|
||||||
|
"nbformat": 4,
|
||||||
|
"nbformat_minor": 0,
|
||||||
|
"metadata": {
|
||||||
|
"colab": {
|
||||||
|
"name": "Ultralytics HUB",
|
||||||
|
"provenance": []
|
||||||
|
},
|
||||||
|
"kernelspec": {
|
||||||
|
"name": "python3",
|
||||||
|
"display_name": "Python 3"
|
||||||
|
},
|
||||||
|
"language_info": {
|
||||||
|
"name": "python"
|
||||||
|
},
|
||||||
|
"accelerator": "GPU"
|
||||||
|
},
|
||||||
|
"cells": [
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {
|
||||||
|
"id": "FIzICjaph_Wy"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"<a align=\"center\" href=\"https://hub.ultralytics.com\" target=\"_blank\">\n",
|
||||||
|
"<img width=\"1024\", src=\"https://github.com/ultralytics/assets/raw/main/im/ultralytics-hub.png\"></a>\n",
|
||||||
|
"\n",
|
||||||
|
"<div align=\"center\">\n",
|
||||||
|
" <a href=\"https://github.com/ultralytics/hub/actions/workflows/ci.yaml\">\n",
|
||||||
|
" <img src=\"https://github.com/ultralytics/hub/actions/workflows/ci.yaml/badge.svg\" alt=\"CI CPU\"></a>\n",
|
||||||
|
" <a href=\"https://colab.research.google.com/github/ultralytics/hub/blob/master/hub.ipynb\">\n",
|
||||||
|
" <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"></a>\n",
|
||||||
|
"\n",
|
||||||
|
"Welcome to the [Ultralytics](https://ultralytics.com/) HUB notebook! \n",
|
||||||
|
"\n",
|
||||||
|
"This notebook allows you to train [YOLOv5](https://github.com/ultralytics/yolov5) and [YOLOv8](https://github.com/ultralytics/ultralytics) 🚀 models using [HUB](https://hub.ultralytics.com/). Please browse the YOLOv8 <a href=\"https://docs.ultralytics.com\">Docs</a> for details, raise an issue on <a href=\"https://github.com/ultralytics/hub/issues/new/choose\">GitHub</a> for support, and join our <a href=\"https://discord.gg/n6cFeSPZdD\">Discord</a> community for questions and discussions!\n",
|
||||||
|
"</div>"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {
|
||||||
|
"id": "eRQ2ow94MiOv"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"# Setup\n",
|
||||||
|
"\n",
|
||||||
|
"Pip install `ultralytics` and [dependencies](https://github.com/ultralytics/ultralytics/blob/main/requirements.txt) and check software and hardware."
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"metadata": {
|
||||||
|
"id": "FyDnXd-n4c7Y",
|
||||||
|
"colab": {
|
||||||
|
"base_uri": "https://localhost:8080/"
|
||||||
|
},
|
||||||
|
"outputId": "22dcbc27-9c6f-44fb-9745-620431f93793"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"%pip install ultralytics # install\n",
|
||||||
|
"from ultralytics import YOLO, checks, hub\n",
|
||||||
|
"checks() # checks"
|
||||||
|
],
|
||||||
|
"execution_count": null,
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"output_type": "stream",
|
||||||
|
"name": "stderr",
|
||||||
|
"text": [
|
||||||
|
"Ultralytics YOLOv8.0.64 🚀 Python-3.9.16 torch-2.0.0+cu118 CUDA:0 (Tesla T4, 15102MiB)\n",
|
||||||
|
"Setup complete ✅ (2 CPUs, 12.7 GB RAM, 28.3/166.8 GB disk)\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {
|
||||||
|
"id": "cQ9BwaAqxAm4"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"# Start\n",
|
||||||
|
"\n",
|
||||||
|
"Login with your [API key](https://hub.ultralytics.com/settings?tab=api+keys), select your YOLO 🚀 model and start training!"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"metadata": {
|
||||||
|
"id": "XSlZaJ9Iw_iZ"
|
||||||
|
},
|
||||||
|
"source": [
|
||||||
|
"hub.login('API_KEY') # use your API key\n",
|
||||||
|
"\n",
|
||||||
|
"model = YOLO('https://hub.ultralytics.com/MODEL_ID') # use your model URL\n",
|
||||||
|
"model.train() # train model"
|
||||||
|
],
|
||||||
|
"execution_count": null,
|
||||||
|
"outputs": []
|
||||||
|
}
|
||||||
|
]
|
||||||
|
}
|
@ -1,6 +1,6 @@
|
|||||||
# Ultralytics YOLO 🚀, GPL-3.0 license
|
# Ultralytics YOLO 🚀, GPL-3.0 license
|
||||||
|
|
||||||
__version__ = '8.0.72'
|
__version__ = '8.0.73'
|
||||||
|
|
||||||
from ultralytics.hub import start
|
from ultralytics.hub import start
|
||||||
from ultralytics.yolo.engine.model import YOLO
|
from ultralytics.yolo.engine.model import YOLO
|
||||||
|
@ -112,10 +112,8 @@ class HUBTrainingSession:
|
|||||||
raise ValueError('Dataset may still be processing. Please wait a minute and try again.') # RF fix
|
raise ValueError('Dataset may still be processing. Please wait a minute and try again.') # RF fix
|
||||||
self.model_id = data['id']
|
self.model_id = data['id']
|
||||||
|
|
||||||
# TODO: restore when server keys when dataset URL and GPU train is working
|
|
||||||
|
|
||||||
self.train_args = {
|
self.train_args = {
|
||||||
'batch': data['batch_size'],
|
'batch': data['batch' if ('batch' in data) else 'batch_size'], # TODO: deprecate 'batch_size' in 3Q23
|
||||||
'epochs': data['epochs'],
|
'epochs': data['epochs'],
|
||||||
'imgsz': data['imgsz'],
|
'imgsz': data['imgsz'],
|
||||||
'patience': data['patience'],
|
'patience': data['patience'],
|
||||||
|
@ -539,7 +539,7 @@ def guess_model_task(model):
|
|||||||
model (nn.Module) or (dict): PyTorch model or model configuration in YAML format.
|
model (nn.Module) or (dict): PyTorch model or model configuration in YAML format.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
str: Task of the model ('detect', 'segment', 'classify').
|
str: Task of the model ('detect', 'segment', 'classify', 'pose').
|
||||||
|
|
||||||
Raises:
|
Raises:
|
||||||
SyntaxError: If the task of the model could not be determined.
|
SyntaxError: If the task of the model could not be determined.
|
||||||
|
@ -180,10 +180,8 @@ class BaseDataset(Dataset):
|
|||||||
label = self.labels[index].copy()
|
label = self.labels[index].copy()
|
||||||
label.pop('shape', None) # shape is for rect, remove it
|
label.pop('shape', None) # shape is for rect, remove it
|
||||||
label['img'], label['ori_shape'], label['resized_shape'] = self.load_image(index)
|
label['img'], label['ori_shape'], label['resized_shape'] = self.load_image(index)
|
||||||
label['ratio_pad'] = (
|
label['ratio_pad'] = (label['resized_shape'][0] / label['ori_shape'][0],
|
||||||
label['resized_shape'][0] / label['ori_shape'][0],
|
label['resized_shape'][1] / label['ori_shape'][1]) # for evaluation
|
||||||
label['resized_shape'][1] / label['ori_shape'][1],
|
|
||||||
) # for evaluation
|
|
||||||
if self.rect:
|
if self.rect:
|
||||||
label['rect_shape'] = self.batch_shapes[self.batch[index]]
|
label['rect_shape'] = self.batch_shapes[self.batch[index]]
|
||||||
label = self.update_labels_info(label)
|
label = self.update_labels_info(label)
|
||||||
|
@ -17,7 +17,8 @@ from PIL import ExifTags, Image, ImageOps
|
|||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
|
|
||||||
from ultralytics.nn.autobackend import check_class_names
|
from ultralytics.nn.autobackend import check_class_names
|
||||||
from ultralytics.yolo.utils import DATASETS_DIR, LOGGER, NUM_THREADS, ROOT, clean_url, colorstr, emojis, yaml_load
|
from ultralytics.yolo.utils import (DATASETS_DIR, LOGGER, NUM_THREADS, ROOT, SETTINGS_YAML, clean_url, colorstr, emojis,
|
||||||
|
yaml_load)
|
||||||
from ultralytics.yolo.utils.checks import check_file, check_font, is_ascii
|
from ultralytics.yolo.utils.checks import check_file, check_font, is_ascii
|
||||||
from ultralytics.yolo.utils.downloads import download, safe_download, unzip_file
|
from ultralytics.yolo.utils.downloads import download, safe_download, unzip_file
|
||||||
from ultralytics.yolo.utils.ops import segments2boxes
|
from ultralytics.yolo.utils.ops import segments2boxes
|
||||||
@ -246,6 +247,7 @@ def check_det_dataset(dataset, autodownload=True):
|
|||||||
if s and autodownload:
|
if s and autodownload:
|
||||||
LOGGER.warning(m)
|
LOGGER.warning(m)
|
||||||
else:
|
else:
|
||||||
|
m += f"\nNote dataset download directory is '{DATASETS_DIR}'. You can update this in '{SETTINGS_YAML}'"
|
||||||
raise FileNotFoundError(m)
|
raise FileNotFoundError(m)
|
||||||
t = time.time()
|
t = time.time()
|
||||||
if s.startswith('http') and s.endswith('.zip'): # URL
|
if s.startswith('http') and s.endswith('.zip'): # URL
|
||||||
|
@ -356,7 +356,6 @@ class YOLO:
|
|||||||
raise AttributeError("Dataset required but missing, i.e. pass 'data=coco128.yaml'")
|
raise AttributeError("Dataset required but missing, i.e. pass 'data=coco128.yaml'")
|
||||||
if overrides.get('resume'):
|
if overrides.get('resume'):
|
||||||
overrides['resume'] = self.ckpt_path
|
overrides['resume'] = self.ckpt_path
|
||||||
|
|
||||||
self.task = overrides.get('task') or self.task
|
self.task = overrides.get('task') or self.task
|
||||||
self.trainer = TASK_MAP[self.task][1](overrides=overrides, _callbacks=self.callbacks)
|
self.trainer = TASK_MAP[self.task][1](overrides=overrides, _callbacks=self.callbacks)
|
||||||
if not overrides.get('resume'): # manually set model only if not resuming
|
if not overrides.get('resume'): # manually set model only if not resuming
|
||||||
|
@ -109,8 +109,35 @@ class BasePredictor:
|
|||||||
def preprocess(self, img):
|
def preprocess(self, img):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
def write_results(self, results, batch, print_string):
|
def write_results(self, idx, results, batch):
|
||||||
raise NotImplementedError('print_results function needs to be implemented')
|
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):
|
def postprocess(self, preds, img, orig_img):
|
||||||
return preds
|
return preds
|
||||||
|
@ -7,13 +7,14 @@ Usage: See https://docs.ultralytics.com/modes/predict/
|
|||||||
|
|
||||||
from copy import deepcopy
|
from copy import deepcopy
|
||||||
from functools import lru_cache
|
from functools import lru_cache
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
from ultralytics.yolo.data.augment import LetterBox
|
from ultralytics.yolo.data.augment import LetterBox
|
||||||
from ultralytics.yolo.utils import LOGGER, SimpleClass, deprecation_warn, ops
|
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):
|
class BaseTensor(SimpleClass):
|
||||||
@ -233,6 +234,80 @@ class Results(SimpleClass):
|
|||||||
|
|
||||||
return annotator.result()
|
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):
|
class Boxes(BaseTensor):
|
||||||
"""
|
"""
|
||||||
@ -339,6 +414,8 @@ class Masks(BaseTensor):
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, masks, orig_shape) -> None:
|
def __init__(self, masks, orig_shape) -> None:
|
||||||
|
if masks.ndim == 2:
|
||||||
|
masks = masks[None, :]
|
||||||
self.masks = masks # N, h, w
|
self.masks = masks # N, h, w
|
||||||
self.orig_shape = orig_shape
|
self.orig_shape = orig_shape
|
||||||
|
|
||||||
|
@ -552,7 +552,7 @@ class BaseTrainer:
|
|||||||
if self.resume:
|
if self.resume:
|
||||||
assert start_epoch > 0, \
|
assert start_epoch > 0, \
|
||||||
f'{self.args.model} training to {self.epochs} epochs is finished, nothing to resume.\n' \
|
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(
|
LOGGER.info(
|
||||||
f'Resuming training from {self.args.model} from epoch {start_epoch + 1} to {self.epochs} total epochs')
|
f'Resuming training from {self.args.model} from epoch {start_epoch + 1} to {self.epochs} total epochs')
|
||||||
if self.epochs < start_epoch:
|
if self.epochs < start_epoch:
|
||||||
|
@ -490,6 +490,7 @@ def get_user_config_dir(sub_dir='Ultralytics'):
|
|||||||
|
|
||||||
|
|
||||||
USER_CONFIG_DIR = Path(os.getenv('YOLO_CONFIG_DIR', get_user_config_dir())) # Ultralytics settings dir
|
USER_CONFIG_DIR = Path(os.getenv('YOLO_CONFIG_DIR', get_user_config_dir())) # Ultralytics settings dir
|
||||||
|
SETTINGS_YAML = USER_CONFIG_DIR / 'settings.yaml'
|
||||||
|
|
||||||
|
|
||||||
def emojis(string=''):
|
def emojis(string=''):
|
||||||
@ -591,7 +592,7 @@ def set_sentry():
|
|||||||
logging.getLogger(logger).setLevel(logging.CRITICAL)
|
logging.getLogger(logger).setLevel(logging.CRITICAL)
|
||||||
|
|
||||||
|
|
||||||
def get_settings(file=USER_CONFIG_DIR / 'settings.yaml', version='0.0.3'):
|
def get_settings(file=SETTINGS_YAML, version='0.0.3'):
|
||||||
"""
|
"""
|
||||||
Loads a global Ultralytics settings YAML file or creates one with default values if it does not exist.
|
Loads a global Ultralytics settings YAML file or creates one with default values if it does not exist.
|
||||||
|
|
||||||
@ -640,7 +641,7 @@ def get_settings(file=USER_CONFIG_DIR / 'settings.yaml', version='0.0.3'):
|
|||||||
return settings
|
return settings
|
||||||
|
|
||||||
|
|
||||||
def set_settings(kwargs, file=USER_CONFIG_DIR / 'settings.yaml'):
|
def set_settings(kwargs, file=SETTINGS_YAML):
|
||||||
"""
|
"""
|
||||||
Function that runs on a first-time ultralytics package installation to set up global settings and create necessary
|
Function that runs on a first-time ultralytics package installation to set up global settings and create necessary
|
||||||
directories.
|
directories.
|
||||||
|
@ -5,14 +5,10 @@ import torch
|
|||||||
from ultralytics.yolo.engine.predictor import BasePredictor
|
from ultralytics.yolo.engine.predictor import BasePredictor
|
||||||
from ultralytics.yolo.engine.results import Results
|
from ultralytics.yolo.engine.results import Results
|
||||||
from ultralytics.yolo.utils import DEFAULT_CFG, ROOT
|
from ultralytics.yolo.utils import DEFAULT_CFG, ROOT
|
||||||
from ultralytics.yolo.utils.plotting import Annotator
|
|
||||||
|
|
||||||
|
|
||||||
class ClassificationPredictor(BasePredictor):
|
class ClassificationPredictor(BasePredictor):
|
||||||
|
|
||||||
def get_annotator(self, img):
|
|
||||||
return Annotator(img, example=str(self.model.names), pil=True)
|
|
||||||
|
|
||||||
def preprocess(self, img):
|
def preprocess(self, img):
|
||||||
img = (img if isinstance(img, torch.Tensor) else torch.from_numpy(img)).to(self.model.device)
|
img = (img if isinstance(img, torch.Tensor) else torch.from_numpy(img)).to(self.model.device)
|
||||||
return img.half() if self.model.fp16 else img.float() # uint8 to fp16/32
|
return img.half() if self.model.fp16 else img.float() # uint8 to fp16/32
|
||||||
@ -27,43 +23,6 @@ class ClassificationPredictor(BasePredictor):
|
|||||||
|
|
||||||
return results
|
return results
|
||||||
|
|
||||||
def write_results(self, idx, results, batch):
|
|
||||||
p, im, im0 = batch
|
|
||||||
log_string = ''
|
|
||||||
if len(im.shape) == 3:
|
|
||||||
im = im[None] # expand for batch dim
|
|
||||||
self.seen += 1
|
|
||||||
im0 = im0.copy()
|
|
||||||
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
|
|
||||||
# save_path = str(self.save_dir / p.name) # im.jpg
|
|
||||||
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]
|
|
||||||
if len(result) == 0:
|
|
||||||
return log_string
|
|
||||||
prob = result.probs
|
|
||||||
# Print results
|
|
||||||
n5 = min(len(self.model.names), 5)
|
|
||||||
top5i = prob.argsort(0, descending=True)[:n5].tolist() # top 5 indices
|
|
||||||
log_string += f"{', '.join(f'{self.model.names[j]} {prob[j]:.2f}' for j in top5i)}, "
|
|
||||||
|
|
||||||
# write
|
|
||||||
if self.args.save or self.args.show: # Add bbox to image
|
|
||||||
self.plotted_img = result.plot()
|
|
||||||
if self.args.save_txt: # Write to file
|
|
||||||
text = '\n'.join(f'{prob[j]:.2f} {self.model.names[j]}' for j in top5i)
|
|
||||||
with open(f'{self.txt_path}.txt', 'a') as f:
|
|
||||||
f.write(text + '\n')
|
|
||||||
|
|
||||||
return log_string
|
|
||||||
|
|
||||||
|
|
||||||
def predict(cfg=DEFAULT_CFG, use_python=False):
|
def predict(cfg=DEFAULT_CFG, use_python=False):
|
||||||
model = cfg.model or 'yolov8n-cls.pt' # or "resnet18"
|
model = cfg.model or 'yolov8n-cls.pt' # or "resnet18"
|
||||||
|
@ -5,7 +5,6 @@ import torch
|
|||||||
from ultralytics.yolo.engine.predictor import BasePredictor
|
from ultralytics.yolo.engine.predictor import BasePredictor
|
||||||
from ultralytics.yolo.engine.results import Results
|
from ultralytics.yolo.engine.results import Results
|
||||||
from ultralytics.yolo.utils import DEFAULT_CFG, ROOT, ops
|
from ultralytics.yolo.utils import DEFAULT_CFG, ROOT, ops
|
||||||
from ultralytics.yolo.utils.plotting import save_one_box
|
|
||||||
|
|
||||||
|
|
||||||
class DetectionPredictor(BasePredictor):
|
class DetectionPredictor(BasePredictor):
|
||||||
@ -34,48 +33,6 @@ class DetectionPredictor(BasePredictor):
|
|||||||
results.append(Results(orig_img=orig_img, path=img_path, names=self.model.names, boxes=pred))
|
results.append(Results(orig_img=orig_img, path=img_path, names=self.model.names, boxes=pred))
|
||||||
return results
|
return results
|
||||||
|
|
||||||
def write_results(self, idx, results, batch):
|
|
||||||
p, im, im0 = batch
|
|
||||||
log_string = ''
|
|
||||||
if len(im.shape) == 3:
|
|
||||||
im = im[None] # expand for batch dim
|
|
||||||
self.seen += 1
|
|
||||||
imc = im0.copy() if self.args.save_crop else im0
|
|
||||||
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] # TODO: make boxes inherit from tensors
|
|
||||||
if len(result) == 0:
|
|
||||||
return f'{log_string}(no detections), '
|
|
||||||
det = result.boxes
|
|
||||||
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)}, "
|
|
||||||
|
|
||||||
if self.args.save or self.args.show: # Add bbox to image
|
|
||||||
self.plotted_img = result.plot(line_width=self.args.line_thickness)
|
|
||||||
|
|
||||||
# write
|
|
||||||
for d in reversed(det):
|
|
||||||
c, conf, id = int(d.cls), float(d.conf), None if d.id is None else int(d.id.item())
|
|
||||||
if self.args.save_txt: # Write to file
|
|
||||||
line = (c, *d.xywhn.view(-1)) + (conf, ) * self.args.save_conf + (() if id is None else (id, ))
|
|
||||||
with open(f'{self.txt_path}.txt', 'a') as f:
|
|
||||||
f.write(('%g ' * len(line)).rstrip() % line + '\n')
|
|
||||||
if self.args.save_crop:
|
|
||||||
save_one_box(d.xyxy,
|
|
||||||
imc,
|
|
||||||
file=self.save_dir / 'crops' / self.model.names[c] / f'{self.data_path.stem}.jpg',
|
|
||||||
BGR=True)
|
|
||||||
|
|
||||||
return log_string
|
|
||||||
|
|
||||||
|
|
||||||
def predict(cfg=DEFAULT_CFG, use_python=False):
|
def predict(cfg=DEFAULT_CFG, use_python=False):
|
||||||
model = cfg.model or 'yolov8n.pt'
|
model = cfg.model or 'yolov8n.pt'
|
||||||
|
@ -2,7 +2,6 @@
|
|||||||
|
|
||||||
from ultralytics.yolo.engine.results import Results
|
from ultralytics.yolo.engine.results import Results
|
||||||
from ultralytics.yolo.utils import DEFAULT_CFG, ROOT, ops
|
from ultralytics.yolo.utils import DEFAULT_CFG, ROOT, ops
|
||||||
from ultralytics.yolo.utils.plotting import save_one_box
|
|
||||||
from ultralytics.yolo.v8.detect.predict import DetectionPredictor
|
from ultralytics.yolo.v8.detect.predict import DetectionPredictor
|
||||||
|
|
||||||
|
|
||||||
@ -34,50 +33,6 @@ class PosePredictor(DetectionPredictor):
|
|||||||
keypoints=pred_kpts))
|
keypoints=pred_kpts))
|
||||||
return results
|
return results
|
||||||
|
|
||||||
def write_results(self, idx, results, batch):
|
|
||||||
p, im, im0 = batch
|
|
||||||
log_string = ''
|
|
||||||
if len(im.shape) == 3:
|
|
||||||
im = im[None] # expand for batch dim
|
|
||||||
self.seen += 1
|
|
||||||
imc = im0.copy() if self.args.save_crop else im0
|
|
||||||
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] # TODO: make boxes inherit from tensors
|
|
||||||
if len(result) == 0:
|
|
||||||
return f'{log_string}(no detections), '
|
|
||||||
det = result.boxes
|
|
||||||
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)}, "
|
|
||||||
|
|
||||||
if self.args.save or self.args.show: # Add bbox to image
|
|
||||||
self.plotted_img = result.plot(line_width=self.args.line_thickness, boxes=self.args.boxes)
|
|
||||||
|
|
||||||
# write
|
|
||||||
for j, d in enumerate(reversed(det)):
|
|
||||||
c, conf, id = int(d.cls), float(d.conf), None if d.id is None else int(d.id.item())
|
|
||||||
if self.args.save_txt: # Write to file
|
|
||||||
kpt = (result[j].keypoints[:, :2] / d.orig_shape[[1, 0]]).reshape(-1).tolist()
|
|
||||||
box = d.xywhn.view(-1).tolist()
|
|
||||||
line = (c, *box, *kpt) + (conf, ) * self.args.save_conf + (() if id is None else (id, ))
|
|
||||||
with open(f'{self.txt_path}.txt', 'a') as f:
|
|
||||||
f.write(('%g ' * len(line)).rstrip() % line + '\n')
|
|
||||||
if self.args.save_crop:
|
|
||||||
save_one_box(d.xyxy,
|
|
||||||
imc,
|
|
||||||
file=self.save_dir / 'crops' / self.model.model.names[c] / f'{self.data_path.stem}.jpg',
|
|
||||||
BGR=True)
|
|
||||||
|
|
||||||
return log_string
|
|
||||||
|
|
||||||
|
|
||||||
def predict(cfg=DEFAULT_CFG, use_python=False):
|
def predict(cfg=DEFAULT_CFG, use_python=False):
|
||||||
model = cfg.model or 'yolov8n-pose.pt'
|
model = cfg.model or 'yolov8n-pose.pt'
|
||||||
|
@ -4,7 +4,6 @@ import torch
|
|||||||
|
|
||||||
from ultralytics.yolo.engine.results import Results
|
from ultralytics.yolo.engine.results import Results
|
||||||
from ultralytics.yolo.utils import DEFAULT_CFG, ROOT, ops
|
from ultralytics.yolo.utils import DEFAULT_CFG, ROOT, ops
|
||||||
from ultralytics.yolo.utils.plotting import save_one_box
|
|
||||||
from ultralytics.yolo.v8.detect.predict import DetectionPredictor
|
from ultralytics.yolo.v8.detect.predict import DetectionPredictor
|
||||||
|
|
||||||
|
|
||||||
@ -40,55 +39,6 @@ class SegmentationPredictor(DetectionPredictor):
|
|||||||
Results(orig_img=orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], masks=masks))
|
Results(orig_img=orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], masks=masks))
|
||||||
return results
|
return results
|
||||||
|
|
||||||
def write_results(self, idx, results, batch):
|
|
||||||
p, im, im0 = batch
|
|
||||||
log_string = ''
|
|
||||||
if len(im.shape) == 3:
|
|
||||||
im = im[None] # expand for batch dim
|
|
||||||
self.seen += 1
|
|
||||||
imc = im0.copy() if self.args.save_crop else im0
|
|
||||||
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]
|
|
||||||
if len(result) == 0:
|
|
||||||
return f'{log_string}(no detections), '
|
|
||||||
det, mask = result.boxes, result.masks # getting tensors TODO: mask mask,box inherit for tensor
|
|
||||||
|
|
||||||
# Print results
|
|
||||||
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)}, "
|
|
||||||
|
|
||||||
# Mask plotting
|
|
||||||
if self.args.save or self.args.show:
|
|
||||||
im_gpu = torch.as_tensor(im0, dtype=torch.float16, device=mask.masks.device).permute(
|
|
||||||
2, 0, 1).flip(0).contiguous() / 255 if self.args.retina_masks else im[idx]
|
|
||||||
self.plotted_img = result.plot(line_width=self.args.line_thickness, im_gpu=im_gpu, boxes=self.args.boxes)
|
|
||||||
|
|
||||||
# Write results
|
|
||||||
for j, d in enumerate(reversed(det)):
|
|
||||||
c, conf, id = int(d.cls), float(d.conf), None if d.id is None else int(d.id.item())
|
|
||||||
if self.args.save_txt: # Write to file
|
|
||||||
seg = mask.xyn[len(det) - j - 1].copy().reshape(-1) # reversed mask.xyn, (n,2) to (n*2)
|
|
||||||
line = (c, *seg) + (conf, ) * self.args.save_conf + (() if id is None else (id, ))
|
|
||||||
with open(f'{self.txt_path}.txt', 'a') as f:
|
|
||||||
f.write(('%g ' * len(line)).rstrip() % line + '\n')
|
|
||||||
if self.args.save_crop:
|
|
||||||
save_one_box(d.xyxy,
|
|
||||||
imc,
|
|
||||||
file=self.save_dir / 'crops' / self.model.names[c] / f'{self.data_path.stem}.jpg',
|
|
||||||
BGR=True)
|
|
||||||
|
|
||||||
return log_string
|
|
||||||
|
|
||||||
|
|
||||||
def predict(cfg=DEFAULT_CFG, use_python=False):
|
def predict(cfg=DEFAULT_CFG, use_python=False):
|
||||||
model = cfg.model or 'yolov8n-seg.pt'
|
model = cfg.model or 'yolov8n-seg.pt'
|
||||||
|
Loading…
x
Reference in New Issue
Block a user