modelinhos 0.0.0__py3-none-any.whl

modelinhos/blaze/blazenet.py

@@ -0,0 +1,278 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BlazeBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1):
+        super(BlazeBlock, self).__init__()
+
+        self.stride = stride
+        self.channel_pad = out_channels - in_channels
+
+        # TFLite uses slightly different padding than PyTorch
+        # on the depthwise conv layer when the stride is 2.
+        if stride == 2:
+            self.max_pool = nn.MaxPool2d(kernel_size=stride, stride=stride)
+            padding = 0
+        else:
+            padding = (kernel_size - 1) // 2
+
+        self.convs = nn.Sequential(
+            nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=padding,
+                groups=in_channels,
+                bias=True,
+            ),
+            nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                bias=True,
+            ),
+        )
+
+        self.act = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        if self.stride == 2:
+            h = F.pad(x, (0, 2, 0, 2), "constant", 0)
+            x = self.max_pool(x)
+        else:
+            h = x
+
+        if self.channel_pad > 0:
+            x = F.pad(x, (0, 0, 0, 0, 0, self.channel_pad), "constant", 0)
+
+        return self.act(self.convs(h) + x)
+
+
+class FinalBlazeBlock(nn.Module):
+    def __init__(self, channels, kernel_size=3):
+        super(FinalBlazeBlock, self).__init__()
+        # TFLite uses slightly different padding than PyTorch
+        # on the depthwise conv layer when the stride is 2.
+        self.convs = nn.Sequential(
+            nn.Conv2d(
+                in_channels=channels,
+                out_channels=channels,
+                kernel_size=kernel_size,
+                stride=2,
+                padding=0,
+                groups=channels,
+                bias=True,
+            ),
+            nn.Conv2d(
+                in_channels=channels,
+                out_channels=channels,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                bias=True,
+            ),
+        )
+
+        self.act = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        h = F.pad(x, (0, 2, 0, 2), "constant", 0)
+
+        return self.act(self.convs(h))
+
+
+class BlazeNet(nn.Module):
+    """The BlazeFace face detection model from MediaPipe.
+
+    The version from MediaPipe is simpler than the one in the paper;
+    it does not use the "double" BlazeBlocks.
+
+    Because we won't be training this model, it doesn't need to have
+    batchnorm layers. These have already been "folded" into the conv
+    weights by TFLite.
+
+    The conversion to PyTorch is fairly straightforward, but there are
+    some small differences between TFLite and PyTorch in how they handle
+    padding on conv layers with stride 2.
+
+    This version works on batches, while the MediaPipe version can only
+    handle a single image at a time.
+
+    Based on code from https://github.com/tkat0/PyTorch_BlazeFace/ and
+    https://github.com/google/mediapipe/
+    """
+
+    def __init__(self, back_model=False):
+        super(BlazeNet, self).__init__()
+
+        # These are the settings from the MediaPipe example graphs
+        # mediapipe/graphs/face_detection/face_detection_mobile_gpu.pbtxt
+        # and
+        # mediapipe/graphs/face_detection/face_detection_back_mobile_gpu.pbtxt
+        self.num_classes = 1
+        self.num_anchors = 896
+        self.num_coords = 16
+        self.score_clipping_thresh = 100.0
+        self.back_model = back_model
+        if back_model:
+            self.x_scale = 256.0
+            self.y_scale = 256.0
+            self.h_scale = 256.0
+            self.w_scale = 256.0
+            self.min_score_thresh = 0.65
+        else:
+            self.x_scale = 128.0
+            self.y_scale = 128.0
+            self.h_scale = 128.0
+            self.w_scale = 128.0
+            self.min_score_thresh = 0.75
+        self.min_suppression_threshold = 0.3
+
+        self._define_layers()
+
+    def _define_layers(self):
+        if self.back_model:
+            self.backbone = nn.Sequential(
+                nn.Conv2d(
+                    in_channels=3,
+                    out_channels=24,
+                    kernel_size=5,
+                    stride=2,
+                    padding=0,
+                    bias=True,
+                ),
+                nn.ReLU(inplace=True),
+                BlazeBlock(24, 24),
+                BlazeBlock(24, 24),
+                BlazeBlock(24, 24),
+                BlazeBlock(24, 24),
+                BlazeBlock(24, 24),
+                BlazeBlock(24, 24),
+                BlazeBlock(24, 24),
+                BlazeBlock(24, 24, stride=2),
+                BlazeBlock(24, 24),
+                BlazeBlock(24, 24),
+                BlazeBlock(24, 24),
+                BlazeBlock(24, 24),
+                BlazeBlock(24, 24),
+                BlazeBlock(24, 24),
+                BlazeBlock(24, 24),
+                BlazeBlock(24, 48, stride=2),
+                BlazeBlock(48, 48),
+                BlazeBlock(48, 48),
+                BlazeBlock(48, 48),
+                BlazeBlock(48, 48),
+                BlazeBlock(48, 48),
+                BlazeBlock(48, 48),
+                BlazeBlock(48, 48),
+                BlazeBlock(48, 96, stride=2),
+                BlazeBlock(96, 96),
+                BlazeBlock(96, 96),
+                BlazeBlock(96, 96),
+                BlazeBlock(96, 96),
+                BlazeBlock(96, 96),
+                BlazeBlock(96, 96),
+                BlazeBlock(96, 96),
+            )
+            self.final = FinalBlazeBlock(96)
+            self.classifier_8 = nn.Conv2d(96, 2, 1, bias=True)
+            self.classifier_16 = nn.Conv2d(96, 6, 1, bias=True)
+
+            self.regressor_8 = nn.Conv2d(96, 32, 1, bias=True)
+            self.regressor_16 = nn.Conv2d(96, 96, 1, bias=True)
+        else:
+            self.backbone1 = nn.Sequential(
+                nn.Conv2d(
+                    in_channels=3,
+                    out_channels=24,
+                    kernel_size=5,
+                    stride=2,
+                    padding=0,
+                    bias=True,
+                ),
+                nn.ReLU(inplace=True),
+                BlazeBlock(24, 24),
+                BlazeBlock(24, 28),
+                BlazeBlock(28, 32, stride=2),
+                BlazeBlock(32, 36),
+                BlazeBlock(36, 42),
+                BlazeBlock(42, 48, stride=2),
+                BlazeBlock(48, 56),
+                BlazeBlock(56, 64),
+                BlazeBlock(64, 72),
+                BlazeBlock(72, 80),
+                BlazeBlock(80, 88),
+            )
+
+            self.backbone2 = nn.Sequential(
+                BlazeBlock(88, 96, stride=2),
+                BlazeBlock(96, 96),
+                BlazeBlock(96, 96),
+                BlazeBlock(96, 96),
+                BlazeBlock(96, 96),
+            )
+            self.classifier_8 = nn.Conv2d(88, 2, 1, bias=True)
+            self.classifier_16 = nn.Conv2d(96, 6, 1, bias=True)
+
+            self.regressor_8 = nn.Conv2d(88, 32, 1, bias=True)
+            self.regressor_16 = nn.Conv2d(96, 96, 1, bias=True)
+
+    def forward(self, image):
+        # TFLite uses slightly different padding on the first conv layer
+        # than PyTorch, so do it manually.
+        x = F.pad(image, (1, 2, 1, 2), "constant", 0)
+
+        b = x.shape[0]  # batch size, needed for reshaping later
+
+        if self.back_model:
+            x = self.backbone(x)  # (b, 16, 16, 96)
+            h = self.final(x)  # (b, 8, 8, 96)
+        else:
+            x = self.backbone1(x)  # (b, 88, 16, 16)
+            h = self.backbone2(x)  # (b, 96, 8, 8)
+
+        # Note: Because PyTorch is NCHW but TFLite is NHWC, we need to
+        # permute the output from the conv layers before reshaping it.
+
+        c1 = self.classifier_8(x)  # (b, 2, 16, 16)
+        c1 = c1.permute(0, 2, 3, 1)  # (b, 16, 16, 2)
+        c1 = c1.reshape(b, -1, 1)  # (b, 512, 1)
+
+        c2 = self.classifier_16(h)  # (b, 6, 8, 8)
+        c2 = c2.permute(0, 2, 3, 1)  # (b, 8, 8, 6)
+        c2 = c2.reshape(b, -1, 1)  # (b, 384, 1)
+
+        c = torch.cat((c1, c2), dim=1)  # (b, 896, 1)
+
+        r1 = self.regressor_8(x)  # (b, 32, 16, 16)
+        r1 = r1.permute(0, 2, 3, 1)  # (b, 16, 16, 32)
+        r1 = r1.reshape(b, -1, 16)  # (b, 512, 16)
+
+        r2 = self.regressor_16(h)  # (b, 96, 8, 8)
+        r2 = r2.permute(0, 2, 3, 1)  # (b, 8, 8, 96)
+        r2 = r2.reshape(b, -1, 16)  # (b, 384, 16)
+
+        r = torch.cat((r1, r2), dim=1)  # (b, 896, 16)
+        return [r, c]
+
+
+def load_weights(model: BlazeNet, path):
+    model.load_state_dict(torch.load(path))
+    model.eval()
+
+
+def load_anchors(model: BlazeNet, path):
+    model.anchors = torch.tensor(
+        np.load(path),
+        dtype=torch.float32,
+        device=model.classifier_8.weight.device,
+    )
+    assert model.anchors.ndimension() == 2
+    assert model.anchors.shape[0] == model.num_anchors
+    assert model.anchors.shape[1] == 4

modelinhos/blaze/infer.py

@@ -0,0 +1,93 @@
+import cv2
+import numpy as np
+import torch
+
+from modelinhos.blazenet import BlazeNet
+
+EXPECTED = np.array(
+    [
+        [
+            0.2763,
+            0.3182,
+            0.4465,
+            0.4884,
+            0.3830,
+            0.3150,
+            0.4561,
+            0.3202,
+            0.4309,
+            0.3526,
+            0.4229,
+            0.3913,
+            0.3182,
+            0.3373,
+            0.4769,
+            0.3464,
+            0.9308,
+        ]
+    ],
+)
+
+
+def plot(image, detections, with_keypoints=True):
+    visualized = image.copy()
+
+    if isinstance(detections, torch.Tensor):
+        detections = detections.cpu().numpy()
+
+    if detections.ndim == 1:
+        detections = np.expand_dims(detections, axis=0)
+
+    print("Found %d faces" % detections.shape[0])
+
+    for i in range(detections.shape[0]):
+        ymin = int(detections[i, 0] * image.shape[0])
+        xmin = int(detections[i, 1] * image.shape[1])
+        ymax = int(detections[i, 2] * image.shape[0])
+        xmax = int(detections[i, 3] * image.shape[1])
+
+        cv2.rectangle(
+            visualized,
+            (xmin, ymin),
+            (xmax, ymax),
+            color=(0, 0, 255),  # red in BGR
+            thickness=1,
+        )
+
+        if with_keypoints:
+            for k in range(6):
+                kp_x = int(detections[i, 4 + k * 2] * image.shape[1])
+                kp_y = int(detections[i, 4 + k * 2 + 1] * image.shape[0])
+
+                cv2.circle(
+                    visualized,
+                    (kp_x, kp_y),
+                    radius=2,
+                    color=(255, 200, 100),  # light-sky-blue-ish in BGR
+                    thickness=1,
+                )
+
+    return visualized
+
+
+def main():
+    front_net = BlazeNet()
+    front_net.load_weights("blazeface.pth")
+    front_net.load_anchors("anchors.npy")
+    front_net.min_score_thresh = 0.75
+    front_net.min_suppression_threshold = 0.3
+    image = cv2.imread("1face.png")
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    predictions = front_net.predict_on_image(image)
+    np.testing.assert_almost_equal(
+        predictions.cpu().numpy(),
+        EXPECTED,
+    )
+    visualized = plot(image, predictions)
+    cv2.imshow("Detections", visualized)
+    cv2.waitKey(0)
+    cv2.destroyAllWindows()
+
+
+if __name__ == "__main__":
+    main()

modelinhos/blaze/postprocessing.py

@@ -0,0 +1,302 @@
+import numpy as np
+import torch
+
+from modelinhos.blaze.blazenet import BlazeNet
+
+
+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 overlap_similarity(box, other_boxes):
+    """Computes the IOU between a bounding box and set of other boxes."""
+    return jaccard(box.unsqueeze(0), other_boxes).squeeze(0)
+
+
+def _weighted_non_max_suppression(
+    model: BlazeNet,
+    detections,
+    min_suppression_threshold: int,
+):
+    """The alternative NMS method as mentioned in the BlazeFace paper:
+
+    "We replace the suppression algorithm with a blending strategy that
+    estimates the regression parameters of a bounding box as a weighted
+    mean between the overlapping predictions."
+
+    The original MediaPipe code assigns the score of the most confident
+    detection to the weighted detection, but we take the average score
+    of the overlapping detections.
+
+    The input detections should be a Tensor of shape (count, 17).
+
+    Returns a list of PyTorch tensors, one for each detected face.
+
+    This is based on the source code from:
+    mediapipe/calculators/util/non_max_suppression_calculator.cc
+    mediapipe/calculators/util/non_max_suppression_calculator.proto
+    """
+    if len(detections) == 0:
+        return []
+
+    output_detections = []
+
+    # Sort the detections from highest to lowest score.
+    remaining = torch.argsort(detections[:, 16], descending=True)
+
+    while len(remaining) > 0:
+        detection = detections[remaining[0]]
+
+        # Compute the overlap between the first box and the other
+        # remaining boxes. (Note that the other_boxes also include
+        # the first_box.)
+        first_box = detection[:4]
+        other_boxes = detections[remaining, :4]
+        ious = overlap_similarity(first_box, other_boxes)
+
+        # If two detections don't overlap enough, they are considered
+        # to be from different faces.
+        mask = ious > min_suppression_threshold
+        overlapping = remaining[mask]
+        remaining = remaining[~mask]
+
+        # Take an average of the coordinates from the overlapping
+        # detections, weighted by their confidence scores.
+        weighted_detection = detection.clone()
+        if len(overlapping) > 1:
+            coordinates = detections[overlapping, :16]
+            scores = detections[overlapping, 16:17]
+            total_score = scores.sum()
+            weighted = (coordinates * scores).sum(dim=0) / total_score
+            weighted_detection[:16] = weighted
+            weighted_detection[16] = total_score / len(overlapping)
+
+        output_detections.append(weighted_detection)
+
+    return output_detections
+
+
+def _decode_boxes(model: BlazeNet, raw, anchors):
+    """Converts the predictions into actual coordinates using
+    the anchor boxes. Processes the entire batch at once.
+    """
+    boxes = torch.zeros_like(raw)
+
+    x_center = raw[..., 0] / model.x_scale * anchors[:, 2] + anchors[:, 0]
+    y_center = raw[..., 1] / model.y_scale * anchors[:, 3] + anchors[:, 1]
+
+    w = raw[..., 2] / model.w_scale * anchors[:, 2]
+    h = raw[..., 3] / model.h_scale * anchors[:, 3]
+
+    boxes[..., 0] = y_center - h / 2.0  # ymin
+    boxes[..., 1] = x_center - w / 2.0  # xmin
+    boxes[..., 2] = y_center + h / 2.0  # ymax
+    boxes[..., 3] = x_center + w / 2.0  # xmax
+
+    for k in range(6):
+        offset = 4 + k * 2
+        keypoint_x = (
+            raw[..., offset] / model.x_scale * anchors[:, 2] + anchors[:, 0]
+        )  # noqa
+        keypoint_y = (
+            raw[..., offset + 1] / model.y_scale * anchors[:, 3]
+            + anchors[:, 1]  # noqa
+        )
+        boxes[..., offset] = keypoint_x
+        boxes[..., offset + 1] = keypoint_y
+
+    return boxes
+
+
+def predict_on_batch(
+    model: BlazeNet,
+    x,
+    back_model,
+    min_suppression_threshold: int,
+    min_score_thresh: float,
+):
+    """Makes a prediction on a batch of images.
+
+    Arguments:
+        x: a NumPy array of shape (b, H, W, 3) or a PyTorch tensor of
+            shape (b, 3, H, W). The height and width should be 128 pixels.
+
+    Returns:
+        A list containing a tensor of face detections for each image in
+        the batch. If no faces are found for an image, returns a tensor
+        of shape (0, 17).
+
+    Each face detection is a PyTorch tensor consisting of 17 numbers:
+        - ymin, xmin, ymax, xmax
+        - x,y-coordinates for the 6 keypoints
+        - confidence score
+    """
+    if isinstance(x, np.ndarray):
+        x = torch.from_numpy(x).permute((0, 3, 1, 2))
+
+    assert x.shape[1] == 3
+    if back_model:
+        assert x.shape[2] == 256
+        assert x.shape[3] == 256
+    else:
+        assert x.shape[2] == 128
+        assert x.shape[3] == 128
+
+    # 1. Preprocess the images into tensors:
+    x = x.to(model.classifier_8.weight.device)
+    x = _preprocess(x)
+
+    # 2. Run the neural network:
+    with torch.no_grad():
+        out = model(x)
+
+    # 3. Postprocess the raw predictions:
+    detections = _tensors_to_detections(
+        model,
+        out[0],
+        out[1],
+        model.anchors,
+        min_score_thresh,
+    )
+
+    for i in range(len(detections)):
+        faces = _weighted_non_max_suppression(
+            model,
+            detections[i],
+            min_suppression_threshold=min_suppression_threshold,
+        )
+    faces = torch.stack(faces) if len(faces) > 0 else torch.zeros((0, 17))
+    return [faces]
+
+
+def _tensors_to_detections(
+    model: BlazeNet,
+    raw_box_tensor,
+    raw_score_tensor,
+    anchors,
+    min_score_thresh,
+):
+    """The output of the neural network is a tensor of shape (b, 896, 16)
+    containing the bounding box regressor predictions, as well as a tensor
+    of shape (b, 896, 1) with the classification confidences.
+
+    This function converts these two "raw" tensors into proper detections.
+    Returns a list of (num_detections, 17) tensors, one for each image in
+    the batch.
+
+    This is based on the source code from:
+    mediapipe/calculators/tflite/tflite_tensors_to_detections_calculator.cc
+    mediapipe/calculators/tflite/tflite_tensors_to_detections_calculator.proto
+    """
+    assert raw_box_tensor.ndimension() == 3
+    assert raw_box_tensor.shape[1] == model.num_anchors
+    assert raw_box_tensor.shape[2] == model.num_coords
+
+    assert raw_score_tensor.ndimension() == 3
+    assert raw_score_tensor.shape[1] == model.num_anchors
+    assert raw_score_tensor.shape[2] == model.num_classes
+
+    assert raw_box_tensor.shape[0] == raw_score_tensor.shape[0]
+
+    detection_boxes = _decode_boxes(model, raw_box_tensor, anchors)
+
+    thresh = model.score_clipping_thresh
+    raw_score_tensor = raw_score_tensor.clamp(-thresh, thresh)
+    detection_scores = raw_score_tensor.sigmoid().squeeze(dim=-1)
+
+    # Note: we stripped off the last dimension from the scores tensor
+    # because there is only has one class. Now we can simply use a mask
+    # to filter out the boxes with too low confidence.
+    mask = detection_scores >= min_score_thresh
+
+    # Because each image from the batch can have a different number of
+    # detections, process them one at a time using a loop.
+    output_detections = []
+    for i in range(raw_box_tensor.shape[0]):
+        boxes = detection_boxes[i, mask[i]]
+        scores = detection_scores[i, mask[i]].unsqueeze(dim=-1)
+        output_detections.append(torch.cat((boxes, scores), dim=-1))
+
+    return output_detections
+
+
+def predict_on_image(
+    model: BlazeNet,
+    image,
+    back_model,
+    min_suppression_threshold: int,
+    min_score_thresh: float,
+):
+    """Makes a prediction on a single image.
+
+    Arguments:
+        img: a NumPy array of shape (H, W, 3) or a PyTorch tensor of
+                shape (3, H, W). The image's height and width should be
+                128 pixels.
+
+    Returns:
+        A tensor with face detections.
+    """
+    if isinstance(image, np.ndarray):
+        image = torch.from_numpy(image).permute((2, 0, 1))
+
+    return predict_on_batch(
+        model,
+        image.unsqueeze(0),
+        back_model=back_model,
+        min_suppression_threshold=min_suppression_threshold,
+        min_score_thresh=min_score_thresh,
+    )[0]
+
+
+def _preprocess(x):
+    """Converts the image pixels to the range [-1, 1]."""
+    return x.float() / 127.5 - 1.0