diff --git a/ultralytics/yolo/engine/model.py b/ultralytics/yolo/engine/model.py index 80e3c23..5a57022 100644 --- a/ultralytics/yolo/engine/model.py +++ b/ultralytics/yolo/engine/model.py @@ -118,6 +118,7 @@ class YOLO: **kwargs : Any other args accepted by the predictors. To see all args check 'configuration' section in docs """ overrides = self.overrides.copy() + overrides["conf"] = 0.25 overrides.update(kwargs) overrides["mode"] = "predict" predictor = self.PredictorClass(overrides=overrides) diff --git a/ultralytics/yolo/v8/detect/train.py b/ultralytics/yolo/v8/detect/train.py index 1e3dd18..acf2b5a 100644 --- a/ultralytics/yolo/v8/detect/train.py +++ b/ultralytics/yolo/v8/detect/train.py @@ -186,9 +186,9 @@ class Loss: loss[0], loss[2] = self.bbox_loss(pred_distri, pred_bboxes, anchor_points, target_bboxes, target_scores, target_scores_sum, fg_mask) - loss[0] *= 7.5 # box gain - loss[1] *= 0.5 # cls gain - loss[2] *= 1.5 # dfl gain + loss[0] *= self.hyp.box # box gain + loss[1] *= self.hyp.cls # cls gain + loss[2] *= self.hyp.dfl # dfl gain return loss.sum() * batch_size, loss.detach() # loss(box, cls, dfl) diff --git a/ultralytics/yolo/v8/segment/train.py b/ultralytics/yolo/v8/segment/train.py index 7e32bed..3242154 100644 --- a/ultralytics/yolo/v8/segment/train.py +++ b/ultralytics/yolo/v8/segment/train.py @@ -2,18 +2,18 @@ from copy import copy import hydra import torch -import torch.nn as nn import torch.nn.functional as F from ultralytics.nn.tasks import SegmentationModel from ultralytics.yolo import v8 from ultralytics.yolo.utils import DEFAULT_CONFIG -from ultralytics.yolo.utils.loss import BboxLoss -from ultralytics.yolo.utils.ops import crop_mask, xywh2xyxy, xyxy2xywh +from ultralytics.yolo.utils.ops import crop_mask, xyxy2xywh from ultralytics.yolo.utils.plotting import plot_images, plot_results -from ultralytics.yolo.utils.tal import TaskAlignedAssigner, dist2bbox, make_anchors +from ultralytics.yolo.utils.tal import make_anchors from ultralytics.yolo.utils.torch_utils import de_parallel +from ..detect.train import Loss + # BaseTrainer python usage class SegmentationTrainer(v8.detect.DetectionTrainer): @@ -55,51 +55,12 @@ class SegmentationTrainer(v8.detect.DetectionTrainer): # Criterion class for computing training losses -class SegLoss: +class SegLoss(Loss): def __init__(self, model, overlap=True): # model must be de-paralleled - - device = next(model.parameters()).device # get model device - h = model.args # hyperparameters - - m = model.model[-1] # Detect() module - self.bce = nn.BCEWithLogitsLoss(reduction='none') - self.hyp = h - self.stride = m.stride # model strides - self.nc = m.nc # number of classes - self.no = m.no - self.nm = m.nm # number of masks - self.reg_max = m.reg_max + super().__init__(model) + self.nm = model.model[-1].nm # number of masks self.overlap = overlap - self.device = device - - self.use_dfl = m.reg_max > 1 - self.assigner = TaskAlignedAssigner(topk=10, num_classes=self.nc, alpha=0.5, beta=6.0) - self.bbox_loss = BboxLoss(m.reg_max - 1, use_dfl=self.use_dfl).to(device) - self.proj = torch.arange(m.reg_max, dtype=torch.float, device=device) - - def preprocess(self, targets, batch_size, scale_tensor): - if targets.shape[0] == 0: - out = torch.zeros(batch_size, 0, 5, device=self.device) - else: - i = targets[:, 0] # image index - _, counts = i.unique(return_counts=True) - out = torch.zeros(batch_size, counts.max(), 5, device=self.device) - for j in range(batch_size): - matches = i == j - n = matches.sum() - if n: - out[j, :n] = targets[matches, 1:] - out[..., 1:5] = xywh2xyxy(out[..., 1:5].mul_(scale_tensor)) - return out - - def bbox_decode(self, anchor_points, pred_dist): - if self.use_dfl: - b, a, c = pred_dist.shape # batch, anchors, channels - pred_dist = pred_dist.view(b, a, 4, c // 4).softmax(3).matmul(self.proj.type(pred_dist.dtype)) - # pred_dist = pred_dist.view(b, a, c // 4, 4).transpose(2,3).softmax(3).matmul(self.proj.type(pred_dist.dtype)) - # pred_dist = (pred_dist.view(b, a, c // 4, 4).softmax(2) * self.proj.type(pred_dist.dtype).view(1, 1, -1, 1)).sum(2) - return dist2bbox(pred_dist, anchor_points, xywh=False) def __call__(self, preds, batch): loss = torch.zeros(4, device=self.device) # box, cls, dfl @@ -163,10 +124,10 @@ class SegLoss: # else: # loss[1] += proto.sum() * 0 - loss[0] *= 7.5 # box gain - loss[1] *= 7.5 / batch_size # seg gain - loss[2] *= 0.5 # cls gain - loss[3] *= 1.5 # dfl gain + loss[0] *= self.hyp.box # box gain + loss[1] *= self.hyp.box / batch_size # seg gain + loss[2] *= self.hyp.cls # cls gain + loss[3] *= self.hyp.dfl # dfl gain return loss.sum() * batch_size, loss.detach() # loss(box, cls, dfl)