first commit
This commit is contained in:
commit
b680e09bbe
568 files changed
+54067
No files matched your search
@@ -0,0 +1,7 @@
|
||||
# Ultralytics YOLO 🚀, AGPL-3.0 license
|
||||
|
||||
from .predict import DetectionPredictor
|
||||
from .train import DetectionTrainer
|
||||
from .val import DetectionValidator
|
||||
|
||||
__all__ = "DetectionPredictor", "DetectionTrainer", "DetectionValidator"
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,121 @@
|
||||
# Ultralytics YOLO 🚀, AGPL-3.0 license
|
||||
import torch
|
||||
from ultralytics.engine.predictor import BasePredictor
|
||||
from ultralytics.engine.results import Results
|
||||
from ultralytics.utils import ops
|
||||
|
||||
|
||||
class DetectionPredictor(BasePredictor):
|
||||
"""
|
||||
A class extending the BasePredictor class for prediction based on a detection model.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from ultralytics.utils import ASSETS
|
||||
from ultralytics.models.yolo.detect import DetectionPredictor
|
||||
|
||||
args = dict(model='yolov8n.pt', source=ASSETS)
|
||||
predictor = DetectionPredictor(overrides=args)
|
||||
predictor.predict_cli()
|
||||
```
|
||||
"""
|
||||
|
||||
def postprocess(self, preds, img, orig_imgs):
|
||||
"""Post-processes predictions and returns a list of Results objects."""
|
||||
# if isinstance(preds, (list, tuple)): # YOLOv8 model in validation model, output = (inference_out, loss_out)
|
||||
# preds = preds[0]
|
||||
# cls_max = torch.max(preds[:, 4:8, :], dim=1, keepdim=True).values
|
||||
# idx = (cls_max > self.args.conf)
|
||||
# preds[:, -1:, :][idx] = 0
|
||||
|
||||
# preds = ops.non_max_suppression(
|
||||
# preds,
|
||||
# self.args.conf,
|
||||
# self.args.iou,
|
||||
# agnostic=self.args.agnostic_nms,
|
||||
# max_det=self.args.max_det,
|
||||
# classes=self.args.classes,
|
||||
# )
|
||||
|
||||
if isinstance(preds, (list, tuple)): # YOLOv8 model in validation model, output = (inference_out, loss_out)
|
||||
preds = preds[0]
|
||||
preds_c = preds.clone()
|
||||
|
||||
nms_cls = ops.non_max_suppression(
|
||||
preds[:, :8, :],
|
||||
self.args.conf,
|
||||
self.args.iou,
|
||||
agnostic=self.args.agnostic_nms,
|
||||
max_det=self.args.max_det,
|
||||
classes=self.args.classes,
|
||||
)
|
||||
|
||||
nms_obj = ops.non_max_suppression(
|
||||
torch.cat((preds_c[:, :4, :], preds_c[:, -1:, :]), dim=1),
|
||||
self.args.conf,
|
||||
self.args.iou,
|
||||
agnostic=self.args.agnostic_nms,
|
||||
max_det=self.args.max_det,
|
||||
classes=self.args.classes,
|
||||
)
|
||||
|
||||
preds = [n.clone() for n in nms_cls]
|
||||
for i in range(len(nms_cls)):
|
||||
nms_obj[i][:, 5] = len(self.model.names) - 1
|
||||
if nms_obj[i].size(0) == 0: # Skip if there are no objects
|
||||
continue
|
||||
if nms_cls[i].size(0) == 0:
|
||||
preds[i] = nms_obj[i]
|
||||
continue
|
||||
iou = self.box_iou(nms_cls[i][:, :4], nms_obj[i][:, :4])
|
||||
idxs = []
|
||||
# 将多分类置信度低、二分类置信度高,且IOU较大的框替换
|
||||
matchs = torch.ones_like(iou)
|
||||
matchs[nms_cls[i][:, 4] > 0.3, :] = 0 # new_track_thresh
|
||||
matchs[:, nms_obj[i][:, 4] < 0.3] = 0
|
||||
matchs[iou < self.args.iou] = 0
|
||||
rows, cols = torch.where(matchs == 1)
|
||||
preds[i][rows, :] = nms_obj[i][cols, :]
|
||||
idxs.append(cols)
|
||||
|
||||
# 将二分类无对应多分类结果的框附后
|
||||
max_values = torch.max(iou, dim=0)[0]
|
||||
idxs = torch.where(max_values < self.args.iou)[0]
|
||||
idxs = torch.tensor([i for i in idxs if i not in cols]).to(torch.int64)
|
||||
preds[i] = torch.cat((preds[i], nms_obj[i][idxs, :]), dim=0)
|
||||
|
||||
# for m in range(len(nms_obj[i])):
|
||||
# if (torch.max(iou[:, m]) < self.args.iou) and (m not in idxs):
|
||||
# preds[i] = torch.cat((preds[i], nms_obj[i][m, :]), dim=0)
|
||||
|
||||
if not isinstance(orig_imgs, list): # input images are a torch.Tensor, not a list
|
||||
orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
|
||||
|
||||
results = []
|
||||
for i, pred in enumerate(preds):
|
||||
orig_img = orig_imgs[i]
|
||||
pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
|
||||
img_path = self.batch[0][i]
|
||||
results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred))
|
||||
return results
|
||||
|
||||
def box_iou(self, box1, box2, eps=1e-7):
|
||||
"""
|
||||
Calculate intersection-over-union (IoU) of boxes. Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
|
||||
Based on https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
|
||||
|
||||
Args:
|
||||
box1 (torch.Tensor): A tensor of shape (N, 4) representing N bounding boxes.
|
||||
box2 (torch.Tensor): A tensor of shape (M, 4) representing M bounding boxes.
|
||||
eps (float, optional): A small value to avoid division by zero. Defaults to 1e-7.
|
||||
|
||||
Returns:
|
||||
(torch.Tensor): An NxM tensor containing the pairwise IoU values for every element in box1 and box2.
|
||||
"""
|
||||
|
||||
# inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
|
||||
(a1, a2), (b1, b2) = box1.unsqueeze(1).chunk(2, 2), box2.unsqueeze(0).chunk(2, 2)
|
||||
inter = (torch.min(a2, b2) - torch.max(a1, b1)).clamp_(0).prod(2)
|
||||
|
||||
# IoU = inter / (area1 + area2 - inter)
|
||||
return inter / ((a2 - a1).prod(2) + (b2 - b1).prod(2) - inter + eps)
|
||||
@@ -0,0 +1,143 @@
|
||||
# Ultralytics YOLO 🚀, AGPL-3.0 license
|
||||
|
||||
import math
|
||||
import random
|
||||
from copy import copy
|
||||
|
||||
import numpy as np
|
||||
import torch.nn as nn
|
||||
|
||||
from ultralytics.data import build_dataloader, build_yolo_dataset
|
||||
from ultralytics.engine.trainer import BaseTrainer
|
||||
from ultralytics.models import yolo
|
||||
from ultralytics.nn.tasks import DetectionBinModel
|
||||
from ultralytics.utils import LOGGER, RANK
|
||||
from ultralytics.utils.plotting import plot_images, plot_labels, plot_results
|
||||
from ultralytics.utils.torch_utils import de_parallel, torch_distributed_zero_first
|
||||
|
||||
|
||||
class DetectionTrainer(BaseTrainer):
|
||||
"""
|
||||
A class extending the BaseTrainer class for training based on a detection model.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from ultralytics.models.yolo.detect import DetectionTrainer
|
||||
|
||||
args = dict(model='yolov8n.pt', data='coco8.yaml', epochs=3)
|
||||
trainer = DetectionTrainer(overrides=args)
|
||||
trainer.train()
|
||||
```
|
||||
"""
|
||||
|
||||
def build_dataset(self, img_path, mode="train", batch=None):
|
||||
"""
|
||||
Build YOLO Dataset.
|
||||
|
||||
Args:
|
||||
img_path (str): Path to the folder containing images.
|
||||
mode (str): `train` mode or `val` mode, users are able to customize different augmentations for each mode.
|
||||
batch (int, optional): Size of batches, this is for `rect`. Defaults to None.
|
||||
"""
|
||||
gs = max(int(de_parallel(self.model).stride.max() if self.model else 0), 32)
|
||||
return build_yolo_dataset(self.args, img_path, batch, self.data, mode=mode, rect=mode == "val", stride=gs)
|
||||
|
||||
def get_dataloader(self, dataset_path, batch_size=16, rank=0, mode="train"):
|
||||
"""Construct and return dataloader."""
|
||||
assert mode in {"train", "val"}, f"Mode must be 'train' or 'val', not {mode}."
|
||||
with torch_distributed_zero_first(rank): # init dataset *.cache only once if DDP
|
||||
dataset = self.build_dataset(dataset_path, mode, batch_size)
|
||||
shuffle = mode == "train"
|
||||
if getattr(dataset, "rect", False) and shuffle:
|
||||
LOGGER.warning("WARNING ⚠️ 'rect=True' is incompatible with DataLoader shuffle, setting shuffle=False")
|
||||
shuffle = False
|
||||
workers = self.args.workers if mode == "train" else self.args.workers * 2
|
||||
return build_dataloader(dataset, batch_size, workers, shuffle, rank) # return dataloader
|
||||
|
||||
def preprocess_batch(self, batch):
|
||||
"""Preprocesses a batch of images by scaling and converting to float."""
|
||||
batch["img"] = batch["img"].to(self.device, non_blocking=True).float() / 255
|
||||
if self.args.multi_scale:
|
||||
imgs = batch["img"]
|
||||
sz = (
|
||||
random.randrange(self.args.imgsz * 0.5, self.args.imgsz * 1.5 + self.stride)
|
||||
// self.stride
|
||||
* self.stride
|
||||
) # size
|
||||
sf = sz / max(imgs.shape[2:]) # scale factor
|
||||
if sf != 1:
|
||||
ns = [
|
||||
math.ceil(x * sf / self.stride) * self.stride for x in imgs.shape[2:]
|
||||
] # new shape (stretched to gs-multiple)
|
||||
imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False)
|
||||
batch["img"] = imgs
|
||||
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.cls *= (self.args.imgsz / 640) ** 2 * 3 / nl # scale to image size and layers
|
||||
self.model.nc = self.data["nc"] # attach number of classes to model
|
||||
self.model.names = self.data["names"] # attach class names 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
|
||||
|
||||
def get_model(self, cfg=None, weights=None, verbose=True):
|
||||
"""Return a YOLO detection model."""
|
||||
model = DetectionBinModel(cfg, nc=self.data["nc"], verbose=verbose and RANK == -1)
|
||||
if weights:
|
||||
model.load(weights)
|
||||
return model
|
||||
|
||||
def get_validator(self):
|
||||
"""Returns a DetectionValidator for YOLO model validation."""
|
||||
self.loss_names = "box_loss", "cls_loss", "obj_loss", "dfl_loss"
|
||||
return yolo.detectbin.DetectionValidator(
|
||||
self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
|
||||
)
|
||||
|
||||
def label_loss_items(self, loss_items=None, prefix="train"):
|
||||
"""
|
||||
Returns a loss dict with labelled training loss items tensor.
|
||||
|
||||
Not needed for classification but necessary for segmentation & detection
|
||||
"""
|
||||
keys = [f"{prefix}/{x}" for x in self.loss_names]
|
||||
if loss_items is not None:
|
||||
loss_items = [round(float(x), 5) for x in loss_items] # convert tensors to 5 decimal place floats
|
||||
return dict(zip(keys, loss_items))
|
||||
else:
|
||||
return keys
|
||||
|
||||
def progress_string(self):
|
||||
"""Returns a formatted string of training progress with epoch, GPU memory, loss, instances and size."""
|
||||
return ("\n" + "%11s" * (4 + len(self.loss_names))) % (
|
||||
"Epoch",
|
||||
"GPU_mem",
|
||||
*self.loss_names,
|
||||
"Instances",
|
||||
"Size",
|
||||
)
|
||||
|
||||
def plot_training_samples(self, batch, ni):
|
||||
"""Plots training samples with their annotations."""
|
||||
plot_images(
|
||||
images=batch["img"],
|
||||
batch_idx=batch["batch_idx"],
|
||||
cls=batch["cls"].squeeze(-1),
|
||||
bboxes=batch["bboxes"],
|
||||
paths=batch["im_file"],
|
||||
fname=self.save_dir / f"train_batch{ni}.jpg",
|
||||
on_plot=self.on_plot,
|
||||
)
|
||||
|
||||
def plot_metrics(self):
|
||||
"""Plots metrics from a CSV file."""
|
||||
plot_results(file=self.csv, on_plot=self.on_plot) # save results.png
|
||||
|
||||
def plot_training_labels(self):
|
||||
"""Create a labeled training plot of the YOLO model."""
|
||||
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, on_plot=self.on_plot)
|
||||
@@ -0,0 +1,344 @@
|
||||
# Ultralytics YOLO 🚀, AGPL-3.0 license
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from ultralytics.data import build_dataloader, build_yolo_dataset, converter
|
||||
from ultralytics.engine.validator import BaseValidator
|
||||
from ultralytics.utils import LOGGER, ops
|
||||
from ultralytics.utils.checks import check_requirements
|
||||
from ultralytics.utils.metrics import ConfusionMatrix, DetMetrics, box_iou
|
||||
from ultralytics.utils.plotting import output_to_target, plot_images
|
||||
|
||||
|
||||
class DetectionValidator(BaseValidator):
|
||||
"""
|
||||
A class extending the BaseValidator class for validation based on a detection model.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from ultralytics.models.yolo.detect import DetectionValidator
|
||||
|
||||
args = dict(model='yolov8n.pt', data='coco8.yaml')
|
||||
validator = DetectionValidator(args=args)
|
||||
validator()
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, dataloader=None, save_dir=None, pbar=None, args=None, _callbacks=None):
|
||||
"""Initialize detection model with necessary variables and settings."""
|
||||
super().__init__(dataloader, save_dir, pbar, args, _callbacks)
|
||||
self.nt_per_class = None
|
||||
self.is_coco = False
|
||||
self.is_lvis = False
|
||||
self.class_map = None
|
||||
self.args.task = "detectbin"
|
||||
self.metrics = DetMetrics(save_dir=self.save_dir, on_plot=self.on_plot)
|
||||
self.iouv = torch.linspace(0.5, 0.95, 10) # IoU vector for mAP@0.5:0.95
|
||||
self.niou = self.iouv.numel()
|
||||
self.lb = [] # for autolabelling
|
||||
|
||||
def preprocess(self, batch):
|
||||
"""Preprocesses batch of images for YOLO training."""
|
||||
batch["img"] = batch["img"].to(self.device, non_blocking=True)
|
||||
batch["img"] = (batch["img"].half() if self.args.half else batch["img"].float()) / 255
|
||||
for k in ["batch_idx", "cls", "bboxes"]:
|
||||
batch[k] = batch[k].to(self.device)
|
||||
|
||||
if self.args.save_hybrid:
|
||||
height, width = batch["img"].shape[2:]
|
||||
nb = len(batch["img"])
|
||||
bboxes = batch["bboxes"] * torch.tensor((width, height, width, height), device=self.device)
|
||||
self.lb = (
|
||||
[
|
||||
torch.cat([batch["cls"][batch["batch_idx"] == i], bboxes[batch["batch_idx"] == i]], dim=-1)
|
||||
for i in range(nb)
|
||||
]
|
||||
if self.args.save_hybrid
|
||||
else []
|
||||
) # for autolabelling
|
||||
|
||||
return batch
|
||||
|
||||
def init_metrics(self, model):
|
||||
"""Initialize evaluation metrics for YOLO."""
|
||||
val = self.data.get(self.args.split, "") # validation path
|
||||
self.is_coco = isinstance(val, str) and "coco" in val and val.endswith(f"{os.sep}val2017.txt") # is COCO
|
||||
self.is_lvis = isinstance(val, str) and "lvis" in val and not self.is_coco # is LVIS
|
||||
self.class_map = converter.coco80_to_coco91_class() if self.is_coco else list(range(len(model.names)))
|
||||
self.args.save_json |= (self.is_coco or self.is_lvis) and not self.training # run on final val if training COCO
|
||||
self.names = model.names
|
||||
self.nc = len(model.names)
|
||||
|
||||
self.names.update({self.nc: "Object"})
|
||||
|
||||
self.metrics.names = self.names
|
||||
self.metrics.plot = self.args.plots
|
||||
self.confusion_matrix = ConfusionMatrix(nc=self.nc+1, conf=self.args.conf)
|
||||
self.seen = 0
|
||||
self.jdict = []
|
||||
self.stats = dict(tp=[], conf=[], pred_cls=[], target_cls=[])
|
||||
|
||||
def get_desc(self):
|
||||
"""Return a formatted string summarizing class metrics of YOLO model."""
|
||||
return ("%22s" + "%11s" * 7) % ("Class", "Images", "Instances", "Box(P", "R", "mAP50", "mAP75", "mAP50-95)")
|
||||
|
||||
def postprocess(self, preds):
|
||||
"""Apply Non-maximum suppression to prediction outputs."""
|
||||
if isinstance(preds, (list, tuple)): # YOLOv8 model in validation model, output = (inference_out, loss_out)
|
||||
preds = preds[0]
|
||||
preds_c = preds.clone()
|
||||
|
||||
nms_cls = ops.non_max_suppression(
|
||||
preds[:, :8, :],
|
||||
self.args.conf,
|
||||
self.args.iou,
|
||||
labels=self.lb,
|
||||
multi_label=True,
|
||||
agnostic=self.args.single_cls,
|
||||
max_det=self.args.max_det,
|
||||
)
|
||||
|
||||
nms_obj = ops.non_max_suppression(
|
||||
torch.cat((preds_c[:, :4, :], preds_c[:, -1:, :]), dim=1),
|
||||
self.args.conf,
|
||||
self.args.iou,
|
||||
labels=self.lb,
|
||||
multi_label=False,
|
||||
agnostic=True,#self.args.single_cls,
|
||||
max_det=self.args.max_det,
|
||||
)
|
||||
|
||||
nms_out = nms_cls
|
||||
for i in range(0, len(nms_cls)):
|
||||
nms_obj[i][:, 5] = self.nc
|
||||
nms_out[i] = torch.cat((nms_cls[i], nms_obj[i]), dim=0)
|
||||
|
||||
return nms_out
|
||||
|
||||
def _prepare_batch(self, si, batch):
|
||||
"""Prepares a batch of images and annotations for validation."""
|
||||
idx = batch["batch_idx"] == si
|
||||
cls = batch["cls"][idx].squeeze(-1)
|
||||
bbox = batch["bboxes"][idx]
|
||||
ori_shape = batch["ori_shape"][si]
|
||||
imgsz = batch["img"].shape[2:]
|
||||
ratio_pad = batch["ratio_pad"][si]
|
||||
if len(cls):
|
||||
bbox = ops.xywh2xyxy(bbox) * torch.tensor(imgsz, device=self.device)[[1, 0, 1, 0]] # target boxes
|
||||
ops.scale_boxes(imgsz, bbox, ori_shape, ratio_pad=ratio_pad) # native-space labels
|
||||
return {"cls": cls, "bbox": bbox, "ori_shape": ori_shape, "imgsz": imgsz, "ratio_pad": ratio_pad}
|
||||
|
||||
def _prepare_pred(self, pred, pbatch):
|
||||
"""Prepares a batch of images and annotations for validation."""
|
||||
predn = pred.clone()
|
||||
ops.scale_boxes(
|
||||
pbatch["imgsz"], predn[:, :4], pbatch["ori_shape"], ratio_pad=pbatch["ratio_pad"]
|
||||
) # native-space pred
|
||||
return predn
|
||||
|
||||
def update_metrics(self, preds, batch):
|
||||
"""Metrics."""
|
||||
for si, pred in enumerate(preds):
|
||||
self.seen += 1
|
||||
npr = len(pred)
|
||||
stat = dict(
|
||||
conf=torch.zeros(0, device=self.device),
|
||||
pred_cls=torch.zeros(0, device=self.device),
|
||||
tp=torch.zeros(npr, self.niou, dtype=torch.bool, device=self.device),
|
||||
)
|
||||
pbatch = self._prepare_batch(si, batch)
|
||||
cls, bbox = pbatch.pop("cls"), pbatch.pop("bbox")
|
||||
cls = torch.cat((cls, torch.full(cls.shape, self.nc).to(cls.device)), dim=0)
|
||||
bbox = torch.cat((bbox, bbox), dim=0)
|
||||
nl = len(cls)
|
||||
stat["target_cls"] = cls
|
||||
if npr == 0:
|
||||
if nl:
|
||||
for k in self.stats.keys():
|
||||
self.stats[k].append(stat[k])
|
||||
if self.args.plots:
|
||||
self.confusion_matrix.process_batch(detections=None, gt_bboxes=bbox, gt_cls=cls)
|
||||
continue
|
||||
|
||||
# Predictions
|
||||
if self.args.single_cls:
|
||||
pred[:, 5] = 0
|
||||
predn = self._prepare_pred(pred, pbatch)
|
||||
stat["conf"] = predn[:, 4]
|
||||
stat["pred_cls"] = predn[:, 5]
|
||||
|
||||
# Evaluate
|
||||
if nl:
|
||||
stat["tp"] = self._process_batch(predn, bbox, cls)
|
||||
if self.args.plots:
|
||||
self.confusion_matrix.process_batch(predn, bbox, cls)
|
||||
for k in self.stats.keys():
|
||||
self.stats[k].append(stat[k])
|
||||
|
||||
# Save
|
||||
if self.args.save_json:
|
||||
self.pred_to_json(predn, batch["im_file"][si])
|
||||
if self.args.save_txt:
|
||||
file = self.save_dir / "labels" / f'{Path(batch["im_file"][si]).stem}.txt'
|
||||
self.save_one_txt(predn, self.args.save_conf, pbatch["ori_shape"], file)
|
||||
|
||||
def finalize_metrics(self, *args, **kwargs):
|
||||
"""Set final values for metrics speed and confusion matrix."""
|
||||
self.metrics.speed = self.speed
|
||||
self.metrics.confusion_matrix = self.confusion_matrix
|
||||
|
||||
def get_stats(self):
|
||||
"""Returns metrics statistics and results dictionary."""
|
||||
stats = {k: torch.cat(v, 0).cpu().numpy() for k, v in self.stats.items()} # to numpy
|
||||
if len(stats) and stats["tp"].any():
|
||||
self.metrics.process(**stats)
|
||||
self.nt_per_class = np.bincount(
|
||||
stats["target_cls"].astype(int), minlength=self.nc+1
|
||||
) # number of targets per class
|
||||
return self.metrics.results_dict
|
||||
|
||||
def print_results(self):
|
||||
"""Prints training/validation set metrics per class."""
|
||||
pf = "%22s" + "%11i" * 2 + "%11.4g" * len(self.metrics.keys) # print format
|
||||
LOGGER.info(pf % ("all", self.seen, self.nt_per_class.sum(), *self.metrics.mean_results()))
|
||||
if self.nt_per_class.sum() == 0:
|
||||
LOGGER.warning(f"WARNING ⚠️ no labels found in {self.args.task} set, can not compute metrics without labels")
|
||||
|
||||
# Print results per class
|
||||
if self.args.verbose and not self.training and self.nc > 1 and len(self.stats):
|
||||
for i, c in enumerate(self.metrics.ap_class_index):
|
||||
LOGGER.info(pf % (self.names[c], self.seen, self.nt_per_class[c], *self.metrics.class_result(i)))
|
||||
|
||||
if self.args.plots:
|
||||
for normalize in True, False:
|
||||
self.confusion_matrix.plot(
|
||||
save_dir=self.save_dir, names=self.names.values(), normalize=normalize, on_plot=self.on_plot
|
||||
)
|
||||
|
||||
def _process_batch(self, detections, gt_bboxes, gt_cls):
|
||||
"""
|
||||
Return correct prediction matrix.
|
||||
|
||||
Args:
|
||||
detections (torch.Tensor): Tensor of shape [N, 6] representing detections.
|
||||
Each detection is of the format: x1, y1, x2, y2, conf, class.
|
||||
labels (torch.Tensor): Tensor of shape [M, 5] representing labels.
|
||||
Each label is of the format: class, x1, y1, x2, y2.
|
||||
|
||||
Returns:
|
||||
(torch.Tensor): Correct prediction matrix of shape [N, 10] for 10 IoU levels.
|
||||
"""
|
||||
iou = box_iou(gt_bboxes, detections[:, :4])
|
||||
return self.match_predictions(detections[:, 5], gt_cls, iou)
|
||||
|
||||
def build_dataset(self, img_path, mode="val", batch=None):
|
||||
"""
|
||||
Build YOLO Dataset.
|
||||
|
||||
Args:
|
||||
img_path (str): Path to the folder containing images.
|
||||
mode (str): `train` mode or `val` mode, users are able to customize different augmentations for each mode.
|
||||
batch (int, optional): Size of batches, this is for `rect`. Defaults to None.
|
||||
"""
|
||||
return build_yolo_dataset(self.args, img_path, batch, self.data, mode=mode, stride=self.stride)
|
||||
|
||||
def get_dataloader(self, dataset_path, batch_size):
|
||||
"""Construct and return dataloader."""
|
||||
dataset = self.build_dataset(dataset_path, batch=batch_size, mode="val")
|
||||
return build_dataloader(dataset, batch_size, self.args.workers, shuffle=False, rank=-1) # return dataloader
|
||||
|
||||
def plot_val_samples(self, batch, ni):
|
||||
"""Plot validation image samples."""
|
||||
plot_images(
|
||||
batch["img"],
|
||||
batch["batch_idx"],
|
||||
batch["cls"].squeeze(-1),
|
||||
batch["bboxes"],
|
||||
paths=batch["im_file"],
|
||||
fname=self.save_dir / f"val_batch{ni}_labels.jpg",
|
||||
names=self.names,
|
||||
on_plot=self.on_plot,
|
||||
)
|
||||
|
||||
def plot_predictions(self, batch, preds, ni):
|
||||
"""Plots predicted bounding boxes on input images and saves the result."""
|
||||
plot_images(
|
||||
batch["img"],
|
||||
*output_to_target(preds, max_det=self.args.max_det),
|
||||
paths=batch["im_file"],
|
||||
fname=self.save_dir / f"val_batch{ni}_pred.jpg",
|
||||
names=self.names,
|
||||
on_plot=self.on_plot,
|
||||
) # pred
|
||||
|
||||
def save_one_txt(self, predn, save_conf, shape, file):
|
||||
"""Save YOLO detections to a txt file in normalized coordinates in a specific format."""
|
||||
gn = torch.tensor(shape)[[1, 0, 1, 0]] # normalization gain whwh
|
||||
for *xyxy, conf, cls in predn.tolist():
|
||||
xywh = (ops.xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
|
||||
line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
|
||||
with open(file, "a") as f:
|
||||
f.write(("%g " * len(line)).rstrip() % line + "\n")
|
||||
|
||||
def pred_to_json(self, predn, filename):
|
||||
"""Serialize YOLO predictions to COCO json format."""
|
||||
stem = Path(filename).stem
|
||||
image_id = int(stem) if stem.isnumeric() else stem
|
||||
box = ops.xyxy2xywh(predn[:, :4]) # xywh
|
||||
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
|
||||
for p, b in zip(predn.tolist(), box.tolist()):
|
||||
self.jdict.append(
|
||||
{
|
||||
"image_id": image_id,
|
||||
"category_id": self.class_map[int(p[5])]
|
||||
+ (1 if self.is_lvis else 0), # index starts from 1 if it's lvis
|
||||
"bbox": [round(x, 3) for x in b],
|
||||
"score": round(p[4], 5),
|
||||
}
|
||||
)
|
||||
|
||||
def eval_json(self, stats):
|
||||
"""Evaluates YOLO output in JSON format and returns performance statistics."""
|
||||
if self.args.save_json and (self.is_coco or self.is_lvis) and len(self.jdict):
|
||||
pred_json = self.save_dir / "predictions.json" # predictions
|
||||
anno_json = (
|
||||
self.data["path"]
|
||||
/ "annotations"
|
||||
/ ("instances_val2017.json" if self.is_coco else f"lvis_v1_{self.args.split}.json")
|
||||
) # annotations
|
||||
pkg = "pycocotools" if self.is_coco else "lvis"
|
||||
LOGGER.info(f"\nEvaluating {pkg} mAP using {pred_json} and {anno_json}...")
|
||||
try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
|
||||
for x in pred_json, anno_json:
|
||||
assert x.is_file(), f"{x} file not found"
|
||||
check_requirements("pycocotools>=2.0.6" if self.is_coco else "lvis>=0.5.3")
|
||||
if self.is_coco:
|
||||
from pycocotools.coco import COCO # noqa
|
||||
from pycocotools.cocoeval import COCOeval # noqa
|
||||
|
||||
anno = COCO(str(anno_json)) # init annotations api
|
||||
pred = anno.loadRes(str(pred_json)) # init predictions api (must pass string, not Path)
|
||||
eval = COCOeval(anno, pred, "bbox")
|
||||
else:
|
||||
from lvis import LVIS, LVISEval
|
||||
|
||||
anno = LVIS(str(anno_json)) # init annotations api
|
||||
pred = anno._load_json(str(pred_json)) # init predictions api (must pass string, not Path)
|
||||
eval = LVISEval(anno, pred, "bbox")
|
||||
eval.params.imgIds = [int(Path(x).stem) for x in self.dataloader.dataset.im_files] # images to eval
|
||||
eval.evaluate()
|
||||
eval.accumulate()
|
||||
eval.summarize()
|
||||
if self.is_lvis:
|
||||
eval.print_results() # explicitly call print_results
|
||||
# update mAP50-95 and mAP50
|
||||
stats[self.metrics.keys[-1]], stats[self.metrics.keys[-2]], stats[self.metrics.keys[-3]] = (
|
||||
eval.stats[:3] if self.is_coco else [eval.results["AP50"], eval.results["AP75"], eval.results["AP"]]
|
||||
)
|
||||
except Exception as e:
|
||||
LOGGER.warning(f"{pkg} unable to run: {e}")
|
||||
return stats
|
||||
Reference in new issue
Block a user