Add files via uploadHEADmain

3 files changed, 1006 insertions, 0 deletions

diff --git a/r_facelib/detection/retinaface/retinaface.py b/r_facelib/detection/retinaface/retinaface.py
new file mode 100644
index 0000000..5d9770a
--- /dev/null
+++ b/r_facelib/detection/retinaface/retinaface.py
@@ -0,0 +1,389 @@
+import cv2

+import numpy as np

+import torch

+import torch.nn as nn

+import torch.nn.functional as F

+from PIL import Image

+from torchvision.models._utils import IntermediateLayerGetter as IntermediateLayerGetter

+

+from modules import shared

+

+from r_facelib.detection.align_trans import get_reference_facial_points, warp_and_crop_face

+from r_facelib.detection.retinaface.retinaface_net import FPN, SSH, MobileNetV1, make_bbox_head, make_class_head, make_landmark_head

+from r_facelib.detection.retinaface.retinaface_utils import (PriorBox, batched_decode, batched_decode_landm, decode, decode_landm,

+                                                 py_cpu_nms)

+

+#device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

+if torch.cuda.is_available():

+    device = torch.device('cuda')

+elif torch.backends.mps.is_available():

+    device = torch.device('mps')

+# elif hasattr(torch,'dml'):

+#     device = torch.device('dml')

+elif hasattr(torch,'dml') or hasattr(torch,'privateuseone'): # AMD

+    if shared.cmd_opts is not None: # A1111

+        if shared.cmd_opts.device_id is not None:

+            device = torch.device(f'privateuseone:{shared.cmd_opts.device_id}')

+        else:

+            device = torch.device('privateuseone:0')

+    else:

+        device = torch.device('privateuseone:0')

+else:

+    device = torch.device('cpu')

+

+

+def generate_config(network_name):

+

+    cfg_mnet = {

+        'name': 'mobilenet0.25',

+        'min_sizes': [[16, 32], [64, 128], [256, 512]],

+        'steps': [8, 16, 32],

+        'variance': [0.1, 0.2],

+        'clip': False,

+        'loc_weight': 2.0,

+        'gpu_train': True,

+        'batch_size': 32,

+        'ngpu': 1,

+        'epoch': 250,

+        'decay1': 190,

+        'decay2': 220,

+        'image_size': 640,

+        'return_layers': {

+            'stage1': 1,

+            'stage2': 2,

+            'stage3': 3

+        },

+        'in_channel': 32,

+        'out_channel': 64

+    }

+

+    cfg_re50 = {

+        'name': 'Resnet50',

+        'min_sizes': [[16, 32], [64, 128], [256, 512]],

+        'steps': [8, 16, 32],

+        'variance': [0.1, 0.2],

+        'clip': False,

+        'loc_weight': 2.0,

+        'gpu_train': True,

+        'batch_size': 24,

+        'ngpu': 4,

+        'epoch': 100,

+        'decay1': 70,

+        'decay2': 90,

+        'image_size': 840,

+        'return_layers': {

+            'layer2': 1,

+            'layer3': 2,

+            'layer4': 3

+        },

+        'in_channel': 256,

+        'out_channel': 256

+    }

+

+    if network_name == 'mobile0.25':

+        return cfg_mnet

+    elif network_name == 'resnet50':

+        return cfg_re50

+    else:

+        raise NotImplementedError(f'network_name={network_name}')

+

+

+class RetinaFace(nn.Module):

+

+    def __init__(self, network_name='resnet50', half=False, phase='test'):

+        super(RetinaFace, self).__init__()

+        self.half_inference = half

+        cfg = generate_config(network_name)

+        self.backbone = cfg['name']

+

+        self.model_name = f'retinaface_{network_name}'

+        self.cfg = cfg

+        self.phase = phase

+        self.target_size, self.max_size = 1600, 2150

+        self.resize, self.scale, self.scale1 = 1., None, None

+        self.mean_tensor = torch.tensor([[[[104.]], [[117.]], [[123.]]]]).to(device)

+        self.reference = get_reference_facial_points(default_square=True)

+        # Build network.

+        backbone = None

+        if cfg['name'] == 'mobilenet0.25':

+            backbone = MobileNetV1()

+            self.body = IntermediateLayerGetter(backbone, cfg['return_layers'])

+        elif cfg['name'] == 'Resnet50':

+            import torchvision.models as models

+            backbone = models.resnet50(pretrained=False)

+            self.body = IntermediateLayerGetter(backbone, cfg['return_layers'])

+

+        in_channels_stage2 = cfg['in_channel']

+        in_channels_list = [

+            in_channels_stage2 * 2,

+            in_channels_stage2 * 4,

+            in_channels_stage2 * 8,

+        ]

+

+        out_channels = cfg['out_channel']

+        self.fpn = FPN(in_channels_list, out_channels)

+        self.ssh1 = SSH(out_channels, out_channels)

+        self.ssh2 = SSH(out_channels, out_channels)

+        self.ssh3 = SSH(out_channels, out_channels)

+

+        self.ClassHead = make_class_head(fpn_num=3, inchannels=cfg['out_channel'])

+        self.BboxHead = make_bbox_head(fpn_num=3, inchannels=cfg['out_channel'])

+        self.LandmarkHead = make_landmark_head(fpn_num=3, inchannels=cfg['out_channel'])

+

+        self.to(device)

+        self.eval()

+        if self.half_inference:

+            self.half()

+

+    def forward(self, inputs):

+        self.to(device)

+        out = self.body(inputs)

+

+        if self.backbone == 'mobilenet0.25' or self.backbone == 'Resnet50':

+            out = list(out.values())

+        # FPN

+        fpn = self.fpn(out)

+

+        # SSH

+        feature1 = self.ssh1(fpn[0])

+        feature2 = self.ssh2(fpn[1])

+        feature3 = self.ssh3(fpn[2])

+        features = [feature1, feature2, feature3]

+

+        bbox_regressions = torch.cat([self.BboxHead[i](feature) for i, feature in enumerate(features)], dim=1)

+        classifications = torch.cat([self.ClassHead[i](feature) for i, feature in enumerate(features)], dim=1)

+        tmp = [self.LandmarkHead[i](feature) for i, feature in enumerate(features)]

+        ldm_regressions = (torch.cat(tmp, dim=1))

+

+        if self.phase == 'train':

+            output = (bbox_regressions, classifications, ldm_regressions)

+        else:

+            output = (bbox_regressions, F.softmax(classifications, dim=-1), ldm_regressions)

+        return output

+

+    def __detect_faces(self, inputs):

+        # get scale

+        height, width = inputs.shape[2:]

+        self.scale = torch.tensor([width, height, width, height], dtype=torch.float32).to(device)

+        tmp = [width, height, width, height, width, height, width, height, width, height]

+        self.scale1 = torch.tensor(tmp, dtype=torch.float32).to(device)

+

+        # forawrd

+        inputs = inputs.to(device)

+        if self.half_inference:

+            inputs = inputs.half()

+        loc, conf, landmarks = self(inputs)

+

+        # get priorbox

+        priorbox = PriorBox(self.cfg, image_size=inputs.shape[2:])

+        priors = priorbox.forward().to(device)

+

+        return loc, conf, landmarks, priors

+

+    # single image detection

+    def transform(self, image, use_origin_size):

+        # convert to opencv format

+        if isinstance(image, Image.Image):

+            image = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)

+        image = image.astype(np.float32)

+

+        # testing scale

+        im_size_min = np.min(image.shape[0:2])

+        im_size_max = np.max(image.shape[0:2])

+        resize = float(self.target_size) / float(im_size_min)

+

+        # prevent bigger axis from being more than max_size

+        if np.round(resize * im_size_max) > self.max_size:

+            resize = float(self.max_size) / float(im_size_max)

+        resize = 1 if use_origin_size else resize

+

+        # resize

+        if resize != 1:

+            image = cv2.resize(image, None, None, fx=resize, fy=resize, interpolation=cv2.INTER_LINEAR)

+

+        # convert to torch.tensor format

+        # image -= (104, 117, 123)

+        image = image.transpose(2, 0, 1)

+        image = torch.from_numpy(image).unsqueeze(0)

+

+        return image, resize

+

+    def detect_faces(

+        self,

+        image,

+        conf_threshold=0.8,

+        nms_threshold=0.4,

+        use_origin_size=True,

+    ):

+        """

+        Params:

+            imgs: BGR image

+        """

+        image, self.resize = self.transform(image, use_origin_size)

+        image = image.to(device)

+        if self.half_inference:

+            image = image.half()

+        image = image - self.mean_tensor

+

+        loc, conf, landmarks, priors = self.__detect_faces(image)

+

+        boxes = decode(loc.data.squeeze(0), priors.data, self.cfg['variance'])

+        boxes = boxes * self.scale / self.resize

+        boxes = boxes.cpu().numpy()

+

+        scores = conf.squeeze(0).data.cpu().numpy()[:, 1]

+

+        landmarks = decode_landm(landmarks.squeeze(0), priors, self.cfg['variance'])

+        landmarks = landmarks * self.scale1 / self.resize

+        landmarks = landmarks.cpu().numpy()

+

+        # ignore low scores

+        inds = np.where(scores > conf_threshold)[0]

+        boxes, landmarks, scores = boxes[inds], landmarks[inds], scores[inds]

+

+        # sort

+        order = scores.argsort()[::-1]

+        boxes, landmarks, scores = boxes[order], landmarks[order], scores[order]

+

+        # do NMS

+        bounding_boxes = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)

+        keep = py_cpu_nms(bounding_boxes, nms_threshold)

+        bounding_boxes, landmarks = bounding_boxes[keep, :], landmarks[keep]

+        # self.t['forward_pass'].toc()

+        # print(self.t['forward_pass'].average_time)

+        # import sys

+        # sys.stdout.flush()

+        return np.concatenate((bounding_boxes, landmarks), axis=1)

+

+    def __align_multi(self, image, boxes, landmarks, limit=None):

+

+        if len(boxes) < 1:

+            return [], []

+

+        if limit:

+            boxes = boxes[:limit]

+            landmarks = landmarks[:limit]

+

+        faces = []

+        for landmark in landmarks:

+            facial5points = [[landmark[2 * j], landmark[2 * j + 1]] for j in range(5)]

+

+            warped_face = warp_and_crop_face(np.array(image), facial5points, self.reference, crop_size=(112, 112))

+            faces.append(warped_face)

+

+        return np.concatenate((boxes, landmarks), axis=1), faces

+

+    def align_multi(self, img, conf_threshold=0.8, limit=None):

+

+        rlt = self.detect_faces(img, conf_threshold=conf_threshold)

+        boxes, landmarks = rlt[:, 0:5], rlt[:, 5:]

+

+        return self.__align_multi(img, boxes, landmarks, limit)

+

+    # batched detection

+    def batched_transform(self, frames, use_origin_size):

+        """

+        Arguments:

+            frames: a list of PIL.Image, or torch.Tensor(shape=[n, h, w, c],

+                type=np.float32, BGR format).

+            use_origin_size: whether to use origin size.

+        """

+        from_PIL = True if isinstance(frames[0], Image.Image) else False

+

+        # convert to opencv format

+        if from_PIL:

+            frames = [cv2.cvtColor(np.asarray(frame), cv2.COLOR_RGB2BGR) for frame in frames]

+            frames = np.asarray(frames, dtype=np.float32)

+

+        # testing scale

+        im_size_min = np.min(frames[0].shape[0:2])

+        im_size_max = np.max(frames[0].shape[0:2])

+        resize = float(self.target_size) / float(im_size_min)

+

+        # prevent bigger axis from being more than max_size

+        if np.round(resize * im_size_max) > self.max_size:

+            resize = float(self.max_size) / float(im_size_max)

+        resize = 1 if use_origin_size else resize

+

+        # resize

+        if resize != 1:

+            if not from_PIL:

+                frames = F.interpolate(frames, scale_factor=resize)

+            else:

+                frames = [

+                    cv2.resize(frame, None, None, fx=resize, fy=resize, interpolation=cv2.INTER_LINEAR)

+                    for frame in frames

+                ]

+

+        # convert to torch.tensor format

+        if not from_PIL:

+            frames = frames.transpose(1, 2).transpose(1, 3).contiguous()

+        else:

+            frames = frames.transpose((0, 3, 1, 2))

+            frames = torch.from_numpy(frames)

+

+        return frames, resize

+

+    def batched_detect_faces(self, frames, conf_threshold=0.8, nms_threshold=0.4, use_origin_size=True):

+        """

+        Arguments:

+            frames: a list of PIL.Image, or np.array(shape=[n, h, w, c],

+                type=np.uint8, BGR format).

+            conf_threshold: confidence threshold.

+            nms_threshold: nms threshold.

+            use_origin_size: whether to use origin size.

+        Returns:

+            final_bounding_boxes: list of np.array ([n_boxes, 5],

+                type=np.float32).

+            final_landmarks: list of np.array ([n_boxes, 10], type=np.float32).

+        """

+        # self.t['forward_pass'].tic()

+        frames, self.resize = self.batched_transform(frames, use_origin_size)

+        frames = frames.to(device)

+        frames = frames - self.mean_tensor

+

+        b_loc, b_conf, b_landmarks, priors = self.__detect_faces(frames)

+

+        final_bounding_boxes, final_landmarks = [], []

+

+        # decode

+        priors = priors.unsqueeze(0)

+        b_loc = batched_decode(b_loc, priors, self.cfg['variance']) * self.scale / self.resize

+        b_landmarks = batched_decode_landm(b_landmarks, priors, self.cfg['variance']) * self.scale1 / self.resize

+        b_conf = b_conf[:, :, 1]

+

+        # index for selection

+        b_indice = b_conf > conf_threshold

+

+        # concat

+        b_loc_and_conf = torch.cat((b_loc, b_conf.unsqueeze(-1)), dim=2).float()

+

+        for pred, landm, inds in zip(b_loc_and_conf, b_landmarks, b_indice):

+

+            # ignore low scores

+            pred, landm = pred[inds, :], landm[inds, :]

+            if pred.shape[0] == 0:

+                final_bounding_boxes.append(np.array([], dtype=np.float32))

+                final_landmarks.append(np.array([], dtype=np.float32))

+                continue

+

+            # sort

+            # order = score.argsort(descending=True)

+            # box, landm, score = box[order], landm[order], score[order]

+

+            # to CPU

+            bounding_boxes, landm = pred.cpu().numpy(), landm.cpu().numpy()

+

+            # NMS

+            keep = py_cpu_nms(bounding_boxes, nms_threshold)

+            bounding_boxes, landmarks = bounding_boxes[keep, :], landm[keep]

+

+            # append

+            final_bounding_boxes.append(bounding_boxes)

+            final_landmarks.append(landmarks)

+        # self.t['forward_pass'].toc(average=True)

+        # self.batch_time += self.t['forward_pass'].diff

+        # self.total_frame += len(frames)

+        # print(self.batch_time / self.total_frame)

+

+        return final_bounding_boxes, final_landmarks

diff --git a/r_facelib/detection/retinaface/retinaface_net.py b/r_facelib/detection/retinaface/retinaface_net.py
new file mode 100644
index 0000000..c52535e
--- /dev/null
+++ b/r_facelib/detection/retinaface/retinaface_net.py
@@ -0,0 +1,196 @@
+import torch

+import torch.nn as nn

+import torch.nn.functional as F

+

+

+def conv_bn(inp, oup, stride=1, leaky=0):

+    return nn.Sequential(

+        nn.Conv2d(inp, oup, 3, stride, 1, bias=False), nn.BatchNorm2d(oup),

+        nn.LeakyReLU(negative_slope=leaky, inplace=True))

+

+

+def conv_bn_no_relu(inp, oup, stride):

+    return nn.Sequential(

+        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),

+        nn.BatchNorm2d(oup),

+    )

+

+

+def conv_bn1X1(inp, oup, stride, leaky=0):

+    return nn.Sequential(

+        nn.Conv2d(inp, oup, 1, stride, padding=0, bias=False), nn.BatchNorm2d(oup),

+        nn.LeakyReLU(negative_slope=leaky, inplace=True))

+

+

+def conv_dw(inp, oup, stride, leaky=0.1):

+    return nn.Sequential(

+        nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),

+        nn.BatchNorm2d(inp),

+        nn.LeakyReLU(negative_slope=leaky, inplace=True),

+        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),

+        nn.BatchNorm2d(oup),

+        nn.LeakyReLU(negative_slope=leaky, inplace=True),

+    )

+

+

+class SSH(nn.Module):

+

+    def __init__(self, in_channel, out_channel):

+        super(SSH, self).__init__()

+        assert out_channel % 4 == 0

+        leaky = 0

+        if (out_channel <= 64):

+            leaky = 0.1

+        self.conv3X3 = conv_bn_no_relu(in_channel, out_channel // 2, stride=1)

+

+        self.conv5X5_1 = conv_bn(in_channel, out_channel // 4, stride=1, leaky=leaky)

+        self.conv5X5_2 = conv_bn_no_relu(out_channel // 4, out_channel // 4, stride=1)

+

+        self.conv7X7_2 = conv_bn(out_channel // 4, out_channel // 4, stride=1, leaky=leaky)

+        self.conv7x7_3 = conv_bn_no_relu(out_channel // 4, out_channel // 4, stride=1)

+

+    def forward(self, input):

+        conv3X3 = self.conv3X3(input)

+

+        conv5X5_1 = self.conv5X5_1(input)

+        conv5X5 = self.conv5X5_2(conv5X5_1)

+

+        conv7X7_2 = self.conv7X7_2(conv5X5_1)

+        conv7X7 = self.conv7x7_3(conv7X7_2)

+

+        out = torch.cat([conv3X3, conv5X5, conv7X7], dim=1)

+        out = F.relu(out)

+        return out

+

+

+class FPN(nn.Module):

+

+    def __init__(self, in_channels_list, out_channels):

+        super(FPN, self).__init__()

+        leaky = 0

+        if (out_channels <= 64):

+            leaky = 0.1

+        self.output1 = conv_bn1X1(in_channels_list[0], out_channels, stride=1, leaky=leaky)

+        self.output2 = conv_bn1X1(in_channels_list[1], out_channels, stride=1, leaky=leaky)

+        self.output3 = conv_bn1X1(in_channels_list[2], out_channels, stride=1, leaky=leaky)

+

+        self.merge1 = conv_bn(out_channels, out_channels, leaky=leaky)

+        self.merge2 = conv_bn(out_channels, out_channels, leaky=leaky)

+

+    def forward(self, input):

+        # names = list(input.keys())

+        # input = list(input.values())

+

+        output1 = self.output1(input[0])

+        output2 = self.output2(input[1])

+        output3 = self.output3(input[2])

+

+        up3 = F.interpolate(output3, size=[output2.size(2), output2.size(3)], mode='nearest')

+        output2 = output2 + up3

+        output2 = self.merge2(output2)

+

+        up2 = F.interpolate(output2, size=[output1.size(2), output1.size(3)], mode='nearest')

+        output1 = output1 + up2

+        output1 = self.merge1(output1)

+

+        out = [output1, output2, output3]

+        return out

+

+

+class MobileNetV1(nn.Module):

+

+    def __init__(self):

+        super(MobileNetV1, self).__init__()

+        self.stage1 = nn.Sequential(

+            conv_bn(3, 8, 2, leaky=0.1),  # 3

+            conv_dw(8, 16, 1),  # 7

+            conv_dw(16, 32, 2),  # 11

+            conv_dw(32, 32, 1),  # 19

+            conv_dw(32, 64, 2),  # 27

+            conv_dw(64, 64, 1),  # 43

+        )

+        self.stage2 = nn.Sequential(

+            conv_dw(64, 128, 2),  # 43 + 16 = 59

+            conv_dw(128, 128, 1),  # 59 + 32 = 91

+            conv_dw(128, 128, 1),  # 91 + 32 = 123

+            conv_dw(128, 128, 1),  # 123 + 32 = 155

+            conv_dw(128, 128, 1),  # 155 + 32 = 187

+            conv_dw(128, 128, 1),  # 187 + 32 = 219

+        )

+        self.stage3 = nn.Sequential(

+            conv_dw(128, 256, 2),  # 219 +3 2 = 241

+            conv_dw(256, 256, 1),  # 241 + 64 = 301

+        )

+        self.avg = nn.AdaptiveAvgPool2d((1, 1))

+        self.fc = nn.Linear(256, 1000)

+

+    def forward(self, x):

+        x = self.stage1(x)

+        x = self.stage2(x)

+        x = self.stage3(x)

+        x = self.avg(x)

+        # x = self.model(x)

+        x = x.view(-1, 256)

+        x = self.fc(x)

+        return x

+

+

+class ClassHead(nn.Module):

+

+    def __init__(self, inchannels=512, num_anchors=3):

+        super(ClassHead, self).__init__()

+        self.num_anchors = num_anchors

+        self.conv1x1 = nn.Conv2d(inchannels, self.num_anchors * 2, kernel_size=(1, 1), stride=1, padding=0)

+

+    def forward(self, x):

+        out = self.conv1x1(x)

+        out = out.permute(0, 2, 3, 1).contiguous()

+

+        return out.view(out.shape[0], -1, 2)

+

+

+class BboxHead(nn.Module):

+

+    def __init__(self, inchannels=512, num_anchors=3):

+        super(BboxHead, self).__init__()

+        self.conv1x1 = nn.Conv2d(inchannels, num_anchors * 4, kernel_size=(1, 1), stride=1, padding=0)

+

+    def forward(self, x):

+        out = self.conv1x1(x)

+        out = out.permute(0, 2, 3, 1).contiguous()

+

+        return out.view(out.shape[0], -1, 4)

+

+

+class LandmarkHead(nn.Module):

+

+    def __init__(self, inchannels=512, num_anchors=3):

+        super(LandmarkHead, self).__init__()

+        self.conv1x1 = nn.Conv2d(inchannels, num_anchors * 10, kernel_size=(1, 1), stride=1, padding=0)

+

+    def forward(self, x):

+        out = self.conv1x1(x)

+        out = out.permute(0, 2, 3, 1).contiguous()

+

+        return out.view(out.shape[0], -1, 10)

+

+

+def make_class_head(fpn_num=3, inchannels=64, anchor_num=2):

+    classhead = nn.ModuleList()

+    for i in range(fpn_num):

+        classhead.append(ClassHead(inchannels, anchor_num))

+    return classhead

+

+

+def make_bbox_head(fpn_num=3, inchannels=64, anchor_num=2):

+    bboxhead = nn.ModuleList()

+    for i in range(fpn_num):

+        bboxhead.append(BboxHead(inchannels, anchor_num))

+    return bboxhead

+

+

+def make_landmark_head(fpn_num=3, inchannels=64, anchor_num=2):

+    landmarkhead = nn.ModuleList()

+    for i in range(fpn_num):

+        landmarkhead.append(LandmarkHead(inchannels, anchor_num))

+    return landmarkhead

diff --git a/r_facelib/detection/retinaface/retinaface_utils.py b/r_facelib/detection/retinaface/retinaface_utils.py
new file mode 100644
index 0000000..f19e320
--- /dev/null
+++ b/r_facelib/detection/retinaface/retinaface_utils.py
@@ -0,0 +1,421 @@
+import numpy as np

+import torch

+import torchvision

+from itertools import product as product

+from math import ceil

+

+

+class PriorBox(object):

+

+    def __init__(self, cfg, image_size=None, phase='train'):

+        super(PriorBox, self).__init__()

+        self.min_sizes = cfg['min_sizes']

+        self.steps = cfg['steps']

+        self.clip = cfg['clip']

+        self.image_size = image_size

+        self.feature_maps = [[ceil(self.image_size[0] / step), ceil(self.image_size[1] / step)] for step in self.steps]

+        self.name = 's'

+

+    def forward(self):

+        anchors = []

+        for k, f in enumerate(self.feature_maps):

+            min_sizes = self.min_sizes[k]

+            for i, j in product(range(f[0]), range(f[1])):

+                for min_size in min_sizes:

+                    s_kx = min_size / self.image_size[1]

+                    s_ky = min_size / self.image_size[0]

+                    dense_cx = [x * self.steps[k] / self.image_size[1] for x in [j + 0.5]]

+                    dense_cy = [y * self.steps[k] / self.image_size[0] for y in [i + 0.5]]

+                    for cy, cx in product(dense_cy, dense_cx):

+                        anchors += [cx, cy, s_kx, s_ky]

+

+        # back to torch land

+        output = torch.Tensor(anchors).view(-1, 4)

+        if self.clip:

+            output.clamp_(max=1, min=0)

+        return output

+

+

+def py_cpu_nms(dets, thresh):

+    """Pure Python NMS baseline."""

+    keep = torchvision.ops.nms(

+        boxes=torch.Tensor(dets[:, :4]),

+        scores=torch.Tensor(dets[:, 4]),

+        iou_threshold=thresh,

+    )

+

+    return list(keep)

+

+

+def point_form(boxes):

+    """ Convert prior_boxes to (xmin, ymin, xmax, ymax)

+    representation for comparison to point form ground truth data.

+    Args:

+        boxes: (tensor) center-size default boxes from priorbox layers.

+    Return:

+        boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.

+    """

+    return torch.cat(

+        (

+            boxes[:, :2] - boxes[:, 2:] / 2,  # xmin, ymin

+            boxes[:, :2] + boxes[:, 2:] / 2),

+        1)  # xmax, ymax

+

+

+def center_size(boxes):

+    """ Convert prior_boxes to (cx, cy, w, h)

+    representation for comparison to center-size form ground truth data.

+    Args:

+        boxes: (tensor) point_form boxes

+    Return:

+        boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.

+    """

+    return torch.cat(

+        (boxes[:, 2:] + boxes[:, :2]) / 2,  # cx, cy

+        boxes[:, 2:] - boxes[:, :2],

+        1)  # w, h

+

+

+def intersect(box_a, box_b):

+    """ We resize both tensors to [A,B,2] without new malloc:

+    [A,2] -> [A,1,2] -> [A,B,2]

+    [B,2] -> [1,B,2] -> [A,B,2]

+    Then we compute the area of intersect between box_a and box_b.

+    Args:

+      box_a: (tensor) bounding boxes, Shape: [A,4].

+      box_b: (tensor) bounding boxes, Shape: [B,4].

+    Return:

+      (tensor) intersection area, Shape: [A,B].

+    """

+    A = box_a.size(0)

+    B = box_b.size(0)

+    max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2), box_b[:, 2:].unsqueeze(0).expand(A, B, 2))

+    min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2), box_b[:, :2].unsqueeze(0).expand(A, B, 2))

+    inter = torch.clamp((max_xy - min_xy), min=0)

+    return inter[:, :, 0] * inter[:, :, 1]

+

+

+def jaccard(box_a, box_b):

+    """Compute the jaccard overlap of two sets of boxes.  The jaccard overlap

+    is simply the intersection over union of two boxes.  Here we operate on

+    ground truth boxes and default boxes.

+    E.g.:

+        A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)

+    Args:

+        box_a: (tensor) Ground truth bounding boxes, Shape: [num_objects,4]

+        box_b: (tensor) Prior boxes from priorbox layers, Shape: [num_priors,4]

+    Return:

+        jaccard overlap: (tensor) Shape: [box_a.size(0), box_b.size(0)]

+    """

+    inter = intersect(box_a, box_b)

+    area_a = ((box_a[:, 2] - box_a[:, 0]) * (box_a[:, 3] - box_a[:, 1])).unsqueeze(1).expand_as(inter)  # [A,B]

+    area_b = ((box_b[:, 2] - box_b[:, 0]) * (box_b[:, 3] - box_b[:, 1])).unsqueeze(0).expand_as(inter)  # [A,B]

+    union = area_a + area_b - inter

+    return inter / union  # [A,B]

+

+

+def matrix_iou(a, b):

+    """

+    return iou of a and b, numpy version for data augenmentation

+    """

+    lt = np.maximum(a[:, np.newaxis, :2], b[:, :2])

+    rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])

+

+    area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2)

+    area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)

+    area_b = np.prod(b[:, 2:] - b[:, :2], axis=1)

+    return area_i / (area_a[:, np.newaxis] + area_b - area_i)

+

+

+def matrix_iof(a, b):

+    """

+    return iof of a and b, numpy version for data augenmentation

+    """

+    lt = np.maximum(a[:, np.newaxis, :2], b[:, :2])

+    rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])

+

+    area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2)

+    area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)

+    return area_i / np.maximum(area_a[:, np.newaxis], 1)

+

+

+def match(threshold, truths, priors, variances, labels, landms, loc_t, conf_t, landm_t, idx):

+    """Match each prior box with the ground truth box of the highest jaccard

+    overlap, encode the bounding boxes, then return the matched indices

+    corresponding to both confidence and location preds.

+    Args:

+        threshold: (float) The overlap threshold used when matching boxes.

+        truths: (tensor) Ground truth boxes, Shape: [num_obj, 4].

+        priors: (tensor) Prior boxes from priorbox layers, Shape: [n_priors,4].

+        variances: (tensor) Variances corresponding to each prior coord,

+            Shape: [num_priors, 4].

+        labels: (tensor) All the class labels for the image, Shape: [num_obj].

+        landms: (tensor) Ground truth landms, Shape [num_obj, 10].

+        loc_t: (tensor) Tensor to be filled w/ encoded location targets.

+        conf_t: (tensor) Tensor to be filled w/ matched indices for conf preds.

+        landm_t: (tensor) Tensor to be filled w/ encoded landm targets.

+        idx: (int) current batch index

+    Return:

+        The matched indices corresponding to 1)location 2)confidence

+        3)landm preds.

+    """

+    # jaccard index

+    overlaps = jaccard(truths, point_form(priors))

+    # (Bipartite Matching)

+    # [1,num_objects] best prior for each ground truth

+    best_prior_overlap, best_prior_idx = overlaps.max(1, keepdim=True)

+

+    # ignore hard gt

+    valid_gt_idx = best_prior_overlap[:, 0] >= 0.2

+    best_prior_idx_filter = best_prior_idx[valid_gt_idx, :]

+    if best_prior_idx_filter.shape[0] <= 0:

+        loc_t[idx] = 0

+        conf_t[idx] = 0

+        return

+

+    # [1,num_priors] best ground truth for each prior

+    best_truth_overlap, best_truth_idx = overlaps.max(0, keepdim=True)

+    best_truth_idx.squeeze_(0)

+    best_truth_overlap.squeeze_(0)

+    best_prior_idx.squeeze_(1)

+    best_prior_idx_filter.squeeze_(1)

+    best_prior_overlap.squeeze_(1)

+    best_truth_overlap.index_fill_(0, best_prior_idx_filter, 2)  # ensure best prior

+    # TODO refactor: index  best_prior_idx with long tensor

+    # ensure every gt matches with its prior of max overlap

+    for j in range(best_prior_idx.size(0)):  # 判别此anchor是预测哪一个boxes

+        best_truth_idx[best_prior_idx[j]] = j

+    matches = truths[best_truth_idx]  # Shape: [num_priors,4] 此处为每一个anchor对应的bbox取出来

+    conf = labels[best_truth_idx]  # Shape: [num_priors]      此处为每一个anchor对应的label取出来

+    conf[best_truth_overlap < threshold] = 0  # label as background   overlap<0.35的全部作为负样本

+    loc = encode(matches, priors, variances)

+

+    matches_landm = landms[best_truth_idx]

+    landm = encode_landm(matches_landm, priors, variances)

+    loc_t[idx] = loc  # [num_priors,4] encoded offsets to learn

+    conf_t[idx] = conf  # [num_priors] top class label for each prior

+    landm_t[idx] = landm

+

+

+def encode(matched, priors, variances):

+    """Encode the variances from the priorbox layers into the ground truth boxes

+    we have matched (based on jaccard overlap) with the prior boxes.

+    Args:

+        matched: (tensor) Coords of ground truth for each prior in point-form

+            Shape: [num_priors, 4].

+        priors: (tensor) Prior boxes in center-offset form

+            Shape: [num_priors,4].

+        variances: (list[float]) Variances of priorboxes

+    Return:

+        encoded boxes (tensor), Shape: [num_priors, 4]

+    """

+

+    # dist b/t match center and prior's center

+    g_cxcy = (matched[:, :2] + matched[:, 2:]) / 2 - priors[:, :2]

+    # encode variance

+    g_cxcy /= (variances[0] * priors[:, 2:])

+    # match wh / prior wh

+    g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:]

+    g_wh = torch.log(g_wh) / variances[1]

+    # return target for smooth_l1_loss

+    return torch.cat([g_cxcy, g_wh], 1)  # [num_priors,4]

+

+

+def encode_landm(matched, priors, variances):

+    """Encode the variances from the priorbox layers into the ground truth boxes

+    we have matched (based on jaccard overlap) with the prior boxes.

+    Args:

+        matched: (tensor) Coords of ground truth for each prior in point-form

+            Shape: [num_priors, 10].

+        priors: (tensor) Prior boxes in center-offset form

+            Shape: [num_priors,4].

+        variances: (list[float]) Variances of priorboxes

+    Return:

+        encoded landm (tensor), Shape: [num_priors, 10]

+    """

+

+    # dist b/t match center and prior's center

+    matched = torch.reshape(matched, (matched.size(0), 5, 2))

+    priors_cx = priors[:, 0].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)

+    priors_cy = priors[:, 1].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)

+    priors_w = priors[:, 2].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)

+    priors_h = priors[:, 3].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)

+    priors = torch.cat([priors_cx, priors_cy, priors_w, priors_h], dim=2)

+    g_cxcy = matched[:, :, :2] - priors[:, :, :2]

+    # encode variance

+    g_cxcy /= (variances[0] * priors[:, :, 2:])

+    # g_cxcy /= priors[:, :, 2:]

+    g_cxcy = g_cxcy.reshape(g_cxcy.size(0), -1)

+    # return target for smooth_l1_loss

+    return g_cxcy

+

+

+# Adapted from https://github.com/Hakuyume/chainer-ssd

+def decode(loc, priors, variances):

+    """Decode locations from predictions using priors to undo

+    the encoding we did for offset regression at train time.

+    Args:

+        loc (tensor): location predictions for loc layers,

+            Shape: [num_priors,4]

+        priors (tensor): Prior boxes in center-offset form.

+            Shape: [num_priors,4].

+        variances: (list[float]) Variances of priorboxes

+    Return:

+        decoded bounding box predictions

+    """

+

+    boxes = torch.cat((priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],

+                       priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)

+    boxes[:, :2] -= boxes[:, 2:] / 2

+    boxes[:, 2:] += boxes[:, :2]

+    return boxes

+

+

+def decode_landm(pre, priors, variances):

+    """Decode landm from predictions using priors to undo

+    the encoding we did for offset regression at train time.

+    Args:

+        pre (tensor): landm predictions for loc layers,

+            Shape: [num_priors,10]

+        priors (tensor): Prior boxes in center-offset form.

+            Shape: [num_priors,4].

+        variances: (list[float]) Variances of priorboxes

+    Return:

+        decoded landm predictions

+    """

+    tmp = (

+        priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:],

+        priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:],

+        priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:],

+        priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:],

+        priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:],

+    )

+    landms = torch.cat(tmp, dim=1)

+    return landms

+

+

+def batched_decode(b_loc, priors, variances):

+    """Decode locations from predictions using priors to undo

+    the encoding we did for offset regression at train time.

+    Args:

+        b_loc (tensor): location predictions for loc layers,

+            Shape: [num_batches,num_priors,4]

+        priors (tensor): Prior boxes in center-offset form.

+            Shape: [1,num_priors,4].

+        variances: (list[float]) Variances of priorboxes

+    Return:

+        decoded bounding box predictions

+    """

+    boxes = (

+        priors[:, :, :2] + b_loc[:, :, :2] * variances[0] * priors[:, :, 2:],

+        priors[:, :, 2:] * torch.exp(b_loc[:, :, 2:] * variances[1]),

+    )

+    boxes = torch.cat(boxes, dim=2)

+

+    boxes[:, :, :2] -= boxes[:, :, 2:] / 2

+    boxes[:, :, 2:] += boxes[:, :, :2]

+    return boxes

+

+

+def batched_decode_landm(pre, priors, variances):

+    """Decode landm from predictions using priors to undo

+    the encoding we did for offset regression at train time.

+    Args:

+        pre (tensor): landm predictions for loc layers,

+            Shape: [num_batches,num_priors,10]

+        priors (tensor): Prior boxes in center-offset form.

+            Shape: [1,num_priors,4].

+        variances: (list[float]) Variances of priorboxes

+    Return:

+        decoded landm predictions

+    """

+    landms = (

+        priors[:, :, :2] + pre[:, :, :2] * variances[0] * priors[:, :, 2:],

+        priors[:, :, :2] + pre[:, :, 2:4] * variances[0] * priors[:, :, 2:],

+        priors[:, :, :2] + pre[:, :, 4:6] * variances[0] * priors[:, :, 2:],

+        priors[:, :, :2] + pre[:, :, 6:8] * variances[0] * priors[:, :, 2:],

+        priors[:, :, :2] + pre[:, :, 8:10] * variances[0] * priors[:, :, 2:],

+    )

+    landms = torch.cat(landms, dim=2)

+    return landms

+

+

+def log_sum_exp(x):

+    """Utility function for computing log_sum_exp while determining

+    This will be used to determine unaveraged confidence loss across

+    all examples in a batch.

+    Args:

+        x (Variable(tensor)): conf_preds from conf layers

+    """

+    x_max = x.data.max()

+    return torch.log(torch.sum(torch.exp(x - x_max), 1, keepdim=True)) + x_max

+

+

+# Original author: Francisco Massa:

+# https://github.com/fmassa/object-detection.torch

+# Ported to PyTorch by Max deGroot (02/01/2017)

+def nms(boxes, scores, overlap=0.5, top_k=200):

+    """Apply non-maximum suppression at test time to avoid detecting too many

+    overlapping bounding boxes for a given object.

+    Args:

+        boxes: (tensor) The location preds for the img, Shape: [num_priors,4].

+        scores: (tensor) The class predscores for the img, Shape:[num_priors].

+        overlap: (float) The overlap thresh for suppressing unnecessary boxes.

+        top_k: (int) The Maximum number of box preds to consider.

+    Return:

+        The indices of the kept boxes with respect to num_priors.

+    """

+

+    keep = torch.Tensor(scores.size(0)).fill_(0).long()

+    if boxes.numel() == 0:

+        return keep

+    x1 = boxes[:, 0]

+    y1 = boxes[:, 1]

+    x2 = boxes[:, 2]

+    y2 = boxes[:, 3]

+    area = torch.mul(x2 - x1, y2 - y1)

+    v, idx = scores.sort(0)  # sort in ascending order

+    # I = I[v >= 0.01]

+    idx = idx[-top_k:]  # indices of the top-k largest vals

+    xx1 = boxes.new()

+    yy1 = boxes.new()

+    xx2 = boxes.new()

+    yy2 = boxes.new()

+    w = boxes.new()

+    h = boxes.new()

+

+    # keep = torch.Tensor()

+    count = 0

+    while idx.numel() > 0:

+        i = idx[-1]  # index of current largest val

+        # keep.append(i)

+        keep[count] = i

+        count += 1

+        if idx.size(0) == 1:

+            break

+        idx = idx[:-1]  # remove kept element from view

+        # load bboxes of next highest vals

+        torch.index_select(x1, 0, idx, out=xx1)

+        torch.index_select(y1, 0, idx, out=yy1)

+        torch.index_select(x2, 0, idx, out=xx2)

+        torch.index_select(y2, 0, idx, out=yy2)

+        # store element-wise max with next highest score

+        xx1 = torch.clamp(xx1, min=x1[i])

+        yy1 = torch.clamp(yy1, min=y1[i])

+        xx2 = torch.clamp(xx2, max=x2[i])

+        yy2 = torch.clamp(yy2, max=y2[i])

+        w.resize_as_(xx2)

+        h.resize_as_(yy2)

+        w = xx2 - xx1

+        h = yy2 - yy1

+        # check sizes of xx1 and xx2.. after each iteration

+        w = torch.clamp(w, min=0.0)

+        h = torch.clamp(h, min=0.0)

+        inter = w * h

+        # IoU = i / (area(a) + area(b) - i)

+        rem_areas = torch.index_select(area, 0, idx)  # load remaining areas)

+        union = (rem_areas - inter) + area[i]

+        IoU = inter / union  # store result in iou

+        # keep only elements with an IoU <= overlap

+        idx = idx[IoU.le(overlap)]

+    return keep, count