moved retina face ops to model repo

retinaface/anchor.py

@@ -1,296 +0,0 @@
-"""Anchor utils modified from https://github.com/biubug6/Pytorch_Retinaface"""
-import math
-import tensorflow as tf
-import numpy as np
-from itertools import product as product
-
-
-###############################################################################
-#   Tensorflow / Numpy Priors                                                 #
-###############################################################################
-def prior_box(image_sizes, min_sizes, steps, clip=False):
-    """prior box"""
-    feature_maps = [
-        [math.ceil(image_sizes[0] / step), math.ceil(image_sizes[1] / step)]
-        for step in steps]
-
-    anchors = []
-    for k, f in enumerate(feature_maps):
-        for i, j in product(range(f[0]), range(f[1])):
-            for min_size in min_sizes[k]:
-                s_kx = min_size / image_sizes[1]
-                s_ky = min_size / image_sizes[0]
-                cx = (j + 0.5) * steps[k] / image_sizes[1]
-                cy = (i + 0.5) * steps[k] / image_sizes[0]
-                anchors += [cx, cy, s_kx, s_ky]
-
-    output = np.asarray(anchors).reshape([-1, 4])
-
-    if clip:
-        output = np.clip(output, 0, 1)
-
-    return output
-
-
-def prior_box_tf(image_sizes, min_sizes, steps, clip=False):
-    """prior box"""
-    image_sizes = tf.cast(tf.convert_to_tensor(image_sizes), tf.float32)
-    feature_maps = tf.math.ceil(
-        tf.reshape(image_sizes, [1, 2]) /
-        tf.reshape(tf.cast(steps, tf.float32), [-1, 1]))
-
-    anchors = []
-    for k in range(len(min_sizes)):
-        grid_x, grid_y = _meshgrid_tf(tf.range(feature_maps[k][1]),
-                                      tf.range(feature_maps[k][0]))
-        cx = (grid_x + 0.5) * steps[k] / image_sizes[1]
-        cy = (grid_y + 0.5) * steps[k] / image_sizes[0]
-        cxcy = tf.stack([cx, cy], axis=-1)
-        cxcy = tf.reshape(cxcy, [-1, 2])
-        cxcy = tf.repeat(cxcy, repeats=tf.shape(min_sizes[k])[0], axis=0)
-
-        sx = min_sizes[k] / image_sizes[1]
-        sy = min_sizes[k] / image_sizes[0]
-        sxsy = tf.stack([sx, sy], 1)
-        sxsy = tf.repeat(sxsy[tf.newaxis],
-                         repeats=tf.shape(grid_x)[0] * tf.shape(grid_x)[1],
-                         axis=0)
-        sxsy = tf.reshape(sxsy, [-1, 2])
-
-        anchors.append(tf.concat([cxcy, sxsy], 1))
-
-    output = tf.concat(anchors, axis=0)
-
-    if clip:
-        output = tf.clip_by_value(output, 0, 1)
-
-    return output
-
-
-def _meshgrid_tf(x, y):
-    """ workaround solution of the tf.meshgrid() issue:
-        https://github.com/tensorflow/tensorflow/issues/34470"""
-    grid_shape = [tf.shape(y)[0], tf.shape(x)[0]]
-    grid_x = tf.broadcast_to(tf.reshape(x, [1, -1]), grid_shape)
-    grid_y = tf.broadcast_to(tf.reshape(y, [-1, 1]), grid_shape)
-    return grid_x, grid_y
-
-
-###############################################################################
-#   Tensorflow Encoding                                                       #
-###############################################################################
-def encode_tf(labels, priors, match_thresh, ignore_thresh,
-              variances=[0.1, 0.2]):
-    """tensorflow encoding"""
-    assert ignore_thresh <= match_thresh
-    priors = tf.cast(priors, tf.float32)
-    bbox = labels[:, :4]
-    landm = labels[:, 4:-1]
-    landm_valid = labels[:, -1]  # 1: with landm, 0: w/o landm.
-
-    # jaccard index
-    overlaps = _jaccard(bbox, _point_form(priors))
-
-    # (Bipartite Matching)
-    # [num_objects] best prior for each ground truth
-    best_prior_overlap, best_prior_idx = tf.math.top_k(overlaps, k=1)
-    best_prior_overlap = best_prior_overlap[:, 0]
-    best_prior_idx = best_prior_idx[:, 0]
-
-    # [num_priors] best ground truth for each prior
-    overlaps_t = tf.transpose(overlaps)
-    best_truth_overlap, best_truth_idx = tf.math.top_k(overlaps_t, k=1)
-    best_truth_overlap = best_truth_overlap[:, 0]
-    best_truth_idx = best_truth_idx[:, 0]
-
-    # ensure best prior
-    def _loop_body(i, bt_idx, bt_overlap):
-        bp_mask = tf.one_hot(best_prior_idx[i], tf.shape(bt_idx)[0])
-        bp_mask_int = tf.cast(bp_mask, tf.int32)
-        new_bt_idx = bt_idx * (1 - bp_mask_int) + bp_mask_int * i
-        bp_mask_float = tf.cast(bp_mask, tf.float32)
-        new_bt_overlap = bt_overlap * (1 - bp_mask_float) + bp_mask_float * 2
-        return tf.cond(best_prior_overlap[i] > match_thresh,
-                       lambda: (i + 1, new_bt_idx, new_bt_overlap),
-                       lambda: (i + 1, bt_idx, bt_overlap))
-    _, best_truth_idx, best_truth_overlap = tf.while_loop(
-        lambda i, bt_idx, bt_overlap: tf.less(i, tf.shape(best_prior_idx)[0]),
-        _loop_body, [tf.constant(0), best_truth_idx, best_truth_overlap])
-
-    matches_bbox = tf.gather(bbox, best_truth_idx)  # [num_priors, 4]
-    matches_landm = tf.gather(landm, best_truth_idx)  # [num_priors, 10]
-    matches_landm_v = tf.gather(landm_valid, best_truth_idx)  # [num_priors]
-
-    loc_t = _encode_bbox(matches_bbox, priors, variances)
-    landm_t = _encode_landm(matches_landm, priors, variances)
-    landm_valid_t = tf.cast(matches_landm_v > 0, tf.float32)
-    conf_t = tf.cast(best_truth_overlap > match_thresh, tf.float32)
-    conf_t = tf.where(
-        tf.logical_and(best_truth_overlap < match_thresh,
-                       best_truth_overlap > ignore_thresh),
-        tf.ones_like(conf_t) * -1, conf_t)    # 1: pos, 0: neg, -1: ignore
-
-    return tf.concat([loc_t, landm_t, landm_valid_t[..., tf.newaxis],
-                      conf_t[..., tf.newaxis]], axis=1)
-
-
-def _encode_bbox(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 = tf.math.log(g_wh) / variances[1]
-    # return target for smooth_l1_loss
-    return tf.concat([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 = tf.reshape(matched, [tf.shape(matched)[0], 5, 2])
-    priors = tf.broadcast_to(
-        tf.expand_dims(priors, 1), [tf.shape(matched)[0], 5, 4])
-    g_cxcy = matched[:, :, :2] - priors[:, :, :2]
-    # encode variance
-    g_cxcy /= (variances[0] * priors[:, :, 2:])
-    # g_cxcy /= priors[:, :, 2:]
-    g_cxcy = tf.reshape(g_cxcy, [tf.shape(g_cxcy)[0], -1])
-    # return target for smooth_l1_loss
-    return g_cxcy
-
-
-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 tf.concat((boxes[:, :2] - boxes[:, 2:] / 2,
-                      boxes[:, :2] + boxes[:, 2:] / 2), axis=1)
-
-
-def _intersect(box_a, box_b):
-    """ We resize both tensors to [A,B,2]:
-    [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 = tf.shape(box_a)[0]
-    B = tf.shape(box_b)[0]
-    max_xy = tf.minimum(
-        tf.broadcast_to(tf.expand_dims(box_a[:, 2:], 1), [A, B, 2]),
-        tf.broadcast_to(tf.expand_dims(box_b[:, 2:], 0), [A, B, 2]))
-    min_xy = tf.maximum(
-        tf.broadcast_to(tf.expand_dims(box_a[:, :2], 1), [A, B, 2]),
-        tf.broadcast_to(tf.expand_dims(box_b[:, :2], 0), [A, B, 2]))
-    inter = tf.maximum((max_xy - min_xy), tf.zeros_like(max_xy - min_xy))
-    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 = tf.broadcast_to(
-        tf.expand_dims(
-            (box_a[:, 2] - box_a[:, 0]) * (box_a[:, 3] - box_a[:, 1]), 1),
-        tf.shape(inter))  # [A,B]
-    area_b = tf.broadcast_to(
-        tf.expand_dims(
-            (box_b[:, 2] - box_b[:, 0]) * (box_b[:, 3] - box_b[:, 1]), 0),
-        tf.shape(inter))  # [A,B]
-    union = area_a + area_b - inter
-    return inter / union  # [A,B]
-
-
-###############################################################################
-#   Tensorflow Decoding                                                       #
-###############################################################################
-def decode_tf(labels, priors, variances=[0.1, 0.2]):
-    """tensorflow decoding"""
-    bbox = _decode_bbox(labels[:, :4], priors, variances)
-    landm = _decode_landm(labels[:, 4:14], priors, variances)
-    landm_valid = labels[:, 14][:, tf.newaxis]
-    conf = labels[:, 15][:, tf.newaxis]
-
-    return tf.concat([bbox, landm, landm_valid, conf], axis=1)
-
-
-def _decode_bbox(pre, priors, variances=[0.1, 0.2]):
-    """Decode locations from predictions using priors to undo
-    the encoding we did for offset regression at train time.
-    Args:
-        pre (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
-    """
-    centers = priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:]
-    sides = priors[:, 2:] * tf.math.exp(pre[:, 2:] * variances[1])
-
-    return tf.concat([centers - sides / 2, centers + sides / 2], axis=1)
-
-
-def _decode_landm(pre, priors, variances=[0.1, 0.2]):
-    """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
-    """
-    landms = tf.concat(
-        [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:]], axis=1)
-    return landms

retinaface/models.py

@@ -1,301 +0,0 @@
-import tensorflow as tf
-from tensorflow.keras import Model
-from tensorflow.keras.applications import MobileNetV2, ResNet50
-from tensorflow.keras.layers import Input, Conv2D, ReLU, LeakyReLU
-from retinaface.anchor import decode_tf, prior_box_tf
-
-
-def _regularizer(weights_decay):
-    """l2 regularizer"""
-    return tf.keras.regularizers.l2(weights_decay)
-
-
-def _kernel_init(scale=1.0, seed=None):
-    """He normal initializer"""
-    return tf.keras.initializers.he_normal()
-
-
-class BatchNormalization(tf.keras.layers.BatchNormalization):
-    """Make trainable=False freeze BN for real (the og version is sad).
-       ref: https://github.com/zzh8829/yolov3-tf2
-    """
-    def __init__(self, axis=-1, momentum=0.9, epsilon=1e-5, center=True,
-                 scale=True, name=None, **kwargs):
-        super(BatchNormalization, self).__init__(
-            axis=axis, momentum=momentum, epsilon=epsilon, center=center,
-            scale=scale, name=name, **kwargs)
-
-    def call(self, x, training=False):
-        if training is None:
-            training = tf.constant(False)
-        training = tf.logical_and(training, self.trainable)
-
-        return super().call(x, training)
-
-
-def Backbone(backbone_type='ResNet50', use_pretrain=True):
-    """Backbone Model"""
-    weights = None
-    if use_pretrain:
-        weights = 'imagenet'
-
-    def backbone(x):
-        if backbone_type == 'ResNet50':
-            extractor = ResNet50(
-                input_shape=x.shape[1:], include_top=False, weights=weights)
-            pick_layer1 = 80  # [80, 80, 512]
-            pick_layer2 = 142  # [40, 40, 1024]
-            pick_layer3 = 174  # [20, 20, 2048]
-            preprocess = tf.keras.applications.resnet.preprocess_input
-        elif backbone_type == 'MobileNetV2':
-            extractor = MobileNetV2(
-                input_shape=x.shape[1:], include_top=False, weights=weights)
-            pick_layer1 = 54  # [80, 80, 32]
-            pick_layer2 = 116  # [40, 40, 96]
-            pick_layer3 = 143  # [20, 20, 160]
-            preprocess = tf.keras.applications.mobilenet_v2.preprocess_input
-        else:
-            raise NotImplementedError(
-                'Backbone type {} is not recognized.'.format(backbone_type))
-
-        return Model(extractor.input,
-                     (extractor.layers[pick_layer1].output,
-                      extractor.layers[pick_layer2].output,
-                      extractor.layers[pick_layer3].output),
-                     name=backbone_type + '_extrator')(preprocess(x))
-
-    return backbone
-
-
-class ConvUnit(tf.keras.layers.Layer):
-    """Conv + BN + Act"""
-    def __init__(self, f, k, s, wd, act=None, **kwargs):
-        super(ConvUnit, self).__init__(**kwargs)
-        self.conv = Conv2D(filters=f, kernel_size=k, strides=s, padding='same',
-                           kernel_initializer=_kernel_init(),
-                           kernel_regularizer=_regularizer(wd),
-                           use_bias=False)
-        self.bn = BatchNormalization()
-
-        if act is None:
-            self.act_fn = tf.identity
-        elif act == 'relu':
-            self.act_fn = ReLU()
-        elif act == 'lrelu':
-            self.act_fn = LeakyReLU(0.1)
-        else:
-            raise NotImplementedError(
-                'Activation function type {} is not recognized.'.format(act))
-
-    def call(self, x):
-        return self.act_fn(self.bn(self.conv(x)))
-
-
-class FPN(tf.keras.layers.Layer):
-    """Feature Pyramid Network"""
-    def __init__(self, out_ch, wd, **kwargs):
-        super(FPN, self).__init__(**kwargs)
-        act = 'relu'
-        self.out_ch = out_ch
-        self.wd = wd
-        if (out_ch <= 64):
-            act = 'lrelu'
-
-        self.output1 = ConvUnit(f=out_ch, k=1, s=1, wd=wd, act=act)
-        self.output2 = ConvUnit(f=out_ch, k=1, s=1, wd=wd, act=act)
-        self.output3 = ConvUnit(f=out_ch, k=1, s=1, wd=wd, act=act)
-        self.merge1 = ConvUnit(f=out_ch, k=3, s=1, wd=wd, act=act)
-        self.merge2 = ConvUnit(f=out_ch, k=3, s=1, wd=wd, act=act)
-
-    def call(self, x):
-        output1 = self.output1(x[0])  # [80, 80, out_ch]
-        output2 = self.output2(x[1])  # [40, 40, out_ch]
-        output3 = self.output3(x[2])  # [20, 20, out_ch]
-
-        up_h, up_w = tf.shape(output2)[1], tf.shape(output2)[2]
-        up3 = tf.image.resize(output3, [up_h, up_w], method='nearest')
-        output2 = output2 + up3
-        output2 = self.merge2(output2)
-
-        up_h, up_w = tf.shape(output1)[1], tf.shape(output1)[2]
-        up2 = tf.image.resize(output2, [up_h, up_w], method='nearest')
-        output1 = output1 + up2
-        output1 = self.merge1(output1)
-
-        return output1, output2, output3
-    
-    def get_config(self):
-        config = {
-            'out_ch': self.out_ch,
-            'wd': self.wd,
-        }
-        base_config = super(FPN, self).get_config()
-        return dict(list(base_config.items()) + list(config.items()))
-
-
-class SSH(tf.keras.layers.Layer):
-    """Single Stage Headless Layer"""
-    def __init__(self, out_ch, wd, **kwargs):
-        super(SSH, self).__init__(**kwargs)
-        assert out_ch % 4 == 0
-        self.out_ch = out_ch
-        self.wd = wd
-        act = 'relu'
-        if (out_ch <= 64):
-            act = 'lrelu'
-
-        self.conv_3x3 = ConvUnit(f=out_ch // 2, k=3, s=1, wd=wd, act=None)
-
-        self.conv_5x5_1 = ConvUnit(f=out_ch // 4, k=3, s=1, wd=wd, act=act)
-        self.conv_5x5_2 = ConvUnit(f=out_ch // 4, k=3, s=1, wd=wd, act=None)
-
-        self.conv_7x7_2 = ConvUnit(f=out_ch // 4, k=3, s=1, wd=wd, act=act)
-        self.conv_7x7_3 = ConvUnit(f=out_ch // 4, k=3, s=1, wd=wd, act=None)
-
-        self.relu = ReLU()
-
-    def call(self, x):
-        conv_3x3 = self.conv_3x3(x)
-
-        conv_5x5_1 = self.conv_5x5_1(x)
-        conv_5x5 = self.conv_5x5_2(conv_5x5_1)
-
-        conv_7x7_2 = self.conv_7x7_2(conv_5x5_1)
-        conv_7x7 = self.conv_7x7_3(conv_7x7_2)
-
-        output = tf.concat([conv_3x3, conv_5x5, conv_7x7], axis=3)
-        output = self.relu(output)
-
-        return output
-    
-    def get_config(self):
-        config = {
-            'out_ch': self.out_ch,
-            'wd': self.wd,
-        }
-        base_config = super(SSH, self).get_config()
-        return dict(list(base_config.items()) + list(config.items()))
-
-
-class BboxHead(tf.keras.layers.Layer):
-    """Bbox Head Layer"""
-    def __init__(self, num_anchor, wd, **kwargs):
-        super(BboxHead, self).__init__(**kwargs)
-        self.num_anchor = num_anchor
-        self.wd = wd
-        self.conv = Conv2D(filters=num_anchor * 4, kernel_size=1, strides=1)
-
-    def call(self, x):
-        h, w = tf.shape(x)[1], tf.shape(x)[2]
-        x = self.conv(x)
-
-        return tf.reshape(x, [-1, h * w * self.num_anchor, 4])
-    
-    def get_config(self):
-        config = {
-            'num_anchor': self.num_anchor,
-            'wd': self.wd,
-        }
-        base_config = super(BboxHead, self).get_config()
-        return dict(list(base_config.items()) + list(config.items()))
-
-
-class LandmarkHead(tf.keras.layers.Layer):
-    """Landmark Head Layer"""
-    def __init__(self, num_anchor, wd, name='LandmarkHead', **kwargs):
-        super(LandmarkHead, self).__init__(name=name, **kwargs)
-        self.num_anchor = num_anchor
-        self.wd = wd
-        self.conv = Conv2D(filters=num_anchor * 10, kernel_size=1, strides=1)
-
-    def call(self, x):
-        h, w = tf.shape(x)[1], tf.shape(x)[2]
-        x = self.conv(x)
-
-        return tf.reshape(x, [-1, h * w * self.num_anchor, 10])
-
-    def get_config(self):
-        config = {
-            'num_anchor': self.num_anchor,
-            'wd': self.wd,
-        }
-        base_config = super(LandmarkHead, self).get_config()
-        return dict(list(base_config.items()) + list(config.items()))
-
-
-class ClassHead(tf.keras.layers.Layer):
-    """Class Head Layer"""
-    def __init__(self, num_anchor, wd, name='ClassHead', **kwargs):
-        super(ClassHead, self).__init__(name=name, **kwargs)
-        self.num_anchor = num_anchor
-        self.wd = wd
-        self.conv = Conv2D(filters=num_anchor * 2, kernel_size=1, strides=1)
-
-    def call(self, x):
-        h, w = tf.shape(x)[1], tf.shape(x)[2]
-        x = self.conv(x)
-
-        return tf.reshape(x, [-1, h * w * self.num_anchor, 2])
-
-    def get_config(self):
-        config = {
-            'num_anchor': self.num_anchor,
-            'wd': self.wd,
-        }
-        base_config = super(ClassHead, self).get_config()
-        return dict(list(base_config.items()) + list(config.items()))
-
-
-def RetinaFaceModel(cfg, training=False, iou_th=0.4, score_th=0.02,
-                    name='RetinaFaceModel'):
-    """Retina Face Model"""
-    input_size = cfg['input_size'] if training else None
-    wd = cfg['weights_decay']
-    out_ch = cfg['out_channel']
-    num_anchor = len(cfg['min_sizes'][0])
-    backbone_type = cfg['backbone_type']
-
-    # define model
-    x = inputs = Input([input_size, input_size, 3], name='input_image')
-
-    x = Backbone(backbone_type=backbone_type)(x)
-
-    fpn = FPN(out_ch=out_ch, wd=wd)(x)
-
-    features = [SSH(out_ch=out_ch, wd=wd)(f)
-                for i, f in enumerate(fpn)]
-
-    bbox_regressions = tf.concat(
-        [BboxHead(num_anchor, wd=wd)(f)
-         for i, f in enumerate(features)], axis=1)
-    landm_regressions = tf.concat(
-        [LandmarkHead(num_anchor, wd=wd, name=f'LandmarkHead_{i}')(f)
-         for i, f in enumerate(features)], axis=1)
-    classifications = tf.concat(
-        [ClassHead(num_anchor, wd=wd, name=f'ClassHead_{i}')(f)
-         for i, f in enumerate(features)], axis=1)
-
-    classifications = tf.keras.layers.Softmax(axis=-1)(classifications)
-
-    if training:
-        out = (bbox_regressions, landm_regressions, classifications)
-    else:
-        # only for batch size 1
-        preds = tf.concat(  # [bboxes, landms, landms_valid, conf]
-            [bbox_regressions[0],
-             landm_regressions[0],
-             tf.ones_like(classifications[0, :, 0][..., tf.newaxis]),
-             classifications[0, :, 1][..., tf.newaxis]], 1)
-        priors = prior_box_tf((tf.shape(inputs)[1], tf.shape(inputs)[2]), cfg['min_sizes'], cfg['steps'], cfg['clip'])
-        decode_preds = decode_tf(preds, priors, cfg['variances'])
-
-        selected_indices = tf.image.non_max_suppression(
-            boxes=decode_preds[:, :4],
-            scores=decode_preds[:, -1],
-            max_output_size=tf.shape(decode_preds)[0],
-            iou_threshold=iou_th,
-            score_threshold=score_th)
-
-        out = tf.gather(decode_preds, selected_indices)
-
-    return Model(inputs, out, name=name), Model(inputs, [bbox_regressions, landm_regressions, classifications], name=name + '_bb_only')

retinaface/ops.py

@@ -1,27 +0,0 @@
-from retinaface.anchor import decode_tf, prior_box_tf
-import tensorflow as tf
-
-
-def extract_detections(bbox_regressions, landm_regressions, classifications, image_sizes, iou_th=0.4, score_th=0.02):
-    min_sizes = [[16, 32], [64, 128], [256, 512]]
-    steps = [8, 16, 32]
-    variances = [0.1, 0.2]
-    preds = tf.concat(  # [bboxes, landms, landms_valid, conf]
-        [bbox_regressions,
-         landm_regressions,
-         tf.ones_like(classifications[:, 0][..., tf.newaxis]),
-         classifications[:, 1][..., tf.newaxis]], 1)
-    priors = prior_box_tf(image_sizes, min_sizes, steps, False)
-    decode_preds = decode_tf(preds, priors, variances)
-
-    selected_indices = tf.image.non_max_suppression(
-        boxes=decode_preds[:, :4],
-        scores=decode_preds[:, -1],
-        max_output_size=tf.shape(decode_preds)[0],
-        iou_threshold=iou_th,
-        score_threshold=score_th)
-
-    out = tf.gather(decode_preds, selected_indices)
-
-    return out
-