Segmentation support & other enchancements (#40)
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:
Ayush Chaurasia
GitHub
@ -4,6 +4,7 @@
|
||||
|
||||
# Train settings -------------------------------------------------------------------------------------------------------
|
||||
model: null # i.e. yolov5s.pt
|
||||
cfg: null # i.e. yolov5s.yaml
|
||||
data: null # i.e. coco128.yaml
|
||||
epochs: 300
|
||||
batch_size: 16
|
||||
@ -20,6 +21,23 @@ optimizer: 'SGD' # choices=['SGD', 'Adam', 'AdamW', 'RMSProp']
|
||||
verbose: False
|
||||
seed: 0
|
||||
local_rank: -1
|
||||
single_cls: False # train multi-class data as single-class
|
||||
image_weights: False # use weighted image selection for training
|
||||
shuffle: True
|
||||
rect: False # support rectangular training
|
||||
overlap_mask: True # Segmentation masks overlap
|
||||
mask_ratio: 4 # Segmentation mask downsample ratio
|
||||
|
||||
# Val/Test settings ----------------------------------------------------------------------------------------------------
|
||||
save_json: False
|
||||
save_hybrid: False
|
||||
conf_thres: 0.001
|
||||
iou_thres: 0.6
|
||||
max_det: 300
|
||||
half: True
|
||||
plots: False
|
||||
save_txt: False
|
||||
task: 'val'
|
||||
|
||||
# Hyperparameters ------------------------------------------------------------------------------------------------------
|
||||
lr0: 0.001 # initial learning rate (SGD=1E-2, Adam=1E-3)
|
||||
@ -51,6 +69,7 @@ fliplr: 0.5 # image flip left-right (probability)
|
||||
mosaic: 1.0 # image mosaic (probability)
|
||||
mixup: 0.0 # image mixup (probability)
|
||||
copy_paste: 0.0 # segment copy-paste (probability)
|
||||
label_smoothing: 0.0
|
||||
|
||||
# Hydra configs --------------------------------------------------------------------------------------------------------
|
||||
hydra:
|
||||
|
||||
@ -2,11 +2,19 @@
|
||||
"""
|
||||
Model validation metrics
|
||||
"""
|
||||
import math
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from ultralytics.yolo.utils import TryExcept
|
||||
|
||||
|
||||
# boxes
|
||||
def box_area(box):
|
||||
# box = xyxy(4,n)
|
||||
return (box[2] - box[0]) * (box[3] - box[1])
|
||||
@ -53,3 +61,484 @@ def box_iou(box1, box2, eps=1e-7):
|
||||
|
||||
# IoU = inter / (area1 + area2 - inter)
|
||||
return inter / (box_area(box1.T)[:, None] + box_area(box2.T) - inter + eps)
|
||||
|
||||
|
||||
def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7):
|
||||
# Returns Intersection over Union (IoU) of box1(1,4) to box2(n,4)
|
||||
|
||||
# Get the coordinates of bounding boxes
|
||||
if xywh: # transform from xywh to xyxy
|
||||
(x1, y1, w1, h1), (x2, y2, w2, h2) = box1.chunk(4, 1), box2.chunk(4, 1)
|
||||
w1_, h1_, w2_, h2_ = w1 / 2, h1 / 2, w2 / 2, h2 / 2
|
||||
b1_x1, b1_x2, b1_y1, b1_y2 = x1 - w1_, x1 + w1_, y1 - h1_, y1 + h1_
|
||||
b2_x1, b2_x2, b2_y1, b2_y2 = x2 - w2_, x2 + w2_, y2 - h2_, y2 + h2_
|
||||
else: # x1, y1, x2, y2 = box1
|
||||
b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, 1)
|
||||
b2_x1, b2_y1, b2_x2, b2_y2 = box2.chunk(4, 1)
|
||||
w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
|
||||
w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
|
||||
|
||||
# Intersection area
|
||||
inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
|
||||
(torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)
|
||||
|
||||
# Union Area
|
||||
union = w1 * h1 + w2 * h2 - inter + eps
|
||||
|
||||
# IoU
|
||||
iou = inter / union
|
||||
if CIoU or DIoU or GIoU:
|
||||
cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) # convex (smallest enclosing box) width
|
||||
ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height
|
||||
if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
|
||||
c2 = cw ** 2 + ch ** 2 + eps # convex diagonal squared
|
||||
rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # center dist ** 2
|
||||
if CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
|
||||
v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
|
||||
with torch.no_grad():
|
||||
alpha = v / (v - iou + (1 + eps))
|
||||
return iou - (rho2 / c2 + v * alpha) # CIoU
|
||||
return iou - rho2 / c2 # DIoU
|
||||
c_area = cw * ch + eps # convex area
|
||||
return iou - (c_area - union) / c_area # GIoU https://arxiv.org/pdf/1902.09630.pdf
|
||||
return iou # IoU
|
||||
|
||||
|
||||
def mask_iou(mask1, mask2, eps=1e-7):
|
||||
"""
|
||||
mask1: [N, n] m1 means number of predicted objects
|
||||
mask2: [M, n] m2 means number of gt objects
|
||||
Note: n means image_w x image_h
|
||||
return: masks iou, [N, M]
|
||||
"""
|
||||
intersection = torch.matmul(mask1, mask2.t()).clamp(0)
|
||||
union = (mask1.sum(1)[:, None] + mask2.sum(1)[None]) - intersection # (area1 + area2) - intersection
|
||||
return intersection / (union + eps)
|
||||
|
||||
|
||||
def masks_iou(mask1, mask2, eps=1e-7):
|
||||
"""
|
||||
mask1: [N, n] m1 means number of predicted objects
|
||||
mask2: [N, n] m2 means number of gt objects
|
||||
Note: n means image_w x image_h
|
||||
return: masks iou, (N, )
|
||||
"""
|
||||
intersection = (mask1 * mask2).sum(1).clamp(0) # (N, )
|
||||
union = (mask1.sum(1) + mask2.sum(1))[None] - intersection # (area1 + area2) - intersection
|
||||
return intersection / (union + eps)
|
||||
|
||||
|
||||
def smooth_BCE(eps=0.1): # https://github.com/ultralytics/yolov3/issues/238#issuecomment-598028441
|
||||
# return positive, negative label smoothing BCE targets
|
||||
return 1.0 - 0.5 * eps, 0.5 * eps
|
||||
|
||||
|
||||
# losses
|
||||
class FocalLoss(nn.Module):
|
||||
# Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
|
||||
def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
|
||||
super().__init__()
|
||||
self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss()
|
||||
self.gamma = gamma
|
||||
self.alpha = alpha
|
||||
self.reduction = loss_fcn.reduction
|
||||
self.loss_fcn.reduction = 'none' # required to apply FL to each element
|
||||
|
||||
def forward(self, pred, true):
|
||||
loss = self.loss_fcn(pred, true)
|
||||
# p_t = torch.exp(-loss)
|
||||
# loss *= self.alpha * (1.000001 - p_t) ** self.gamma # non-zero power for gradient stability
|
||||
|
||||
# TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py
|
||||
pred_prob = torch.sigmoid(pred) # prob from logits
|
||||
p_t = true * pred_prob + (1 - true) * (1 - pred_prob)
|
||||
alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
|
||||
modulating_factor = (1.0 - p_t) ** self.gamma
|
||||
loss *= alpha_factor * modulating_factor
|
||||
|
||||
if self.reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss.sum()
|
||||
else: # 'none'
|
||||
return loss
|
||||
|
||||
|
||||
class ConfusionMatrix:
|
||||
# Updated version of https://github.com/kaanakan/object_detection_confusion_matrix
|
||||
def __init__(self, nc, conf=0.25, iou_thres=0.45):
|
||||
self.matrix = np.zeros((nc + 1, nc + 1))
|
||||
self.nc = nc # number of classes
|
||||
self.conf = conf
|
||||
self.iou_thres = iou_thres
|
||||
|
||||
def process_batch(self, detections, labels):
|
||||
"""
|
||||
Return intersection-over-union (Jaccard index) of boxes.
|
||||
Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
|
||||
Arguments:
|
||||
detections (Array[N, 6]), x1, y1, x2, y2, conf, class
|
||||
labels (Array[M, 5]), class, x1, y1, x2, y2
|
||||
Returns:
|
||||
None, updates confusion matrix accordingly
|
||||
"""
|
||||
if detections is None:
|
||||
gt_classes = labels.int()
|
||||
for gc in gt_classes:
|
||||
self.matrix[self.nc, gc] += 1 # background FN
|
||||
return
|
||||
|
||||
detections = detections[detections[:, 4] > self.conf]
|
||||
gt_classes = labels[:, 0].int()
|
||||
detection_classes = detections[:, 5].int()
|
||||
iou = box_iou(labels[:, 1:], detections[:, :4])
|
||||
|
||||
x = torch.where(iou > self.iou_thres)
|
||||
if x[0].shape[0]:
|
||||
matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy()
|
||||
if x[0].shape[0] > 1:
|
||||
matches = matches[matches[:, 2].argsort()[::-1]]
|
||||
matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
|
||||
matches = matches[matches[:, 2].argsort()[::-1]]
|
||||
matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
|
||||
else:
|
||||
matches = np.zeros((0, 3))
|
||||
|
||||
n = matches.shape[0] > 0
|
||||
m0, m1, _ = matches.transpose().astype(int)
|
||||
for i, gc in enumerate(gt_classes):
|
||||
j = m0 == i
|
||||
if n and sum(j) == 1:
|
||||
self.matrix[detection_classes[m1[j]], gc] += 1 # correct
|
||||
else:
|
||||
self.matrix[self.nc, gc] += 1 # true background
|
||||
|
||||
if n:
|
||||
for i, dc in enumerate(detection_classes):
|
||||
if not any(m1 == i):
|
||||
self.matrix[dc, self.nc] += 1 # predicted background
|
||||
|
||||
def matrix(self):
|
||||
return self.matrix
|
||||
|
||||
def tp_fp(self):
|
||||
tp = self.matrix.diagonal() # true positives
|
||||
fp = self.matrix.sum(1) - tp # false positives
|
||||
# fn = self.matrix.sum(0) - tp # false negatives (missed detections)
|
||||
return tp[:-1], fp[:-1] # remove background class
|
||||
|
||||
@TryExcept('WARNING ⚠️ ConfusionMatrix plot failure')
|
||||
def plot(self, normalize=True, save_dir='', names=()):
|
||||
import seaborn as sn
|
||||
|
||||
array = self.matrix / ((self.matrix.sum(0).reshape(1, -1) + 1E-9) if normalize else 1) # normalize columns
|
||||
array[array < 0.005] = np.nan # don't annotate (would appear as 0.00)
|
||||
|
||||
fig, ax = plt.subplots(1, 1, figsize=(12, 9), tight_layout=True)
|
||||
nc, nn = self.nc, len(names) # number of classes, names
|
||||
sn.set(font_scale=1.0 if nc < 50 else 0.8) # for label size
|
||||
labels = (0 < nn < 99) and (nn == nc) # apply names to ticklabels
|
||||
ticklabels = (names + ['background']) if labels else "auto"
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter('ignore') # suppress empty matrix RuntimeWarning: All-NaN slice encountered
|
||||
sn.heatmap(array,
|
||||
ax=ax,
|
||||
annot=nc < 30,
|
||||
annot_kws={
|
||||
"size": 8},
|
||||
cmap='Blues',
|
||||
fmt='.2f',
|
||||
square=True,
|
||||
vmin=0.0,
|
||||
xticklabels=ticklabels,
|
||||
yticklabels=ticklabels).set_facecolor((1, 1, 1))
|
||||
ax.set_ylabel('True')
|
||||
ax.set_ylabel('Predicted')
|
||||
ax.set_title('Confusion Matrix')
|
||||
fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250)
|
||||
plt.close(fig)
|
||||
|
||||
def print(self):
|
||||
for i in range(self.nc + 1):
|
||||
print(' '.join(map(str, self.matrix[i])))
|
||||
|
||||
|
||||
def fitness_detection(x):
|
||||
# Model fitness as a weighted combination of metrics
|
||||
w = [0.0, 0.0, 0.1, 0.9] # weights for [P, R, mAP@0.5, mAP@0.5:0.95]
|
||||
return (x[:, :4] * w).sum(1)
|
||||
|
||||
|
||||
def fitness_segmentation(x):
|
||||
# Model fitness as a weighted combination of metrics
|
||||
w = [0.0, 0.0, 0.1, 0.9, 0.0, 0.0, 0.1, 0.9]
|
||||
return (x[:, :8] * w).sum(1)
|
||||
|
||||
|
||||
def smooth(y, f=0.05):
|
||||
# Box filter of fraction f
|
||||
nf = round(len(y) * f * 2) // 2 + 1 # number of filter elements (must be odd)
|
||||
p = np.ones(nf // 2) # ones padding
|
||||
yp = np.concatenate((p * y[0], y, p * y[-1]), 0) # y padded
|
||||
return np.convolve(yp, np.ones(nf) / nf, mode='valid') # y-smoothed
|
||||
|
||||
|
||||
def compute_ap(recall, precision):
|
||||
""" Compute the average precision, given the recall and precision curves
|
||||
# Arguments
|
||||
recall: The recall curve (list)
|
||||
precision: The precision curve (list)
|
||||
# Returns
|
||||
Average precision, precision curve, recall curve
|
||||
"""
|
||||
|
||||
# Append sentinel values to beginning and end
|
||||
mrec = np.concatenate(([0.0], recall, [1.0]))
|
||||
mpre = np.concatenate(([1.0], precision, [0.0]))
|
||||
|
||||
# Compute the precision envelope
|
||||
mpre = np.flip(np.maximum.accumulate(np.flip(mpre)))
|
||||
|
||||
# Integrate area under curve
|
||||
method = 'interp' # methods: 'continuous', 'interp'
|
||||
if method == 'interp':
|
||||
x = np.linspace(0, 1, 101) # 101-point interp (COCO)
|
||||
ap = np.trapz(np.interp(x, mrec, mpre), x) # integrate
|
||||
else: # 'continuous'
|
||||
i = np.where(mrec[1:] != mrec[:-1])[0] # points where x axis (recall) changes
|
||||
ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1]) # area under curve
|
||||
|
||||
return ap, mpre, mrec
|
||||
|
||||
|
||||
def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names=(), eps=1e-16, prefix=""):
|
||||
""" Compute the average precision, given the recall and precision curves.
|
||||
Source: https://github.com/rafaelpadilla/Object-Detection-Metrics.
|
||||
# Arguments
|
||||
tp: True positives (nparray, nx1 or nx10).
|
||||
conf: Objectness value from 0-1 (nparray).
|
||||
pred_cls: Predicted object classes (nparray).
|
||||
target_cls: True object classes (nparray).
|
||||
plot: Plot precision-recall curve at mAP@0.5
|
||||
save_dir: Plot save directory
|
||||
# Returns
|
||||
The average precision as computed in py-faster-rcnn.
|
||||
"""
|
||||
|
||||
# Sort by objectness
|
||||
i = np.argsort(-conf)
|
||||
tp, conf, pred_cls = tp[i], conf[i], pred_cls[i]
|
||||
|
||||
# Find unique classes
|
||||
unique_classes, nt = np.unique(target_cls, return_counts=True)
|
||||
nc = unique_classes.shape[0] # number of classes, number of detections
|
||||
|
||||
# Create Precision-Recall curve and compute AP for each class
|
||||
px, py = np.linspace(0, 1, 1000), [] # for plotting
|
||||
ap, p, r = np.zeros((nc, tp.shape[1])), np.zeros((nc, 1000)), np.zeros((nc, 1000))
|
||||
for ci, c in enumerate(unique_classes):
|
||||
i = pred_cls == c
|
||||
n_l = nt[ci] # number of labels
|
||||
n_p = i.sum() # number of predictions
|
||||
if n_p == 0 or n_l == 0:
|
||||
continue
|
||||
|
||||
# Accumulate FPs and TPs
|
||||
fpc = (1 - tp[i]).cumsum(0)
|
||||
tpc = tp[i].cumsum(0)
|
||||
|
||||
# Recall
|
||||
recall = tpc / (n_l + eps) # recall curve
|
||||
r[ci] = np.interp(-px, -conf[i], recall[:, 0], left=0) # negative x, xp because xp decreases
|
||||
|
||||
# Precision
|
||||
precision = tpc / (tpc + fpc) # precision curve
|
||||
p[ci] = np.interp(-px, -conf[i], precision[:, 0], left=1) # p at pr_score
|
||||
|
||||
# AP from recall-precision curve
|
||||
for j in range(tp.shape[1]):
|
||||
ap[ci, j], mpre, mrec = compute_ap(recall[:, j], precision[:, j])
|
||||
if plot and j == 0:
|
||||
py.append(np.interp(px, mrec, mpre)) # precision at mAP@0.5
|
||||
|
||||
# Compute F1 (harmonic mean of precision and recall)
|
||||
f1 = 2 * p * r / (p + r + eps)
|
||||
names = [v for k, v in names.items() if k in unique_classes] # list: only classes that have data
|
||||
names = dict(enumerate(names)) # to dict
|
||||
# TODO: plot
|
||||
'''
|
||||
if plot:
|
||||
plot_pr_curve(px, py, ap, Path(save_dir) / f'{prefix}PR_curve.png', names)
|
||||
plot_mc_curve(px, f1, Path(save_dir) / f'{prefix}F1_curve.png', names, ylabel='F1')
|
||||
plot_mc_curve(px, p, Path(save_dir) / f'{prefix}P_curve.png', names, ylabel='Precision')
|
||||
plot_mc_curve(px, r, Path(save_dir) / f'{prefix}R_curve.png', names, ylabel='Recall')
|
||||
'''
|
||||
|
||||
i = smooth(f1.mean(0), 0.1).argmax() # max F1 index
|
||||
p, r, f1 = p[:, i], r[:, i], f1[:, i]
|
||||
tp = (r * nt).round() # true positives
|
||||
fp = (tp / (p + eps) - tp).round() # false positives
|
||||
return tp, fp, p, r, f1, ap, unique_classes.astype(int)
|
||||
|
||||
|
||||
def ap_per_class_box_and_mask(
|
||||
tp_m,
|
||||
tp_b,
|
||||
conf,
|
||||
pred_cls,
|
||||
target_cls,
|
||||
plot=False,
|
||||
save_dir=".",
|
||||
names=(),
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
tp_b: tp of boxes.
|
||||
tp_m: tp of masks.
|
||||
other arguments see `func: ap_per_class`.
|
||||
"""
|
||||
results_boxes = ap_per_class(tp_b,
|
||||
conf,
|
||||
pred_cls,
|
||||
target_cls,
|
||||
plot=plot,
|
||||
save_dir=save_dir,
|
||||
names=names,
|
||||
prefix="Box")[2:]
|
||||
results_masks = ap_per_class(tp_m,
|
||||
conf,
|
||||
pred_cls,
|
||||
target_cls,
|
||||
plot=plot,
|
||||
save_dir=save_dir,
|
||||
names=names,
|
||||
prefix="Mask")[2:]
|
||||
|
||||
results = {
|
||||
"boxes": {
|
||||
"p": results_boxes[0],
|
||||
"r": results_boxes[1],
|
||||
"ap": results_boxes[3],
|
||||
"f1": results_boxes[2],
|
||||
"ap_class": results_boxes[4]},
|
||||
"masks": {
|
||||
"p": results_masks[0],
|
||||
"r": results_masks[1],
|
||||
"ap": results_masks[3],
|
||||
"f1": results_masks[2],
|
||||
"ap_class": results_masks[4]}}
|
||||
return results
|
||||
|
||||
|
||||
class Metric:
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.p = [] # (nc, )
|
||||
self.r = [] # (nc, )
|
||||
self.f1 = [] # (nc, )
|
||||
self.all_ap = [] # (nc, 10)
|
||||
self.ap_class_index = [] # (nc, )
|
||||
|
||||
@property
|
||||
def ap50(self):
|
||||
"""AP@0.5 of all classes.
|
||||
Return:
|
||||
(nc, ) or [].
|
||||
"""
|
||||
return self.all_ap[:, 0] if len(self.all_ap) else []
|
||||
|
||||
@property
|
||||
def ap(self):
|
||||
"""AP@0.5:0.95
|
||||
Return:
|
||||
(nc, ) or [].
|
||||
"""
|
||||
return self.all_ap.mean(1) if len(self.all_ap) else []
|
||||
|
||||
@property
|
||||
def mp(self):
|
||||
"""mean precision of all classes.
|
||||
Return:
|
||||
float.
|
||||
"""
|
||||
return self.p.mean() if len(self.p) else 0.0
|
||||
|
||||
@property
|
||||
def mr(self):
|
||||
"""mean recall of all classes.
|
||||
Return:
|
||||
float.
|
||||
"""
|
||||
return self.r.mean() if len(self.r) else 0.0
|
||||
|
||||
@property
|
||||
def map50(self):
|
||||
"""Mean AP@0.5 of all classes.
|
||||
Return:
|
||||
float.
|
||||
"""
|
||||
return self.all_ap[:, 0].mean() if len(self.all_ap) else 0.0
|
||||
|
||||
@property
|
||||
def map(self):
|
||||
"""Mean AP@0.5:0.95 of all classes.
|
||||
Return:
|
||||
float.
|
||||
"""
|
||||
return self.all_ap.mean() if len(self.all_ap) else 0.0
|
||||
|
||||
def mean_results(self):
|
||||
"""Mean of results, return mp, mr, map50, map"""
|
||||
return (self.mp, self.mr, self.map50, self.map)
|
||||
|
||||
def class_result(self, i):
|
||||
"""class-aware result, return p[i], r[i], ap50[i], ap[i]"""
|
||||
return (self.p[i], self.r[i], self.ap50[i], self.ap[i])
|
||||
|
||||
def get_maps(self, nc):
|
||||
maps = np.zeros(nc) + self.map
|
||||
for i, c in enumerate(self.ap_class_index):
|
||||
maps[c] = self.ap[i]
|
||||
return maps
|
||||
|
||||
def update(self, results):
|
||||
"""
|
||||
Args:
|
||||
results: tuple(p, r, ap, f1, ap_class)
|
||||
"""
|
||||
p, r, all_ap, f1, ap_class_index = results
|
||||
self.p = p
|
||||
self.r = r
|
||||
self.all_ap = all_ap
|
||||
self.f1 = f1
|
||||
self.ap_class_index = ap_class_index
|
||||
|
||||
|
||||
class Metrics:
|
||||
"""Metric for boxes and masks."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.metric_box = Metric()
|
||||
self.metric_mask = Metric()
|
||||
|
||||
def update(self, results):
|
||||
"""
|
||||
Args:
|
||||
results: Dict{'boxes': Dict{}, 'masks': Dict{}}
|
||||
"""
|
||||
self.metric_box.update(list(results["boxes"].values()))
|
||||
self.metric_mask.update(list(results["masks"].values()))
|
||||
|
||||
def mean_results(self):
|
||||
return self.metric_box.mean_results() + self.metric_mask.mean_results()
|
||||
|
||||
def class_result(self, i):
|
||||
return self.metric_box.class_result(i) + self.metric_mask.class_result(i)
|
||||
|
||||
def get_maps(self, nc):
|
||||
return self.metric_box.get_maps(nc) + self.metric_mask.get_maps(nc)
|
||||
|
||||
@property
|
||||
def ap_class_index(self):
|
||||
# boxes and masks have the same ap_class_index
|
||||
return self.metric_box.ap_class_index
|
||||
|
||||
@ -5,6 +5,7 @@ import time
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torchvision
|
||||
|
||||
from ultralytics.yolo.utils import LOGGER
|
||||
@ -32,14 +33,23 @@ class Profile(contextlib.ContextDecorator):
|
||||
return time.time()
|
||||
|
||||
|
||||
def coco80_to_coco91_class(): # converts 80-index (val2014) to 91-index (paper)
|
||||
# https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/
|
||||
# a = np.loadtxt('data/coco.names', dtype='str', delimiter='\n')
|
||||
# b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n')
|
||||
# x1 = [list(a[i] == b).index(True) + 1 for i in range(80)] # darknet to coco
|
||||
# x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)] # coco to darknet
|
||||
return [
|
||||
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34,
|
||||
35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
|
||||
64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]
|
||||
|
||||
|
||||
def segment2box(segment, width=640, height=640):
|
||||
# Convert 1 segment label to 1 box label, applying inside-image constraint, i.e. (xy1, xy2, ...) to (xyxy)
|
||||
x, y = segment.T # segment xy
|
||||
inside = (x >= 0) & (y >= 0) & (x <= width) & (y <= height)
|
||||
x, y, = (
|
||||
x[inside],
|
||||
y[inside],
|
||||
)
|
||||
x, y, = x[inside], y[inside]
|
||||
return np.array([x.min(), y.min(), x.max(), y.max()]) if any(x) else np.zeros(4) # xyxy
|
||||
|
||||
|
||||
@ -304,3 +314,63 @@ def resample_segments(segments, n=1000):
|
||||
xp = np.arange(len(s))
|
||||
segments[i] = np.concatenate([np.interp(x, xp, s[:, i]) for i in range(2)]).reshape(2, -1).T # segment xy
|
||||
return segments
|
||||
|
||||
|
||||
def crop_mask(masks, boxes):
|
||||
"""
|
||||
"Crop" predicted masks by zeroing out everything not in the predicted bbox.
|
||||
Vectorized by Chong (thanks Chong).
|
||||
Args:
|
||||
- masks should be a size [h, w, n] tensor of masks
|
||||
- boxes should be a size [n, 4] tensor of bbox coords in relative point form
|
||||
"""
|
||||
|
||||
n, h, w = masks.shape
|
||||
x1, y1, x2, y2 = torch.chunk(boxes[:, :, None], 4, 1) # x1 shape(1,1,n)
|
||||
r = torch.arange(w, device=masks.device, dtype=x1.dtype)[None, None, :] # rows shape(1,w,1)
|
||||
c = torch.arange(h, device=masks.device, dtype=x1.dtype)[None, :, None] # cols shape(h,1,1)
|
||||
|
||||
return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))
|
||||
|
||||
|
||||
def process_mask_upsample(protos, masks_in, bboxes, shape):
|
||||
"""
|
||||
Crop after upsample.
|
||||
proto_out: [mask_dim, mask_h, mask_w]
|
||||
out_masks: [n, mask_dim], n is number of masks after nms
|
||||
bboxes: [n, 4], n is number of masks after nms
|
||||
shape:input_image_size, (h, w)
|
||||
return: h, w, n
|
||||
"""
|
||||
|
||||
c, mh, mw = protos.shape # CHW
|
||||
masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)
|
||||
masks = F.interpolate(masks[None], shape, mode='bilinear', align_corners=False)[0] # CHW
|
||||
masks = crop_mask(masks, bboxes) # CHW
|
||||
return masks.gt_(0.5)
|
||||
|
||||
|
||||
def process_mask(protos, masks_in, bboxes, shape, upsample=False):
|
||||
"""
|
||||
Crop before upsample.
|
||||
proto_out: [mask_dim, mask_h, mask_w]
|
||||
out_masks: [n, mask_dim], n is number of masks after nms
|
||||
bboxes: [n, 4], n is number of masks after nms
|
||||
shape:input_image_size, (h, w)
|
||||
return: h, w, n
|
||||
"""
|
||||
|
||||
c, mh, mw = protos.shape # CHW
|
||||
ih, iw = shape
|
||||
masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw) # CHW
|
||||
|
||||
downsampled_bboxes = bboxes.clone()
|
||||
downsampled_bboxes[:, 0] *= mw / iw
|
||||
downsampled_bboxes[:, 2] *= mw / iw
|
||||
downsampled_bboxes[:, 3] *= mh / ih
|
||||
downsampled_bboxes[:, 1] *= mh / ih
|
||||
|
||||
masks = crop_mask(masks, downsampled_bboxes) # CHW
|
||||
if upsample:
|
||||
masks = F.interpolate(masks[None], shape, mode='bilinear', align_corners=False)[0] # CHW
|
||||
return masks.gt_(0.5)
|
||||
|
||||
@ -179,3 +179,13 @@ def smart_inference_mode(torch_1_9=check_version(torch.__version__, '1.9.0')):
|
||||
def intersect_state_dicts(da, db, exclude=()):
|
||||
# Dictionary intersection of matching keys and shapes, omitting 'exclude' keys, using da values
|
||||
return {k: v for k, v in da.items() if k in db and all(x not in k for x in exclude) and v.shape == db[k].shape}
|
||||
|
||||
|
||||
def is_parallel(model):
|
||||
# Returns True if model is of type DP or DDP
|
||||
return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)
|
||||
|
||||
|
||||
def de_parallel(model):
|
||||
# De-parallelize a model: returns single-GPU model if model is of type DP or DDP
|
||||
return model.module if is_parallel(model) else model
|
||||
|
||||
Reference in New Issue
Block a user