Clean validator (#144)

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
2023-01-05 21:14:01 +08:00
committed by GitHub
parent 46cb657b64
commit 19334ebb16
9 changed files with 88 additions and 74 deletions

View File

@ -22,17 +22,8 @@ class SegmentationValidator(DetectionValidator):
self.metrics = SegmentMetrics(save_dir=self.save_dir, plot=self.args.plots)
def preprocess(self, batch):
batch["img"] = batch["img"].to(self.device, non_blocking=True)
batch["img"] = (batch["img"].half() if self.args.half else batch["img"].float()) / 255
batch = super().preprocess(batch)
batch["masks"] = batch["masks"].to(self.device).float()
self.nb, _, self.height, self.width = batch["img"].shape # batch size, channels, height, width
self.targets = torch.cat((batch["batch_idx"].view(-1, 1), batch["cls"].view(-1, 1), batch["bboxes"]), 1)
self.targets = self.targets.to(self.device)
height, width = batch["img"].shape[2:]
self.targets[:, 2:] *= torch.tensor((width, height, width, height), device=self.device) # to pixels
self.lb = [self.targets[self.targets[:, 0] == i, 1:]
for i in range(self.nb)] if self.args.save_hybrid else [] # for autolabelling
return batch
def init_metrics(self, model):
@ -72,10 +63,11 @@ class SegmentationValidator(DetectionValidator):
def update_metrics(self, preds, batch):
# Metrics
for si, (pred, proto) in enumerate(zip(preds[0], preds[1])):
labels = self.targets[self.targets[:, 0] == si, 1:]
nl, npr = labels.shape[0], pred.shape[0] # number of labels, predictions
idx = batch["batch_idx"] == si
cls = batch["cls"][idx]
bbox = batch["bboxes"][idx]
nl, npr = cls.shape[0], pred.shape[0] # number of labels, predictions
shape = batch["ori_shape"][si]
# path = batch["shape"][si][0]
correct_masks = torch.zeros(npr, self.niou, dtype=torch.bool, device=self.device) # init
correct_bboxes = torch.zeros(npr, self.niou, dtype=torch.bool, device=self.device) # init
self.seen += 1
@ -83,13 +75,13 @@ class SegmentationValidator(DetectionValidator):
if npr == 0:
if nl:
self.stats.append((correct_masks, correct_bboxes, *torch.zeros(
(2, 0), device=self.device), labels[:, 0]))
(2, 0), device=self.device), cls.squeeze(-1)))
if self.args.plots:
self.confusion_matrix.process_batch(detections=None, labels=labels[:, 0])
self.confusion_matrix.process_batch(detections=None, labels=cls.squeeze(-1))
continue
# Masks
midx = [si] if self.args.overlap_mask else self.targets[:, 0] == si
midx = [si] if self.args.overlap_mask else idx
gt_masks = batch["masks"][midx]
pred_masks = self.process(proto, pred[:, 6:], pred[:, :4], shape=batch["img"][si].shape[1:])
@ -101,9 +93,9 @@ class SegmentationValidator(DetectionValidator):
# Evaluate
if nl:
tbox = ops.xywh2xyxy(labels[:, 1:5]) # target boxes
tbox = ops.xywh2xyxy(bbox) # target boxes
ops.scale_boxes(batch["img"][si].shape[1:], tbox, shape) # native-space labels
labelsn = torch.cat((labels[:, 0:1], tbox), 1) # native-space labels
labelsn = torch.cat((cls, tbox), 1) # native-space labels
correct_bboxes = self._process_batch(predn, labelsn)
# TODO: maybe remove these `self.` arguments as they already are member variable
correct_masks = self._process_batch(predn,
@ -114,7 +106,8 @@ class SegmentationValidator(DetectionValidator):
masks=True)
if self.args.plots:
self.confusion_matrix.process_batch(predn, labelsn)
self.stats.append((correct_masks, correct_bboxes, pred[:, 4], pred[:, 5], labels[:, 0])) # conf, pcls, tcls
self.stats.append((correct_masks, correct_bboxes, pred[:, 4], pred[:,
5], cls.squeeze(-1))) # conf, pcls, tcls
pred_masks = torch.as_tensor(pred_masks, dtype=torch.uint8)
if self.args.plots and self.batch_i < 3:
@ -165,19 +158,6 @@ class SegmentationValidator(DetectionValidator):
correct[matches[:, 1].astype(int), i] = True
return torch.tensor(correct, dtype=torch.bool, device=detections.device)
# TODO: probably add this to class Metrics
@property
def metric_keys(self):
return [
"metrics/precision(B)",
"metrics/recall(B)",
"metrics/mAP50(B)",
"metrics/mAP50-95(B)", # metrics
"metrics/precision(M)",
"metrics/recall(M)",
"metrics/mAP50(M)",
"metrics/mAP50-95(M)",]
def plot_val_samples(self, batch, ni):
plot_images(batch["img"],
batch["batch_idx"],
@ -243,8 +223,8 @@ class SegmentationValidator(DetectionValidator):
eval.accumulate()
eval.summarize()
idx = i * 4 + 2
stats[self.metric_keys[idx + 1]], stats[
self.metric_keys[idx]] = eval.stats[:2] # update mAP50-95 and mAP50
stats[self.metrics.keys[idx + 1]], stats[
self.metrics.keys[idx]] = eval.stats[:2] # update mAP50-95 and mAP50
except Exception as e:
self.logger.warning(f'pycocotools unable to run: {e}')
return stats