change_hair/project/hair_service_sd/momocv/BigResNetStable.py

import torch
import torch.nn as nn
import math
import os
import cv2
import numpy as np
# from utils.DATAIMG import DATAIMG
from utils import landmark_processor
import glob

def op_name(op_name, m):
    m.op_name = op_name
    return m

class Flatten(nn.Module):
    def __init__(self, axis=1):
        super(Flatten, self).__init__()
        self.axis = axis

    def forward(self, x):
        assert self.axis == 1
        x = x.reshape(x.shape[0], -1)
        return x

def flatten(name, axis=1):
    return op_name(name, Flatten(axis))

def conv_relu(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1):
    return nn.Sequential(
        op_name(name, nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, True)),
        op_name(name + '/relu', nn.ReLU()),
    )

class BasicBlock(nn.Module):
    def __init__(self, name, inplanes, planes, stride=2):
        super(BasicBlock, self).__init__()
        self.op_name = name
        self.conv1 = conv_relu(name + '/conv1', inplanes, planes, kernel_size=3, stride=stride, padding=1)
        self.conv2 = conv_relu(name + '/conv2', planes, planes, kernel_size=3, stride=1, padding=1)
        self.downsample = conv_relu(name + '/sc_conv', inplanes, planes, kernel_size=1, stride=stride)

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.conv2(out)
        if self.downsample is not None:
            residual = self.downsample(x)
        ret = out + residual
        return ret


class BigResNetStable(nn.Module):
    def __init__(self, name, in_channels, out_channels):
        super(BigResNetStable, self).__init__()
        self.op_name = name

        op_list = []

        op_list += [conv_relu(name + '/first_conv', in_channels, 32, kernel_size=5, stride=2, padding=2)]

        ch_num = [32, 48, 64, 96, 128]

        op_list += [BasicBlock(name + '/stage%d' % (i + 1), ch_num[i], ch_num[i + 1]) for i in range(len(ch_num) - 1)]

        op_list += [flatten(name + '/flatten', 1)]

        op_list1 = [op_name(name + '/FC1/FC', nn.Linear(2048, 512)),
                    op_name(name + '/FC1/relu', nn.ReLU())]

        fullyconnected1 = [op_name(name + '/FC2', nn.Linear(512, out_channels))]
        poselayer = [op_name('poselayer', nn.Linear(512, 3))]
        tracking_probe = [op_name('tracking_probe', nn.Linear(2048, 1))]
        occlusion_probe = [op_name('occlusion_probe', nn.Linear(2048, 87))]

        self.features = nn.Sequential(*op_list)
        self.fc1 = nn.Sequential(*op_list1)
        self.fc2 = nn.Sequential(*fullyconnected1)
        self.poselayer = nn.Sequential(*poselayer)
        self.tracking_probe = nn.Sequential(*tracking_probe)
        self.occlusion_probe = nn.Sequential(*occlusion_probe)

    def forward(self, x):
        features = self.features(x)
        fc1 = self.fc1(features)
        fullyconnected1 = self.fc2(fc1)
        poselayer = self.poselayer(fc1)
        tracking_probe = self.tracking_probe(features)
        occlusion_probe = self.occlusion_probe(features)
        return fullyconnected1.cpu().numpy(), poselayer.cpu().numpy(), tracking_probe.cpu().numpy(), occlusion_probe.cpu().numpy()

class MomocvFaceAlignment(object):
    def __init__(self, gpu_id=None):
        self.gpu_id = gpu_id
        self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')

        self.face_alignment_net = BigResNetStable('BigResNetStable', 3, 174)

        self.model_path, _ = os.path.split(os.path.realpath(__file__))

        weights = torch.load(os.path.join(self.model_path, 'BigResNetStable.pth'),
                             map_location=lambda storage, loc: storage)
        self.face_alignment_net.load_state_dict(weights)
        self.face_alignment_net.to(self.device)
        self.face_alignment_net.eval()

        self.trackingFaceRects = []


    def crop_img(self, image, labels=None,img_size = 128):
        '''
        :param imageimgs: Original graph
        :param labels: The boxes of the original picture ; labels is a ndarray: [list,list]
        :return: Coordinates and categories relative to the original
        '''
        # Preprocessing
        labels = np.array([labels])
        boxes_ret = np.zeros(labels.shape)
        crop_imgs = []
        ret_M = []
        for i in range(labels.shape[0]):
            box_orig = np.array(labels[i, :])
            box = box_orig.astype(np.int32).copy()
            center = np.array([(box[0] + box[2]) / 2, (box[1] + box[3]) / 2]).astype(np.int32)
            max_lenth = int(max(box[3] - box[1], box[2] - box[0]) / 2 * 1.0)
            up = int(center[1] - max_lenth)
            down = int(center[1] + max_lenth)
            left = int(center[0] - max_lenth)
            right = int(center[0] + max_lenth)
            if up < 0:
                up = 0
                down = max_lenth * 2
            if down > image.shape[0]:
                down = image.shape[0]
            if left < 0:
                left = 0
                right = max_lenth * 2
            if right > image.shape[1]:
                right = image.shape[1]

            crop_img = image[up:down, left:right, :].copy()
            crop_img = cv2.resize(crop_img, (img_size, img_size))
            box_orig[0] -= left
            box_orig[2] -= left
            box_orig[1] -= up
            box_orig[3] -= up
            box_orig[0] *= img_size / (right - left)
            box_orig[2] *= img_size / (right - left)
            box_orig[1] *= img_size / (down - up)
            box_orig[3] *= img_size / (down - up)
            crop_imgs.append(crop_img)
            ret_M.append(np.array([img_size / (right - left), img_size / (down - up), left, up]))
            boxes_ret[i, :] = box_orig
        return crop_imgs, boxes_ret, ret_M


    def detect_from_bbox(self, img, bboxs):
        dst_size = 128
        landmarks_res = []
        with torch.no_grad():
            input_numpy = np.zeros((len(bboxs), 3, dst_size, dst_size), dtype=np.float32)
            for ix, bbox in enumerate(bboxs):
                crop_imgs, boxes_ret, ret_M = self.crop_img(img, bbox, dst_size)
                input_numpy[ix, :, :, :] = crop_imgs[0].transpose((2, 0, 1)).astype(np.float32)

                # cv2.imshow('tmp', tmp)
                # cv2.waitKey()

            in_tensor = torch.from_numpy(input_numpy)
            in_tensor = in_tensor.to(self.device)
            fullyconnected1, poselayer, tracking_probe, occlusion_probe = self.face_alignment_net(in_tensor)
            for ix, pts in enumerate(fullyconnected1):
                orig_pts = ((np.reshape(pts, (2, 87)).transpose((1, 0))) * dst_size)
                # print('pts',orig_pts)
                # landmark2 = landmark2.cpu().detach().numpy()[0].reshape(2, -1).transpose((1, 0)).reshape(-1)
                # orig_pts = ((orig_pts + 0.5) * 128)
                orig_pts[:, 0] = orig_pts[:, 0] / ret_M[0][0] + ret_M[0][2]
                orig_pts[:, 1] = orig_pts[:, 1] / ret_M[0][1] + ret_M[0][3]

                # orig_pts[0] = orig_pts[0] / ret_M[0][0] + ret_M[0][2]
                # orig_pts[2] = orig_pts[2] / ret_M[0][0] + ret_M[0][2]
                # orig_pts[1] = orig_pts[1] / ret_M[0][1] + ret_M[0][3]
                # orig_pts[3] = orig_pts[3] / ret_M[0][1] + ret_M[0][3]
                # cur_landmark2.append(orig_pts)
                # orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True)
                landmarks_res.append(orig_pts)
            tracking_probe = 1 / (1 + np.exp(-tracking_probe))
            occlusion_probe = 1 / (1 + np.exp(-occlusion_probe))
            return landmarks_res, poselayer, tracking_probe, occlusion_probe

    def detect(self, img, landmarks):
        dst_size = 128
        landmarks_res = []
        with torch.no_grad():
            input_numpy = np.zeros((len(landmarks), 3, dst_size, dst_size), dtype=np.float32)
            all_mat = []
            for ix, landmark in enumerate(landmarks):
                M = landmark_processor.get_transform_mat_mmcv(landmark, dst_size)


                all_mat.append(M)

                tmp = cv2.warpAffine(img, M, (dst_size, dst_size))


                input_numpy[ix, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32)
                #
                # cv2.imshow('tmp', tmp)
                # cv2.waitKey()

            in_tensor = torch.from_numpy(input_numpy)
            in_tensor = in_tensor.to(self.device)
            fullyconnected1, poselayer, tracking_probe, occlusion_probe = self.face_alignment_net(in_tensor)
            for ix, pts in enumerate(fullyconnected1):
                orig_pts = (np.reshape(pts, (2, 87)).transpose((1, 0)) * dst_size)
                orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True)
                landmarks_res.append(orig_pts)
            tracking_probe = 1 / (1 + np.exp(-tracking_probe))
            occlusion_probe = 1 / (1 + np.exp(-occlusion_probe))
            return landmarks_res, poselayer, tracking_probe, occlusion_probe

    def forward(self, images, landmarks):
        dst_size = 128
        landmarks_res = []
        with torch.no_grad():
            input_numpy = np.zeros((len(images), 3, dst_size, dst_size), dtype=np.float32)
            assert len(images) == len(landmarks)
            all_mat = []
            for ix, img in enumerate(images):
                landmark = landmarks[ix]

                M = landmark_processor.get_transform_mat_mmcv(landmark, dst_size)

                all_mat.append(M)

                tmp = cv2.warpAffine(img, M, (dst_size, dst_size))

                input_numpy[ix, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32)
                #
                # cv2.imshow('tmp', tmp)
                # cv2.waitKey()

            in_tensor = torch.from_numpy(input_numpy)
            in_tensor = in_tensor.to(self.device)
            fullyconnected1, poselayer, tracking_probe, occlusion_probe = self.face_alignment_net(in_tensor)
            for ix, pts in enumerate(fullyconnected1):
                orig_pts = (np.reshape(pts, (2, 87)).transpose((1, 0)) * dst_size)
                orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True)
                landmarks_res.append(orig_pts)
            tracking_probe = 1 / (1 + np.exp(-tracking_probe))
            occlusion_probe = 1 / (1 + np.exp(-occlusion_probe))
            return landmarks_res, poselayer, tracking_probe, occlusion_probe

    def stable_forward(self, image, detected_faces):
        for face_rect in detected_faces:
            if len(self.trackingFaceRects) == 0:
                new_tracking_rect = [face_rect, True, [0, 0], 0]
                self.trackingFaceRects.append(new_tracking_rect)

        with torch.no_grad():
            landmarks = []
            for tracking_face_rect in self.trackingFaceRects:
                if tracking_face_rect[1] == True:
                    d = tracking_face_rect[0]
                    src_center = np.array([d[2] - (d[2] - d[0]) / 2.0, d[3] - (d[3] - d[1]) / 2.0])
                    rotate_degree = tracking_face_rect[3]
                    scale = 128 * 0.8 / min(d[2] - d[0], d[3] - d[1])
                    dst_center = np.array([0.5, 0.5]) * 128
                    offset = dst_center - src_center
                    print('hello')
                else:
                    dst_left_anchor = np.array([0.395, 0.52]) * 128
                    dst_right_anchor = np.array([1 - 0.395, 0.52]) * 128
                    # use last anchors
                    src_center = (tracking_face_rect[2][0] + tracking_face_rect[2][1]) / 2
                    rotate_radian = math.atan2(tracking_face_rect[2][1][1] - tracking_face_rect[2][0][1], tracking_face_rect[2][1][0] - tracking_face_rect[2][0][0])
                    rotate_degree = rotate_radian / math.pi * 180
                    print('degree', rotate_degree)
                    dst_anchor_len = cv2.norm(dst_left_anchor - dst_right_anchor)
                    src_anchor_len = cv2.norm(tracking_face_rect[2][0], tracking_face_rect[2][1])
                    scale = dst_anchor_len / src_anchor_len
                    dst_center = (dst_left_anchor + dst_right_anchor) / 2
                    offset = dst_center - src_center

                M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale)
                M[:, 2] += offset

                inp = cv2.warpAffine(image, M, (128, 128))

                cv2.imshow('inp_stable', inp)
                # cv2.waitKey()

                inp = inp.transpose((2, 0, 1)).astype(np.float32)
                inp = inp[np.newaxis, :, :, :]

                in_tensor = torch.from_numpy(inp)
                in_tensor = in_tensor.to(self.device)
                fullyconnected1, poselayer, tracking_probe, occlusion_probe = self.face_alignment_net(in_tensor)
                fullyconnected1 = fullyconnected1[0]
                poselayer = poselayer[0]
                tracking_probe = tracking_probe[0]
                occlusion_probe = occlusion_probe[0]
                orig_pts = (np.reshape(fullyconnected1, (2, 87)).transpose((1, 0)) * 128)

                orig_pts = landmark_processor.transform_points(orig_pts, M, invert=True)
                # orig_pts = orig_pts.transpose((1, 0))
                fullyconnected1 = orig_pts
                tracking_probe = 1 / (1 + np.exp(-tracking_probe))
                occlusion_probe = 1 / (1 + np.exp(-occlusion_probe))

                # update tracking infos
                tracking_face_rect[1] = False
                tracking_face_rect[2] = [fullyconnected1[51], fullyconnected1[57]]
                tracking_face_rect[3] = rotate_degree

                landmarks.append(fullyconnected1)
        return landmarks

if __name__ == '__main__':
    all_jpegs = glob.glob(r'E:\deepfacelab_data\expression_dst\7201806132018061208311920180612083119\*.jpg')
    for s_filename_path in all_jpegs:
        img = cv2.imread(s_filename_path)

        dflpng = DATAIMG(str(s_filename_path), print_on_no_embedded_data=True)
        if dflpng is None:
            print('ERROR')

        landmarks = dflpng.get_landmarks()

        mmcv = MomocvFaceAlignment()
        fullyconnected1, poselayer, tracking_probe, occlusion_probe = mmcv.forward([img], [landmarks])
        fullyconnected1 = fullyconnected1[0]
        fullyconnected1 = (np.reshape(fullyconnected1, (2, 87)).transpose((1, 0))).astype(np.int32)
        occlusion_probe = occlusion_probe[0]
        for ix, pt in enumerate(fullyconnected1):
            cv2.circle(img, (int(pt[0]), int(pt[1])), 1, (0, 255, 0) if occlusion_probe[ix] > 0.1 else (0, 0, 255), 2)
            # cv2.putText(img, str(ix), (int(pt[0]), int(pt[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 255, 0), 1)
        cv2.imshow('img', img)
        cv2.waitKey()


    # anchor_dis = 0.445
    # dst_size = 128
    # anchors_template = np.array([[anchor_dis, 0.52], [1 - anchor_dis, 0.52]]) * dst_size
    #
    # src_center = (landmarks[31] + landmarks[35]) / 2
    # src_left_anchor = landmarks[31]
    # src_right_anchor = landmarks[35]
    # rotate_radian = math.atan2(src_right_anchor[1] - src_left_anchor[1], src_right_anchor[0] - src_left_anchor[0])
    # rotate_degree = rotate_radian / math.pi * 180
    # dst_eye_len = np.sqrt(np.sum((anchors_template[0] - anchors_template[1]) ** 2))
    # src_eye_len = np.sqrt(np.sum((src_left_anchor - src_right_anchor) ** 2))
    # scale = dst_eye_len / src_eye_len
    # dst_center = (anchors_template[0] + anchors_template[1]) / 2
    # offset = dst_center - src_center
    #
    # M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale)
    # M[:, 2] += offset
    #
    # tmp = cv2.warpAffine(img, M, (dst_size, dst_size))
    # cv2.imshow('tmp', tmp)
    # cv2.waitKey()
    #
    # mmcv = MomocvFaceAlignment()
    #
    # tmp_input = tmp.transpose((2, 0, 1))[np.newaxis, :, : :].astype(np.float32)
    #
    # # tmp_input = cv2.imread(r'E:\deepfacelab_data\workspace\input.png')
    # # tmp_ori = tmp_input
    # # tmp_input = tmp_input.transpose((2, 0, 1))[np.newaxis, :, :, :].astype(np.float32)
    # fullyconnected1, poselayer, tracking_probe, occlusion_probe = mmcv.forward(torch.from_numpy(tmp_input))
    # fullyconnected1 = fullyconnected1[0]
    # # fullyconnected1 = (np.reshape(fullyconnected1, (2, 87)).transpose((1, 0)) * dst_size).astype(np.int32)
    # for i in range(87):
    #     cv2.circle(tmp, (int(fullyconnected1[i] * 128), int(fullyconnected1[i + 87] * 128)), 1, (255, 0, 0), 1)
    # cv2.imshow('tmp_ori', tmp)
    # cv2.waitKey()
    #
    # for ix, pt in enumerate(landmarks):
    #     cv2.circle(img, (int(pt[0]), int(pt[1])), 1, (255, 0, 0), 2)
    #     cv2.putText(img, str(ix), (int(pt[0]), int(pt[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 255, 0), 1)
    # cv2.imshow('img', img)
    # cv2.waitKey()