paulo061/codeformer

1import math2import time3 4import numpy as np5import torch6import torchvision7 8 9def check_img_size(img_size, s=32):10    # Verify img_size is a multiple of stride s11    new_size = make_divisible(img_size, int(s))  # ceil gs-multiple12    # if new_size != img_size:13    #     print(f"WARNING: --img-size {img_size:g} must be multiple of max stride {s:g}, updating to {new_size:g}")14    return new_size15 16 17def make_divisible(x, divisor):18    # Returns x evenly divisible by divisor19    return math.ceil(x / divisor) * divisor20 21 22def xyxy2xywh(x):23    # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right24    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)25    y[:, 0] = (x[:, 0] + x[:, 2]) / 2  # x center26    y[:, 1] = (x[:, 1] + x[:, 3]) / 2  # y center27    y[:, 2] = x[:, 2] - x[:, 0]  # width28    y[:, 3] = x[:, 3] - x[:, 1]  # height29    return y30 31 32def xywh2xyxy(x):33    # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right34    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)35    y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x36    y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y37    y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x38    y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y39    return y40 41 42def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):43    # Rescale coords (xyxy) from img1_shape to img0_shape44    if ratio_pad is None:  # calculate from img0_shape45        gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new46        pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding47    else:48        gain = ratio_pad[0][0]49        pad = ratio_pad[1]50 51    coords[:, [0, 2]] -= pad[0]  # x padding52    coords[:, [1, 3]] -= pad[1]  # y padding53    coords[:, :4] /= gain54    clip_coords(coords, img0_shape)55    return coords56 57 58def clip_coords(boxes, img_shape):59    # Clip bounding xyxy bounding boxes to image shape (height, width)60    boxes[:, 0].clamp_(0, img_shape[1])  # x161    boxes[:, 1].clamp_(0, img_shape[0])  # y162    boxes[:, 2].clamp_(0, img_shape[1])  # x263    boxes[:, 3].clamp_(0, img_shape[0])  # y264 65 66def box_iou(box1, box2):67    # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py68    """69    Return intersection-over-union (Jaccard index) of boxes.70    Both sets of boxes are expected to be in (x1, y1, x2, y2) format.71    Arguments:72        box1 (Tensor[N, 4])73        box2 (Tensor[M, 4])74    Returns:75        iou (Tensor[N, M]): the NxM matrix containing the pairwise76            IoU values for every element in boxes1 and boxes277    """78 79    def box_area(box):80        return (box[2] - box[0]) * (box[3] - box[1])81 82    area1 = box_area(box1.T)83    area2 = box_area(box2.T)84 85    inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)86    return inter / (area1[:, None] + area2 - inter)87 88 89def non_max_suppression_face(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, labels=()):90    """Performs Non-Maximum Suppression (NMS) on inference results91    Returns:92         detections with shape: nx6 (x1, y1, x2, y2, conf, cls)93    """94 95    nc = prediction.shape[2] - 15  # number of classes96    xc = prediction[..., 4] > conf_thres  # candidates97 98    # Settings99    # (pixels) maximum box width and height100    max_wh = 4096101    time_limit = 10.0  # seconds to quit after102    redundant = True  # require redundant detections103    multi_label = nc > 1  # multiple labels per box (adds 0.5ms/img)104    merge = False  # use merge-NMS105 106    t = time.time()107    output = [torch.zeros((0, 16), device=prediction.device)] * prediction.shape[0]108    for xi, x in enumerate(prediction):  # image index, image inference109        # Apply constraints110        x = x[xc[xi]]  # confidence111 112        # Cat apriori labels if autolabelling113        if labels and len(labels[xi]):114            label = labels[xi]115            v = torch.zeros((len(label), nc + 15), device=x.device)116            v[:, :4] = label[:, 1:5]  # box117            v[:, 4] = 1.0  # conf118            v[range(len(label)), label[:, 0].long() + 15] = 1.0  # cls119            x = torch.cat((x, v), 0)120 121        # If none remain process next image122        if not x.shape[0]:123            continue124 125        # Compute conf126        x[:, 15:] *= x[:, 4:5]  # conf = obj_conf * cls_conf127 128        # Box (center x, center y, width, height) to (x1, y1, x2, y2)129        box = xywh2xyxy(x[:, :4])130 131        # Detections matrix nx6 (xyxy, conf, landmarks, cls)132        if multi_label:133            i, j = (x[:, 15:] > conf_thres).nonzero(as_tuple=False).T134            x = torch.cat((box[i], x[i, j + 15, None], x[:, 5:15], j[:, None].float()), 1)135        else:  # best class only136            conf, j = x[:, 15:].max(1, keepdim=True)137            x = torch.cat((box, conf, x[:, 5:15], j.float()), 1)[conf.view(-1) > conf_thres]138 139        # Filter by class140        if classes is not None:141            x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]142 143        # If none remain process next image144        n = x.shape[0]  # number of boxes145        if not n:146            continue147 148        # Batched NMS149        c = x[:, 15:16] * (0 if agnostic else max_wh)  # classes150        boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores151        i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS152 153        if merge and (1 < n < 3e3):  # Merge NMS (boxes merged using weighted mean)154            # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)155            iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix156            weights = iou * scores[None]  # box weights157            x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes158            if redundant:159                i = i[iou.sum(1) > 1]  # require redundancy160 161        output[xi] = x[i]162        if (time.time() - t) > time_limit:163            break  # time limit exceeded164 165    return output166 167 168def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, labels=()):169    """Performs Non-Maximum Suppression (NMS) on inference results170 171    Returns:172         detections with shape: nx6 (x1, y1, x2, y2, conf, cls)173    """174 175    nc = prediction.shape[2] - 5  # number of classes176    xc = prediction[..., 4] > conf_thres  # candidates177 178    # Settings179    # (pixels) maximum box width and height180    max_wh = 4096181    time_limit = 10.0  # seconds to quit after182    redundant = True  # require redundant detections183    multi_label = nc > 1  # multiple labels per box (adds 0.5ms/img)184    merge = False  # use merge-NMS185 186    t = time.time()187    output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]188    for xi, x in enumerate(prediction):  # image index, image inference189        x = x[xc[xi]]  # confidence190 191        # Cat apriori labels if autolabelling192        if labels and len(labels[xi]):193            label_id = labels[xi]194            v = torch.zeros((len(label_id), nc + 5), device=x.device)195            v[:, :4] = label_id[:, 1:5]  # box196            v[:, 4] = 1.0  # conf197            v[range(len(label_id)), label_id[:, 0].long() + 5] = 1.0  # cls198            x = torch.cat((x, v), 0)199 200        # If none remain process next image201        if not x.shape[0]:202            continue203 204        # Compute conf205        x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf206 207        # Box (center x, center y, width, height) to (x1, y1, x2, y2)208        box = xywh2xyxy(x[:, :4])209 210        # Detections matrix nx6 (xyxy, conf, cls)211        if multi_label:212            i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T213            x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)214        else:  # best class only215            conf, j = x[:, 5:].max(1, keepdim=True)216            x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]217 218        # Filter by class219        if classes is not None:220            x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]221 222        # Check shape223        n = x.shape[0]  # number of boxes224        if not n:  # no boxes225            continue226 227        x = x[x[:, 4].argsort(descending=True)]  # sort by confidence228 229        # Batched NMS230        c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes231        boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores232        i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS233        if merge and (1 < n < 3e3):  # Merge NMS (boxes merged using weighted mean)234            # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)235            iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix236            weights = iou * scores[None]  # box weights237            x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes238            if redundant:239                i = i[iou.sum(1) > 1]  # require redundancy240 241        output[xi] = x[i]242        if (time.time() - t) > time_limit:243            print(f"WARNING: NMS time limit {time_limit}s exceeded")244            break  # time limit exceeded245 246    return output247 248 249def scale_coords_landmarks(img1_shape, coords, img0_shape, ratio_pad=None):250    # Rescale coords (xyxy) from img1_shape to img0_shape251    if ratio_pad is None:  # calculate from img0_shape252        gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new253        pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding254    else:255        gain = ratio_pad[0][0]256        pad = ratio_pad[1]257 258    coords[:, [0, 2, 4, 6, 8]] -= pad[0]  # x padding259    coords[:, [1, 3, 5, 7, 9]] -= pad[1]  # y padding260    coords[:, :10] /= gain261    coords[:, 0].clamp_(0, img0_shape[1])  # x1262    coords[:, 1].clamp_(0, img0_shape[0])  # y1263    coords[:, 2].clamp_(0, img0_shape[1])  # x2264    coords[:, 3].clamp_(0, img0_shape[0])  # y2265    coords[:, 4].clamp_(0, img0_shape[1])  # x3266    coords[:, 5].clamp_(0, img0_shape[0])  # y3267    coords[:, 6].clamp_(0, img0_shape[1])  # x4268    coords[:, 7].clamp_(0, img0_shape[0])  # y4269    coords[:, 8].clamp_(0, img0_shape[1])  # x5270    coords[:, 9].clamp_(0, img0_shape[0])  # y5271    return coords272