xdobro23/YOLOv8-16bit
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

General cleanup (#69)

Browse Source 
Co-authored-by: ayush chaurasia <ayush.chaurarsia@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Laughing
2022-12-08 08:28:13 -06:00
committed by GitHub
parent 7ae45c6cc4
commit d63ee112d4
13 changed files with 265 additions and 433 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
ultralytics/yolo/v8/detect/__init__.py
View File

@ -1,3 +1,3 @@
from ultralytics.yolo.v8.detect.predict import DetectionPredictor, predict
from ultralytics.yolo.v8.detect.train import DetectionTrainer, train
from ultralytics.yolo.v8.detect.val import DetectionValidator, val
from .predict import DetectionPredictor, predict
from .train import DetectionTrainer, train
from .val import DetectionValidator, val

34
ultralytics/yolo/v8/detect/train.py
View File

@ -2,18 +2,37 @@ import hydra
import torch
import torch.nn as nn
from ultralytics.yolo.engine.trainer import DEFAULT_CONFIG
from ultralytics.yolo import v8
from ultralytics.yolo.data import build_dataloader
from ultralytics.yolo.engine.trainer import DEFAULT_CONFIG, BaseTrainer
from ultralytics.yolo.utils.metrics import FocalLoss, bbox_iou, smooth_BCE
from ultralytics.yolo.utils.modeling.tasks import DetectionModel
from ultralytics.yolo.utils.plotting import plot_images, plot_results
from ultralytics.yolo.utils.torch_utils import de_parallel
from ..segment import SegmentationTrainer
from .val import DetectionValidator
# BaseTrainer python usage
class DetectionTrainer(SegmentationTrainer):
class DetectionTrainer(BaseTrainer):
def get_dataloader(self, dataset_path, batch_size, mode="train", rank=0):
# TODO: manage splits differently
# calculate stride - check if model is initialized
gs = max(int(de_parallel(self.model).stride.max() if self.model else 0), 32)
return build_dataloader(self.args, batch_size, img_path=dataset_path, stride=gs, rank=rank, mode=mode)[0]
def preprocess_batch(self, batch):
batch["img"] = batch["img"].to(self.device, non_blocking=True).float() / 255
return batch
def set_model_attributes(self):
nl = de_parallel(self.model).model[-1].nl # number of detection layers (to scale hyps)
self.args.box *= 3 / nl # scale to layers
self.args.cls *= self.data["nc"] / 80 * 3 / nl # scale to classes and layers
self.args.obj *= (self.args.img_size / 640) ** 2 * 3 / nl # scale to image size and layers
self.model.nc = self.data["nc"] # attach number of classes to model
self.model.args = self.args # attach hyperparameters to model
# TODO: self.model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc
self.model.names = self.data["names"]
def load_model(self, model_cfg=None, weights=None):
model = DetectionModel(model_cfg or weights["model"].yaml,
@ -27,7 +46,10 @@ class DetectionTrainer(SegmentationTrainer):
return model
def get_validator(self):
return DetectionValidator(self.test_loader, save_dir=self.save_dir, logger=self.console, args=self.args)
return v8.detect.DetectionValidator(self.test_loader,
save_dir=self.save_dir,
logger=self.console,
args=self.args)
def criterion(self, preds, batch):
head = de_parallel(self.model).model[-1]

18
ultralytics/yolo/v8/detect/val.py
View File

@ -11,7 +11,7 @@ from ultralytics.yolo.engine.validator import BaseValidator
from ultralytics.yolo.utils import ops
from ultralytics.yolo.utils.checks import check_file, check_requirements
from ultralytics.yolo.utils.files import yaml_load
from ultralytics.yolo.utils.metrics import ConfusionMatrix, Metric, ap_per_class, box_iou, fitness_detection
from ultralytics.yolo.utils.metrics import ConfusionMatrix, DetMetrics, box_iou
from ultralytics.yolo.utils.plotting import output_to_target, plot_images
from ultralytics.yolo.utils.torch_utils import de_parallel
@ -62,7 +62,7 @@ class DetectionValidator(BaseValidator):
self.niou = self.iouv.numel()
self.seen = 0
self.confusion_matrix = ConfusionMatrix(nc=self.nc)
self.metrics = Metric()
self.metrics = DetMetrics(save_dir=self.save_dir, plot=self.args.plots, names=self.names)
self.loss = torch.zeros(3, device=self.device)
self.jdict = []
self.stats = []
@ -128,10 +128,9 @@ class DetectionValidator(BaseValidator):
def get_stats(self):
stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*self.stats)] # to numpy
if len(stats) and stats[0].any():
results = ap_per_class(*stats, plot=self.args.plots, save_dir=self.save_dir, names=self.names)
self.metrics.update(results[2:])
self.nt_per_class = np.bincount(stats[3].astype(int), minlength=self.nc) # number of targets per class
metrics = {"fitness": fitness_detection(np.array(self.metrics.mean_results()).reshape(1, -1))}
self.metrics.process(*stats)
self.nt_per_class = np.bincount(stats[-1].astype(int), minlength=self.nc) # number of targets per class
metrics = {"fitness": self.metrics.fitness()}
metrics |= zip(self.metric_keys, self.metrics.mean_results())
return metrics
@ -203,8 +202,11 @@ class DetectionValidator(BaseValidator):
def plot_predictions(self, batch, preds, ni):
images = batch["img"]
paths = batch["im_file"]
plot_images(images, *output_to_target(preds, max_det=15), paths, self.save_dir / f'val_batch{ni}_pred.jpg',
self.names) # pred
plot_images(images,
*output_to_target(preds, max_det=15),
paths=paths,
fname=self.save_dir / f'val_batch{ni}_pred.jpg',
names=self.names) # pred
@hydra.main(version_base=None, config_path=DEFAULT_CONFIG.parent, config_name=DEFAULT_CONFIG.name)

6
ultralytics/yolo/v8/segment/__init__.py
View File

@ -1,3 +1,3 @@
from ultralytics.yolo.v8.segment.predict import SegmentationPredictor, predict
from ultralytics.yolo.v8.segment.train import SegmentationTrainer, train
from ultralytics.yolo.v8.segment.val import SegmentationValidator, val
from .predict import SegmentationPredictor, predict
from .train import SegmentationTrainer, train
from .val import SegmentationValidator, val

37
ultralytics/yolo/v8/segment/train.py
View File

@ -4,27 +4,18 @@ import torch.nn as nn
import torch.nn.functional as F
from ultralytics.yolo import v8
from ultralytics.yolo.data import build_dataloader
from ultralytics.yolo.engine.trainer import DEFAULT_CONFIG, BaseTrainer
from ultralytics.yolo.utils.metrics import FocalLoss, bbox_iou, smooth_BCE
from ultralytics.yolo.utils.modeling.tasks import SegmentationModel
from ultralytics.yolo.utils.ops import crop_mask, xywh2xyxy
from ultralytics.yolo.utils.plotting import plot_images_and_masks, plot_results_with_masks
from ultralytics.yolo.utils.plotting import plot_images, plot_results
from ultralytics.yolo.utils.torch_utils import de_parallel
from ..detect import DetectionTrainer
# BaseTrainer python usage
class SegmentationTrainer(BaseTrainer):
def get_dataloader(self, dataset_path, batch_size, mode="train", rank=0):
# TODO: manage splits differently
# calculate stride - check if model is initialized
gs = max(int(de_parallel(self.model).stride.max() if self.model else 0), 32)
return build_dataloader(self.args, batch_size, img_path=dataset_path, stride=gs, rank=rank, mode=mode)[0]
def preprocess_batch(self, batch):
batch["img"] = batch["img"].to(self.device, non_blocking=True).float() / 255
return batch
class SegmentationTrainer(DetectionTrainer):
def load_model(self, model_cfg=None, weights=None):
model = SegmentationModel(model_cfg or weights["model"].yaml,
@ -37,16 +28,6 @@ class SegmentationTrainer(BaseTrainer):
v.requires_grad = True # train all layers
return model
def set_model_attributes(self):
nl = de_parallel(self.model).model[-1].nl # number of detection layers (to scale hyps)
self.args.box *= 3 / nl # scale to layers
self.args.cls *= self.data["nc"] / 80 * 3 / nl # scale to classes and layers
self.args.obj *= (self.args.img_size / 640) ** 2 * 3 / nl # scale to image size and layers
self.model.nc = self.data["nc"] # attach number of classes to model
self.model.args = self.args # attach hyperparameters to model
# TODO: self.model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc
self.model.names = self.data["names"]
def get_validator(self):
return v8.segment.SegmentationValidator(self.test_loader,
save_dir=self.save_dir,
@ -245,16 +226,10 @@ class SegmentationTrainer(BaseTrainer):
bboxes = batch["bboxes"]
paths = batch["im_file"]
batch_idx = batch["batch_idx"]
plot_images_and_masks(images,
batch_idx,
cls,
bboxes,
masks,
paths=paths,
fname=self.save_dir / f"train_batch{ni}.jpg")
plot_images(images, batch_idx, cls, bboxes, masks, paths=paths, fname=self.save_dir / f"train_batch{ni}.jpg")
def plot_metrics(self):
plot_results_with_masks(file=self.csv) # save results.png
plot_results(file=self.csv, segment=True) # save results.png
@hydra.main(version_base=None, config_path=DEFAULT_CONFIG.parent, config_name=DEFAULT_CONFIG.name)

44
ultralytics/yolo/v8/segment/val.py
View File

@ -7,17 +7,17 @@ import torch.nn.functional as F
from ultralytics.yolo.data import build_dataloader
from ultralytics.yolo.engine.trainer import DEFAULT_CONFIG
from ultralytics.yolo.engine.validator import BaseValidator
from ultralytics.yolo.utils import ops
from ultralytics.yolo.utils.checks import check_file, check_requirements
from ultralytics.yolo.utils.files import yaml_load
from ultralytics.yolo.utils.metrics import (ConfusionMatrix, Metrics, ap_per_class_box_and_mask, box_iou,
fitness_segmentation, mask_iou)
from ultralytics.yolo.utils.plotting import output_to_target, plot_images_and_masks
from ultralytics.yolo.utils.metrics import ConfusionMatrix, SegmentMetrics, box_iou, mask_iou
from ultralytics.yolo.utils.plotting import output_to_target, plot_images
from ultralytics.yolo.utils.torch_utils import de_parallel
from ..detect import DetectionValidator
class SegmentationValidator(BaseValidator):
class SegmentationValidator(DetectionValidator):
def __init__(self, dataloader=None, save_dir=None, pbar=None, logger=None, args=None):
super().__init__(dataloader, save_dir, pbar, logger, args)
@ -65,7 +65,7 @@ class SegmentationValidator(BaseValidator):
self.niou = self.iouv.numel()
self.seen = 0
self.confusion_matrix = ConfusionMatrix(nc=self.nc)
self.metrics = Metrics()
self.metrics = SegmentMetrics(save_dir=self.save_dir, plot=self.args.plots, names=self.names)
self.loss = torch.zeros(4, device=self.device)
self.jdict = []
self.stats = []
@ -150,16 +150,6 @@ class SegmentationValidator(BaseValidator):
# callbacks.run('on_val_image_end', pred, predn, path, names, im[si])
'''
def get_stats(self):
stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*self.stats)] # to numpy
if len(stats) and stats[0].any():
results = ap_per_class_box_and_mask(*stats, plot=self.args.plots, save_dir=self.save_dir, names=self.names)
self.metrics.update(results)
self.nt_per_class = np.bincount(stats[4].astype(int), minlength=self.nc) # number of targets per class
metrics = {"fitness": fitness_segmentation(np.array(self.metrics.mean_results()).reshape(1, -1))}
metrics |= zip(self.metric_keys, self.metrics.mean_results())
return metrics
def print_results(self):
pf = '%22s' + '%11i' * 2 + '%11.3g' * 8 # print format
self.logger.info(pf % ("all", self.seen, self.nt_per_class.sum(), *self.metrics.mean_results()))
@ -218,6 +208,7 @@ class SegmentationValidator(BaseValidator):
gs = max(int(de_parallel(self.model).stride if self.model else 0), 32)
return build_dataloader(self.args, batch_size, img_path=dataset_path, stride=gs, mode="val")[0]
# TODO: probably add this to class Metrics
@property
def metric_keys(self):
return [
@ -237,23 +228,22 @@ class SegmentationValidator(BaseValidator):
bboxes = batch["bboxes"]
paths = batch["im_file"]
batch_idx = batch["batch_idx"]
plot_images_and_masks(images,
batch_idx,
cls,
bboxes,
masks,
paths=paths,
fname=self.save_dir / f"val_batch{ni}_labels.jpg",
names=self.names)
plot_images(images,
batch_idx,
cls,
bboxes,
masks,
paths=paths,
fname=self.save_dir / f"val_batch{ni}_labels.jpg",
names=self.names)
def plot_predictions(self, batch, preds, ni):
images = batch["img"]
paths = batch["im_file"]
if len(self.plot_masks):
plot_masks = torch.cat(self.plot_masks, dim=0)
batch_idx, cls, bboxes, conf = output_to_target(preds[0], max_det=15)
plot_images_and_masks(images, batch_idx, cls, bboxes, plot_masks, conf, paths,
self.save_dir / f'val_batch{ni}_pred.jpg', self.names) # pred
plot_images(images, *output_to_target(preds[0], max_det=15), plot_masks, paths,
self.save_dir / f'val_batch{ni}_pred.jpg', self.names) # pred
self.plot_masks.clear()