add model

app.py

@@ -1,8 +1,65 @@
 import gradio as gr
+from model import SixDRepNet
+import os
+
+import numpy as np
+import torch
+from torchvision import transforms
+import utils
+import time
+
+transformations = transforms.Compose([transforms.Resize(224),
+                                      transforms.CenterCrop(224),
+                                      transforms.ToTensor(),
+                                      transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
+
+model = SixDRepNet(backbone_name='RepVGG-A0',
+                       backbone_file='',
+                       deploy=True,
+                       pretrained=False)
+
+saved_state_dict = torch.load(os.path.join(
+    "weights_ALFW_A0.pth"), map_location='cpu')
+
+if 'model_state_dict' in saved_state_dict:
+    model.load_state_dict(saved_state_dict['model_state_dict'])
+else:
+    model.load_state_dict(saved_state_dict)
+
+# Test the Model
+model.eval()  # Change model to 'eval' mode (BN uses moving mean/var).
+
+th = 15
+
+def predict(img):
+    
+    img = img.convert('RGB')
+    img = transformations(img).unsqueeze(0)
+    with torch.no_grad():
+        start = time.time()
+        R_pred = model(img)
+        end = time.time()
+        timemilis = (end - start)*1000
+
+        euler = utils.compute_euler_angles_from_rotation_matrices(
+            R_pred,use_gpu=False)*180/np.pi
+        p_pred_deg = euler[:, 0].cpu().item()
+        y_pred_deg = euler[:, 1].cpu().item()
+        direction_str = ""
+        if p_pred_deg > th:
+            direction_str = "UP "
+        elif p_pred_deg < th:
+            direction_str ="DOWN "
+        
+        if y_pred_deg > th:
+            direction_str += "LEFT"
+        elif y_pred_deg < th:
+            direction_str += "RIGHT"
+
+        return f"Yaw: {y_pred_deg:0.1f} \n Pitch: {p_pred_deg:0.1f}\n Direction: {direction_str} \n Time: {timemilis:0.2f}ms"
+
 
-def predict(inp):
-    return "Yaw: 30 \n Pitch: 19\n Direction: Left"
 gr.Interface(fn=predict, 
              inputs=gr.Image(type="pil"),
              outputs=gr.Textbox(),
-             examples=["face_left.jpg","face_right.jpg","face_up.jpg","face_down.jpg"]).launch()
+             examples=["face_left.jpg","face_right.jpg","face_up.jpg","face_down.jpg"]).launch(share=True)

backbone/repvgg.py

@@ -0,0 +1,301 @@
+
+import torch.nn as nn
+import numpy as np
+import torch
+import copy
+from backbone.se_block import SEBlock
+
+def conv_bn(in_channels, out_channels, kernel_size, stride, padding, groups=1):
+    result = nn.Sequential()
+    result.add_module('conv', nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
+                                                  kernel_size=kernel_size, stride=stride, padding=padding, groups=groups, bias=False))
+    result.add_module('bn', nn.BatchNorm2d(num_features=out_channels))
+    return result
+
+class RepVGGBlock(nn.Module):
+
+    def __init__(self, in_channels, out_channels, kernel_size,
+                 stride=1, padding=0, dilation=1, groups=1, padding_mode='zeros', deploy=False, use_se=False):
+        super(RepVGGBlock, self).__init__()
+        self.deploy = deploy
+        self.groups = groups
+        self.in_channels = in_channels
+
+        assert kernel_size == 3
+        assert padding == 1
+
+        padding_11 = padding - kernel_size // 2
+
+        self.nonlinearity = nn.ReLU()
+
+        if use_se:
+            self.se = SEBlock(out_channels, internal_neurons=out_channels // 16)
+        else:
+            self.se = nn.Identity()
+
+        if deploy:
+            self.rbr_reparam = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride,
+                                      padding=padding, dilation=dilation, groups=groups, bias=True, padding_mode=padding_mode)
+
+        else:
+            self.rbr_identity = nn.BatchNorm2d(num_features=in_channels) if out_channels == in_channels and stride == 1 else None
+            self.rbr_dense = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, groups=groups)
+            self.rbr_1x1 = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=padding_11, groups=groups)
+            print('RepVGG Block, identity = ', self.rbr_identity)
+
+
+    def forward(self, inputs):
+        if hasattr(self, 'rbr_reparam'):
+            return self.nonlinearity(self.se(self.rbr_reparam(inputs)))
+
+        if self.rbr_identity is None:
+            id_out = 0
+        else:
+            id_out = self.rbr_identity(inputs)
+
+        return self.nonlinearity(self.se(self.rbr_dense(inputs) + self.rbr_1x1(inputs) + id_out))
+
+
+    #   Optional. This improves the accuracy and facilitates quantization.
+    #   1.  Cancel the original weight decay on rbr_dense.conv.weight and rbr_1x1.conv.weight.
+    #   2.  Use like this.
+    #       loss = criterion(....)
+    #       for every RepVGGBlock blk:
+    #           loss += weight_decay_coefficient * 0.5 * blk.get_cust_L2()
+    #       optimizer.zero_grad()
+    #       loss.backward()
+    def get_custom_L2(self):
+        K3 = self.rbr_dense.conv.weight
+        K1 = self.rbr_1x1.conv.weight
+        t3 = (self.rbr_dense.bn.weight / ((self.rbr_dense.bn.running_var + self.rbr_dense.bn.eps).sqrt())).reshape(-1, 1, 1, 1).detach()
+        t1 = (self.rbr_1x1.bn.weight / ((self.rbr_1x1.bn.running_var + self.rbr_1x1.bn.eps).sqrt())).reshape(-1, 1, 1, 1).detach()
+
+        l2_loss_circle = (K3 ** 2).sum() - (K3[:, :, 1:2, 1:2] ** 2).sum()      # The L2 loss of the "circle" of weights in 3x3 kernel. Use regular L2 on them.
+        eq_kernel = K3[:, :, 1:2, 1:2] * t3 + K1 * t1                           # The equivalent resultant central point of 3x3 kernel.
+        l2_loss_eq_kernel = (eq_kernel ** 2 / (t3 ** 2 + t1 ** 2)).sum()        # Normalize for an L2 coefficient comparable to regular L2.
+        return l2_loss_eq_kernel + l2_loss_circle
+
+
+
+#   This func derives the equivalent kernel and bias in a DIFFERENTIABLE way.
+#   You can get the equivalent kernel and bias at any time and do whatever you want,
+    #   for example, apply some penalties or constraints during training, just like you do to the other models.
+#   May be useful for quantization or pruning.
+    def get_equivalent_kernel_bias(self):
+        kernel3x3, bias3x3 = self._fuse_bn_tensor(self.rbr_dense)
+        kernel1x1, bias1x1 = self._fuse_bn_tensor(self.rbr_1x1)
+        kernelid, biasid = self._fuse_bn_tensor(self.rbr_identity)
+        return kernel3x3 + self._pad_1x1_to_3x3_tensor(kernel1x1) + kernelid, bias3x3 + bias1x1 + biasid
+
+    def _pad_1x1_to_3x3_tensor(self, kernel1x1):
+        if kernel1x1 is None:
+            return 0
+        else:
+            return torch.nn.functional.pad(kernel1x1, [1,1,1,1])
+
+    def _fuse_bn_tensor(self, branch):
+        if branch is None:
+            return 0, 0
+        if isinstance(branch, nn.Sequential):
+            kernel = branch.conv.weight
+            running_mean = branch.bn.running_mean
+            running_var = branch.bn.running_var
+            gamma = branch.bn.weight
+            beta = branch.bn.bias
+            eps = branch.bn.eps
+        else:
+            assert isinstance(branch, nn.BatchNorm2d)
+            if not hasattr(self, 'id_tensor'):
+                input_dim = self.in_channels // self.groups
+                kernel_value = np.zeros((self.in_channels, input_dim, 3, 3), dtype=np.float32)
+                for i in range(self.in_channels):
+                    kernel_value[i, i % input_dim, 1, 1] = 1
+                self.id_tensor = torch.from_numpy(kernel_value).to(branch.weight.device)
+            kernel = self.id_tensor
+            running_mean = branch.running_mean
+            running_var = branch.running_var
+            gamma = branch.weight
+            beta = branch.bias
+            eps = branch.eps
+        std = (running_var + eps).sqrt()
+        t = (gamma / std).reshape(-1, 1, 1, 1)
+        return kernel * t, beta - running_mean * gamma / std
+
+    def switch_to_deploy(self):
+        if hasattr(self, 'rbr_reparam'):
+            return
+        kernel, bias = self.get_equivalent_kernel_bias()
+        self.rbr_reparam = nn.Conv2d(in_channels=self.rbr_dense.conv.in_channels, out_channels=self.rbr_dense.conv.out_channels,
+                                     kernel_size=self.rbr_dense.conv.kernel_size, stride=self.rbr_dense.conv.stride,
+                                     padding=self.rbr_dense.conv.padding, dilation=self.rbr_dense.conv.dilation, groups=self.rbr_dense.conv.groups, bias=True)
+        self.rbr_reparam.weight.data = kernel
+        self.rbr_reparam.bias.data = bias
+        for para in self.parameters():
+            para.detach_()
+        self.__delattr__('rbr_dense')
+        self.__delattr__('rbr_1x1')
+        if hasattr(self, 'rbr_identity'):
+            self.__delattr__('rbr_identity')
+        if hasattr(self, 'id_tensor'):
+            self.__delattr__('id_tensor')
+        self.deploy = True
+
+
+
+class RepVGG(nn.Module):
+
+    def __init__(self, num_blocks, num_classes=1000, width_multiplier=None, override_groups_map=None, deploy=False, use_se=False):
+        super(RepVGG, self).__init__()
+
+        assert len(width_multiplier) == 4
+
+        self.deploy = deploy
+        self.override_groups_map = override_groups_map or dict()
+        self.use_se = use_se
+
+        assert 0 not in self.override_groups_map
+
+        self.in_planes = min(64, int(64 * width_multiplier[0]))
+
+        self.stage0 = RepVGGBlock(in_channels=3, out_channels=self.in_planes, kernel_size=3, stride=2, padding=1, deploy=self.deploy, use_se=self.use_se)
+        self.cur_layer_idx = 1
+        self.stage1 = self._make_stage(int(64 * width_multiplier[0]), num_blocks[0], stride=2)
+        self.stage2 = self._make_stage(int(128 * width_multiplier[1]), num_blocks[1], stride=2)
+        self.stage3 = self._make_stage(int(256 * width_multiplier[2]), num_blocks[2], stride=2)
+        self.stage4 = self._make_stage(int(512 * width_multiplier[3]), num_blocks[3], stride=2)
+        self.gap = nn.AdaptiveAvgPool2d(output_size=1)
+        self.linear = nn.Linear(int(512 * width_multiplier[3]), num_classes)
+
+
+    def _make_stage(self, planes, num_blocks, stride):
+        strides = [stride] + [1]*(num_blocks-1)
+        blocks = []
+        for stride in strides:
+            cur_groups = self.override_groups_map.get(self.cur_layer_idx, 1)
+            blocks.append(RepVGGBlock(in_channels=self.in_planes, out_channels=planes, kernel_size=3,
+                                      stride=stride, padding=1, groups=cur_groups, deploy=self.deploy, use_se=self.use_se))
+            self.in_planes = planes
+            self.cur_layer_idx += 1
+        return nn.Sequential(*blocks)
+
+    def forward(self, x):
+        out = self.stage0(x)
+        out = self.stage1(out)
+        out = self.stage2(out)
+        out = self.stage3(out)
+        out = self.stage4(out)
+        out = self.gap(out)
+        out = out.view(out.size(0), -1)
+        out = self.linear(out)
+        return out
+
+
+optional_groupwise_layers = [2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26]
+g2_map = {l: 2 for l in optional_groupwise_layers}
+g4_map = {l: 4 for l in optional_groupwise_layers}
+
+def create_RepVGG_A0(deploy=False):
+    return RepVGG(num_blocks=[2, 4, 14, 1], num_classes=1000,
+                  width_multiplier=[0.75, 0.75, 0.75, 2.5], override_groups_map=None, deploy=deploy)
+
+def create_RepVGG_A1(deploy=False):
+    return RepVGG(num_blocks=[2, 4, 14, 1], num_classes=1000,
+                  width_multiplier=[1, 1, 1, 2.5], override_groups_map=None, deploy=deploy)
+
+def create_RepVGG_A2(deploy=False):
+    return RepVGG(num_blocks=[2, 4, 14, 1], num_classes=1000,
+                  width_multiplier=[1.5, 1.5, 1.5, 2.75], override_groups_map=None, deploy=deploy)
+
+def create_RepVGG_B0(deploy=False):
+    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
+                  width_multiplier=[1, 1, 1, 2.5], override_groups_map=None, deploy=deploy)
+
+def create_RepVGG_B1(deploy=False):
+    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
+                  width_multiplier=[2, 2, 2, 4], override_groups_map=None, deploy=deploy)
+
+def create_RepVGG_B1g2(deploy=False):
+    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
+                  width_multiplier=[2, 2, 2, 4], override_groups_map=g2_map, deploy=deploy)
+
+def create_RepVGG_B1g4(deploy=False):
+    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
+                  width_multiplier=[2, 2, 2, 4], override_groups_map=g4_map, deploy=deploy)
+
+
+def create_RepVGG_B2(deploy=False):
+    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
+                  width_multiplier=[2.5, 2.5, 2.5, 5], override_groups_map=None, deploy=deploy)
+
+def create_RepVGG_B2g2(deploy=False):
+    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
+                  width_multiplier=[2.5, 2.5, 2.5, 5], override_groups_map=g2_map, deploy=deploy)
+
+def create_RepVGG_B2g4(deploy=False):
+    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
+                  width_multiplier=[2.5, 2.5, 2.5, 5], override_groups_map=g4_map, deploy=deploy)
+
+
+def create_RepVGG_B3(deploy=False):
+    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
+                  width_multiplier=[3, 3, 3, 5], override_groups_map=None, deploy=deploy)
+
+def create_RepVGG_B3g2(deploy=False):
+    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
+                  width_multiplier=[3, 3, 3, 5], override_groups_map=g2_map, deploy=deploy)
+
+def create_RepVGG_B3g4(deploy=False):
+    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
+                  width_multiplier=[3, 3, 3, 5], override_groups_map=g4_map, deploy=deploy)
+
+def create_RepVGG_D2se(deploy=False):
+    return RepVGG(num_blocks=[8, 14, 24, 1], num_classes=1000,
+                  width_multiplier=[2.5, 2.5, 2.5, 5], override_groups_map=None, deploy=deploy, use_se=True)
+
+
+func_dict = {
+'RepVGG-A0': create_RepVGG_A0,
+'RepVGG-A1': create_RepVGG_A1,
+'RepVGG-A2': create_RepVGG_A2,
+'RepVGG-B0': create_RepVGG_B0,
+'RepVGG-B1': create_RepVGG_B1,
+'RepVGG-B1g2': create_RepVGG_B1g2,
+'RepVGG-B1g4': create_RepVGG_B1g4,
+'RepVGG-B2': create_RepVGG_B2,
+'RepVGG-B2g2': create_RepVGG_B2g2,
+'RepVGG-B2g4': create_RepVGG_B2g4,
+'RepVGG-B3': create_RepVGG_B3,
+'RepVGG-B3g2': create_RepVGG_B3g2,
+'RepVGG-B3g4': create_RepVGG_B3g4,
+'RepVGG-D2se': create_RepVGG_D2se,      #   Updated at April 25, 2021. This is not reported in the CVPR paper.
+}
+def get_RepVGG_func_by_name(name):
+    return func_dict[name]
+
+
+
+#   Use this for converting a RepVGG model or a bigger model with RepVGG as its component
+#   Use like this
+#   model = create_RepVGG_A0(deploy=False)
+#   train model or load weights
+#   repvgg_model_convert(model, save_path='repvgg_deploy.pth')
+#   If you want to preserve the original model, call with do_copy=True
+
+#   ====================== for using RepVGG as the backbone of a bigger model, e.g., PSPNet, the pseudo code will be like
+#   train_backbone = create_RepVGG_B2(deploy=False)
+#   train_backbone.load_state_dict(torch.load('RepVGG-B2-train.pth'))
+#   train_pspnet = build_pspnet(backbone=train_backbone)
+#   segmentation_train(train_pspnet)
+#   deploy_pspnet = repvgg_model_convert(train_pspnet)
+#   segmentation_test(deploy_pspnet)
+#   =====================   example_pspnet.py shows an example
+
+def repvgg_model_convert(model:torch.nn.Module, save_path=None, do_copy=True):
+    if do_copy:
+        model = copy.deepcopy(model)
+    for module in model.modules():
+        if hasattr(module, 'switch_to_deploy'):
+            module.switch_to_deploy()
+    if save_path is not None:
+        torch.save(model.state_dict(), save_path)
+    return model

backbone/se_block.py

@@ -0,0 +1,22 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+#   https://openaccess.thecvf.com/content_cvpr_2018/html/Hu_Squeeze-and-Excitation_Networks_CVPR_2018_paper.html
+
+class SEBlock(nn.Module):
+
+    def __init__(self, input_channels, internal_neurons):
+        super(SEBlock, self).__init__()
+        self.down = nn.Conv2d(in_channels=input_channels, out_channels=internal_neurons, kernel_size=1, stride=1, bias=True)
+        self.up = nn.Conv2d(in_channels=internal_neurons, out_channels=input_channels, kernel_size=1, stride=1, bias=True)
+        self.input_channels = input_channels
+
+    def forward(self, inputs):
+        x = F.avg_pool2d(inputs, kernel_size=inputs.size(3))
+        x = self.down(x)
+        x = F.relu(x)
+        x = self.up(x)
+        x = torch.sigmoid(x)
+        x = x.view(-1, self.input_channels, 1, 1)
+        return inputs * x

demo.py

@@ -0,0 +1,141 @@
+from face_detection import RetinaFace
+from model import SixDRepNet
+import math
+import re
+from matplotlib import pyplot as plt
+import sys
+import os
+import argparse
+
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+from numpy.lib.function_base import _quantile_unchecked
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torchvision import transforms
+import torch.backends.cudnn as cudnn
+import torchvision
+import torch.nn.functional as F
+import utils
+import matplotlib
+from PIL import Image
+import time
+matplotlib.use('TkAgg')
+
+
+def parse_args():
+    """Parse input arguments."""
+    parser = argparse.ArgumentParser(
+        description='Head pose estimation using the 6DRepNet.')
+    parser.add_argument('--gpu',
+                        dest='gpu_id', help='GPU device id to use [0]',
+                        default=0, type=int)
+    parser.add_argument('--cam',
+                        dest='cam_id', help='Camera device id to use [0]',
+                        default=0, type=int)
+    parser.add_argument('--snapshot',
+                        dest='snapshot', help='Name of model snapshot.',
+                        default='', type=str)
+    parser.add_argument('--save_viz',
+                        dest='save_viz', help='Save images with pose cube.',
+                        default=False, type=bool)
+
+    args = parser.parse_args()
+    return args
+
+
+transformations = transforms.Compose([transforms.Resize(224),
+                                      transforms.CenterCrop(224),
+                                      transforms.ToTensor(),
+                                      transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
+
+if __name__ == '__main__':
+    args = parse_args()
+    cudnn.enabled = True
+    gpu = args.gpu_id
+    cam = args.cam_id
+    snapshot_path = args.snapshot
+    model = SixDRepNet(backbone_name='RepVGG-A0',
+                       backbone_file='',
+                       deploy=True,
+                       pretrained=False)
+
+    print('Loading data.')
+
+    detector = RetinaFace(gpu_id=gpu)
+
+    # Load snapshot
+    saved_state_dict = torch.load(os.path.join(
+        snapshot_path), map_location='cpu')
+
+    if 'model_state_dict' in saved_state_dict:
+        model.load_state_dict(saved_state_dict['model_state_dict'])
+    else:
+        model.load_state_dict(saved_state_dict)
+    if gpu != -1:
+        model.cuda(gpu)
+
+    # Test the Model
+    model.eval()  # Change model to 'eval' mode (BN uses moving mean/var).
+
+    cap = cv2.VideoCapture(cam)
+
+    # Check if the webcam is opened correctly
+    if not cap.isOpened():
+        raise IOError("Cannot open webcam")
+
+    with torch.no_grad():
+        while True:
+            ret, frame = cap.read()
+
+            faces = detector(frame)
+
+            for box, landmarks, score in faces:
+
+                # Print the location of each face in this image
+                if score < .95:
+                    continue
+                x_min = int(box[0])
+                y_min = int(box[1])
+                x_max = int(box[2])
+                y_max = int(box[3])
+                bbox_width = abs(x_max - x_min)
+                bbox_height = abs(y_max - y_min)
+
+                x_min = max(0, x_min-int(0.2*bbox_height))
+                y_min = max(0, y_min-int(0.2*bbox_width))
+                x_max = x_max+int(0.2*bbox_height)
+                y_max = y_max+int(0.2*bbox_width)
+
+                img = frame[y_min:y_max, x_min:x_max]
+                img = Image.fromarray(img)
+                img = img.convert('RGB')
+                img = transformations(img)
+
+                img = torch.Tensor(img[None, :])
+                if gpu != -1:
+                    img = img.cuda(gpu)
+
+                c = cv2.waitKey(1)
+                if c == 27:
+                    break
+
+                start = time.time()
+                R_pred = model(img)
+                end = time.time()
+                print('Head pose estimation: %2f ms' % ((end - start)*1000.))
+
+                euler = utils.compute_euler_angles_from_rotation_matrices(
+                    R_pred,use_gpu=False)*180/np.pi
+                p_pred_deg = euler[:, 0].cpu()
+                y_pred_deg = euler[:, 1].cpu()
+                r_pred_deg = euler[:, 2].cpu()
+
+                #utils.draw_axis(frame, y_pred_deg, p_pred_deg, r_pred_deg, left+int(.5*(right-left)), top, size=100)
+                utils.plot_pose_cube(frame,  y_pred_deg, p_pred_deg, r_pred_deg, x_min + int(.5*(
+                    x_max-x_min)), y_min + int(.5*(y_max-y_min)), size=bbox_width)
+            cv2.imshow("Demo", np.array(frame, dtype = np.uint8))
+            cv2.waitKey(5)

flagged/log.csv

@@ -0,0 +1,2 @@
+'img','output','flag','username','timestamp'
+'','','','','2022-08-17 14:04:57.627952'

model.py

@@ -0,0 +1,117 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+import math
+from backbone.repvgg import get_RepVGG_func_by_name
+import utils
+
+class SixDRepNet(nn.Module):
+    def __init__(self,
+                 backbone_name, backbone_file, deploy,
+                 bins=(1, 2, 3, 6),
+                 droBatchNorm=nn.BatchNorm2d,
+                 pretrained=True):
+        super(SixDRepNet, self).__init__()
+
+        repvgg_fn = get_RepVGG_func_by_name(backbone_name)
+        backbone = repvgg_fn(deploy)
+        if pretrained:
+            checkpoint = torch.load(backbone_file)
+            if 'state_dict' in checkpoint:
+                checkpoint = checkpoint['state_dict']
+            ckpt = {k.replace('module.', ''): v for k,
+                    v in checkpoint.items()}  # strip the names
+            backbone.load_state_dict(ckpt)
+
+        self.layer0, self.layer1, self.layer2, self.layer3, self.layer4 = backbone.stage0, backbone.stage1, backbone.stage2, backbone.stage3, backbone.stage4
+        self.gap = nn.AdaptiveAvgPool2d(output_size=1)
+
+        last_channel = 0
+        for n, m in self.layer4.named_modules():
+            if ('rbr_dense' in n or 'rbr_reparam' in n) and isinstance(m, nn.Conv2d):
+                last_channel = m.out_channels
+
+        fea_dim = last_channel
+
+        self.linear_reg = nn.Linear(fea_dim, 6)
+
+    def forward(self, x):
+
+        x = self.layer0(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x= self.gap(x)
+        x = torch.flatten(x, 1)
+        x = self.linear_reg(x)
+        return utils.compute_rotation_matrix_from_ortho6d(x,use_gpu=False)
+
+
+
+
+class SixDRepNet2(nn.Module):
+    def __init__(self, block, layers, fc_layers=1):
+        self.inplanes = 64
+        super(SixDRepNet2, self).__init__()
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.avgpool = nn.AvgPool2d(7)
+
+        self.linear_reg = nn.Linear(512*block.expansion,6)
+      
+
+
+        # Vestigial layer from previous experiments
+        self.fc_finetune = nn.Linear(512 * block.expansion + 3, 3)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+
+        x = self.linear_reg(x)        
+        out = utils.compute_rotation_matrix_from_ortho6d(x)
+
+        return out

utils.py

@@ -0,0 +1,208 @@
+import numpy as np
+import torch
+#from torch.utils.serialization import load_lua
+import os
+import scipy.io as sio
+import cv2
+import math
+from math import cos, sin
+
+def plot_pose_cube(img, yaw, pitch, roll, tdx=None, tdy=None, size=150.):
+    # Input is a cv2 image
+    # pose_params: (pitch, yaw, roll, tdx, tdy)
+    # Where (tdx, tdy) is the translation of the face.
+    # For pose we have [pitch yaw roll tdx tdy tdz scale_factor]
+
+    p = pitch * np.pi / 180
+    y = -(yaw * np.pi / 180)
+    r = roll * np.pi / 180
+    if tdx != None and tdy != None:
+        face_x = tdx - 0.50 * size 
+        face_y = tdy - 0.50 * size
+
+    else:
+        height, width = img.shape[:2]
+        face_x = width / 2 - 0.5 * size
+        face_y = height / 2 - 0.5 * size
+
+    x1 = size * (cos(y) * cos(r)) + face_x
+    y1 = size * (cos(p) * sin(r) + cos(r) * sin(p) * sin(y)) + face_y 
+    x2 = size * (-cos(y) * sin(r)) + face_x
+    y2 = size * (cos(p) * cos(r) - sin(p) * sin(y) * sin(r)) + face_y
+    x3 = size * (sin(y)) + face_x
+    y3 = size * (-cos(y) * sin(p)) + face_y
+
+
+    # Draw base in red
+    cv2.line(img, (int(face_x), int(face_y)), (int(x1),int(y1)),(0,0,255),3)
+    cv2.line(img, (int(face_x), int(face_y)), (int(x2),int(y2)),(0,0,255),3)
+    cv2.line(img, (int(x2), int(y2)), (int(x2+x1-face_x),int(y2+y1-face_y)),(0,0,255),3)
+    cv2.line(img, (int(x1), int(y1)), (int(x1+x2-face_x),int(y1+y2-face_y)),(0,0,255),3)
+    # Draw pillars in blue
+    cv2.line(img, (int(face_x), int(face_y)), (int(x3),int(y3)),(255,0,0),2)
+    cv2.line(img, (int(x1), int(y1)), (int(x1+x3-face_x),int(y1+y3-face_y)),(255,0,0),2)
+    cv2.line(img, (int(x2), int(y2)), (int(x2+x3-face_x),int(y2+y3-face_y)),(255,0,0),2)
+    cv2.line(img, (int(x2+x1-face_x),int(y2+y1-face_y)), (int(x3+x1+x2-2*face_x),int(y3+y2+y1-2*face_y)),(255,0,0),2)
+    # Draw top in green
+    cv2.line(img, (int(x3+x1-face_x),int(y3+y1-face_y)), (int(x3+x1+x2-2*face_x),int(y3+y2+y1-2*face_y)),(0,255,0),2)
+    cv2.line(img, (int(x2+x3-face_x),int(y2+y3-face_y)), (int(x3+x1+x2-2*face_x),int(y3+y2+y1-2*face_y)),(0,255,0),2)
+    cv2.line(img, (int(x3), int(y3)), (int(x3+x1-face_x),int(y3+y1-face_y)),(0,255,0),2)
+    cv2.line(img, (int(x3), int(y3)), (int(x3+x2-face_x),int(y3+y2-face_y)),(0,255,0),2)
+
+    return img
+
+def draw_axis(img, yaw, pitch, roll, tdx=None, tdy=None, size = 100):
+
+    pitch = pitch * np.pi / 180
+    yaw = -(yaw * np.pi / 180)
+    roll = roll * np.pi / 180
+
+    if tdx != None and tdy != None:
+        tdx = tdx
+        tdy = tdy
+    else:
+        height, width = img.shape[:2]
+        tdx = width / 2
+        tdy = height / 2
+
+    # X-Axis pointing to right. drawn in red
+    x1 = size * (cos(yaw) * cos(roll)) + tdx
+    y1 = size * (cos(pitch) * sin(roll) + cos(roll) * sin(pitch) * sin(yaw)) + tdy
+
+    # Y-Axis | drawn in green
+    #        v
+    x2 = size * (-cos(yaw) * sin(roll)) + tdx
+    y2 = size * (cos(pitch) * cos(roll) - sin(pitch) * sin(yaw) * sin(roll)) + tdy
+
+    # Z-Axis (out of the screen) drawn in blue
+    x3 = size * (sin(yaw)) + tdx
+    y3 = size * (-cos(yaw) * sin(pitch)) + tdy
+
+    cv2.line(img, (int(tdx), int(tdy)), (int(x1),int(y1)),(0,0,255),4)
+    cv2.line(img, (int(tdx), int(tdy)), (int(x2),int(y2)),(0,255,0),4)
+    cv2.line(img, (int(tdx), int(tdy)), (int(x3),int(y3)),(255,0,0),4)
+
+    return img
+
+
+def get_pose_params_from_mat(mat_path):
+    # This functions gets the pose parameters from the .mat
+    # Annotations that come with the Pose_300W_LP dataset.
+    mat = sio.loadmat(mat_path)
+    # [pitch yaw roll tdx tdy tdz scale_factor]
+    pre_pose_params = mat['Pose_Para'][0]
+    # Get [pitch, yaw, roll, tdx, tdy]
+    pose_params = pre_pose_params[:5]
+    return pose_params
+
+def get_ypr_from_mat(mat_path):
+    # Get yaw, pitch, roll from .mat annotation.
+    # They are in radians
+    mat = sio.loadmat(mat_path)
+    # [pitch yaw roll tdx tdy tdz scale_factor]
+    pre_pose_params = mat['Pose_Para'][0]
+    # Get [pitch, yaw, roll]
+    pose_params = pre_pose_params[:3]
+    return pose_params
+
+def get_pt2d_from_mat(mat_path):
+    # Get 2D landmarks
+    mat = sio.loadmat(mat_path)
+    pt2d = mat['pt2d']
+    return pt2d
+
+# batch*n
+def normalize_vector( v, use_gpu=True):
+    batch=v.shape[0]
+    v_mag = torch.sqrt(v.pow(2).sum(1))# batch
+    if use_gpu:
+        v_mag = torch.max(v_mag, torch.autograd.Variable(torch.FloatTensor([1e-8]).cuda()))
+    else:
+        v_mag = torch.max(v_mag, torch.autograd.Variable(torch.FloatTensor([1e-8])))  
+    v_mag = v_mag.view(batch,1).expand(batch,v.shape[1])
+    v = v/v_mag
+    return v
+    
+# u, v batch*n
+def cross_product( u, v):
+    batch = u.shape[0]
+    #print (u.shape)
+    #print (v.shape)
+    i = u[:,1]*v[:,2] - u[:,2]*v[:,1]
+    j = u[:,2]*v[:,0] - u[:,0]*v[:,2]
+    k = u[:,0]*v[:,1] - u[:,1]*v[:,0]
+        
+    out = torch.cat((i.view(batch,1), j.view(batch,1), k.view(batch,1)),1)#batch*3
+        
+    return out
+        
+    
+#poses batch*6
+#poses
+def compute_rotation_matrix_from_ortho6d(poses, use_gpu=True):
+    x_raw = poses[:,0:3]#batch*3
+    y_raw = poses[:,3:6]#batch*3
+        
+    x = normalize_vector(x_raw, use_gpu) #batch*3
+    z = cross_product(x,y_raw) #batch*3
+    z = normalize_vector(z, use_gpu)#batch*3
+    y = cross_product(z,x)#batch*3
+        
+    x = x.view(-1,3,1)
+    y = y.view(-1,3,1)
+    z = z.view(-1,3,1)
+    matrix = torch.cat((x,y,z), 2) #batch*3*3
+    return matrix
+
+
+#input batch*4*4 or batch*3*3
+#output torch batch*3 x, y, z in radiant
+#the rotation is in the sequence of x,y,z
+def compute_euler_angles_from_rotation_matrices(rotation_matrices, use_gpu=True):
+    batch=rotation_matrices.shape[0]
+    R=rotation_matrices
+    sy = torch.sqrt(R[:,0,0]*R[:,0,0]+R[:,1,0]*R[:,1,0])
+    singular= sy<1e-6
+    singular=singular.float()
+        
+    x=torch.atan2(R[:,2,1], R[:,2,2])
+    y=torch.atan2(-R[:,2,0], sy)
+    z=torch.atan2(R[:,1,0],R[:,0,0])
+    
+    xs=torch.atan2(-R[:,1,2], R[:,1,1])
+    ys=torch.atan2(-R[:,2,0], sy)
+    zs=R[:,1,0]*0
+        
+    if use_gpu:
+        out_euler=torch.autograd.Variable(torch.zeros(batch,3).cuda())
+    else:
+        out_euler=torch.autograd.Variable(torch.zeros(batch,3))  
+    out_euler[:,0]=x*(1-singular)+xs*singular
+    out_euler[:,1]=y*(1-singular)+ys*singular
+    out_euler[:,2]=z*(1-singular)+zs*singular
+        
+    return out_euler
+
+
+def get_R(x,y,z):
+    ''' Get rotation matrix from three rotation angles (radians). right-handed.
+    Args:
+        angles: [3,]. x, y, z angles
+    Returns:
+        R: [3, 3]. rotation matrix.
+    '''
+    # x
+    Rx = np.array([[1, 0, 0],
+                   [0, np.cos(x), -np.sin(x)],
+                   [0, np.sin(x), np.cos(x)]])
+    # y
+    Ry = np.array([[np.cos(y), 0, np.sin(y)],
+                   [0, 1, 0],
+                   [-np.sin(y), 0, np.cos(y)]])
+    # z
+    Rz = np.array([[np.cos(z), -np.sin(z), 0],
+                   [np.sin(z), np.cos(z), 0],
+                   [0, 0, 1]])
+
+    R = Rz.dot(Ry.dot(Rx))
+    return R

weights_ALFW_A0.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:46333a6f16f15a95621fbc0db8b00705c7c776172158ef7e671a55f8710ac894
+size 28161179