`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>
single_channel
Glenn Jocher 2 years ago committed by GitHub
parent 95f96dc5bc
commit 5629ed0bb7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -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

@ -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…
Cancel
Save