jotoy / code Goto Github PK

import cv2
# import dlib
import numpy as np
import torch
import face_alignment
import skimage
from skimage import io
import matplotlib.pyplot as plt
from PIL import Image
import os
from tqdm import tqdm

fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=False)

def calculate_center(img, points, gap):
    def center_point(xa, xb, xc, ya, yb, yc, d):
        if ya == yb:
            yb = yb + d
        k_ab = -(xb - xa) / (yb - ya) #if ya != yb else 1e8
        if yb == yc:
            yb = yb + d
        k_bc = -(xc - xb) / (yc - yb) #if yc != yb else 1e11
        if k_ab == k_bc:
            k_ab += 1
        x_mab = float((xa + xb)/2)
        y_mab = float((ya + yb)/2)
        x_mbc = float((xb + xc)/2)
        y_mbc = float((yb + yc)/2)
        c_x = ((y_mbc - y_mab) + (k_ab * x_mab - k_bc * x_mbc)) / (k_ab - k_bc)
        c_y = k_ab * (c_x - x_mab) + y_mab
        return c_x, c_y

    def distance(x, y, c_x, c_y):
        d = ((x - c_x)**2 + (y - c_y)**2) ** 0.5
        return d

    # map = np.zeros((h, w, 1), dtype=np.uint8)
    map = img
    p1, p2, p3, p4, p5, p6 = points
    x1, y1 = int(p1[0]), int(p1[1])
    x2, y2 = int(p2[0]), int(p2[1])
    x3, y3 = int(p3[0]), int(p3[1])
    x4, y4 = int(p4[0]), int(p4[1])
    x5, y5 = int(p5[0]), int(p5[1])
    x6, y6 = int(p6[0]), int(p6[1])

    if (y1 + y4)/2 == y2:
        y2 -= 3
    if (y1 + y4)/2 == y3:
        y3 -= 3
    if (y1 + y4)/2 == y5:
        y5 += 3
    if (y1 + y4)/2 == y6:
        y6 += 3


    y_top = min(y2, y3)
    y_bot = max(y5, y6)
    x_left = x1
    x_right = x4

    c1_x, c1_y = center_point(x1, x2, x4, y1, y2, y4, -3)
    c2_x, c2_y = center_point(x1, x3, x4, y1, y3, y4, -3)
    c3_x, c3_y = center_point(x1, x6, x4, y1, y6, y4, 3)
    c4_x, c4_y = center_point(x1, x5, x4, y1, y5, y4, 3)

    d1 = (distance(x1, y1, c1_x, c1_y))
    d2 = (distance(x4, y4, c2_x, c2_y))
    d3 = (distance(x1, y1, c3_x, c3_y))
    d4 = (distance(x4, x4, c4_x, c4_y))


    for x in range(x_left - gap, x_right + gap + 1):
        for y in range(y_top - gap, y_bot + gap + 1):
            # map[y, x] = 255.0
            if distance(x, y, c1_x, c1_y) < d1 + gap or distance(x, y, c2_x, c2_y) < d2 + gap:
                if (distance(x, y, c3_x, c3_y) < d3 + gap and distance(x, y, c4_x, c4_y) < d4 + gap):
                    map[y, x] = 128

    # cv2.circle(map, (int(c1_x), int(c1_y)), 3, (128), -1, 8)
    # cv2.circle(map, (int(c4_x), int(c4_y)), 3, (128), -1, 8)

    return map

def gen_eye(input_path, save_path):
    def gen(img):
        w = img.shape[0]
        h = img.shape[1]
        preds = fa.get_landmarks(img)
        if preds is None:
            return np.zeros((1,1, 3), dtype=np.uint8)
        # print(preds[0][0])
        preds = preds[0]
        point = []
        # if (len(preds) > 0):
        #     for k, d in enumerate(preds):  # k: index  d: 68 points
        # TODO: Make this nice
        for i in range(36, 42, 1):
            point.append((preds[i, 0], preds[i, 1]))
        x_left, x_right, y_top, y_bot = get_box(point)
        box = np.zeros((abs(y_bot - y_top + 1), abs(x_right - x_left + 1), 3), dtype=np.uint8)
        for i in range(x_left, x_right + 1):
            for j in range(y_top, y_bot + 1):
                for k in range(3):
                    box[j-y_top][i-x_left][k] = img[j][i][k]
        return box
    def get_box(input):
        p1, p2, p3, p4, p5, p6 = input
        x1, y1 = int(p1[0]), int(p1[1])
        x2, y2 = int(p2[0]), int(p2[1])
        x3, y3 = int(p3[0]), int(p3[1])
        x4, y4 = int(p4[0]), int(p4[1])
        x5, y5 = int(p5[0]), int(p5[1])
        x6, y6 = int(p6[0]), int(p6[1])

        x_left = x1 - 5
        x_right = x4 + 5
        y_top = min(y2,y3) - 5
        y_bot = max(y5, y6) + 5
        return x_left, x_right, y_top, y_bot

    img_list = os.listdir(input_path)
    if not os.path.exists(save_path):
        os.mkdir(save_path)
    for img in tqdm(img_list):
        if img[-3:] == 'bmp':
            input = io.imread(input_path + '/' + img)
            print(img)
            res = gen(input)
            if res.all != 0:
                res = Image.fromarray(res.astype(np.uint8))
                resize_img = res.resize((100, 40), Image.BICUBIC)
                # res = np.array(resize_img, dtype=np.uint8)
                resize_img.save(save_path + '/' + img)

    # res = gen(input)
    # img = Image.fromarray(res.astype(np.uint8))
    # resize_img = img.resize((100, 40), Image.BICUBIC)
    # # res = np.array(resize_img, dtype=np.uint8)
    # resize_img.save('res.png')
    # return res
def gen_eye_mask(input):
    def gen(img):
        w = img.shape[0]
        h = img.shape[1]
        preds = fa.get_landmarks(img)
        if preds is None:
            return [[65, 111, 100, 137],[145, 191, 100, 137]]
        # print(preds[0][0])
        preds = preds[0]
        point = []
        mask = []
        # if (len(preds) > 0):
        #     for k, d in enumerate(preds):  # k: index  d: 68 points
        # TODO: Make this nice
        for i in range(36, 42, 1):
            point.append((preds[i, 0], preds[i, 1]))
        x_left, x_right, y_top, y_bot = get_box(point)
        mask.append([x_left, x_right, y_top, y_bot])
        point = []
        for i in range(42, 48, 1):
            point.append((preds[i, 0], preds[i, 1]))
        x_left, x_right, y_top, y_bot = get_box(point)
        mask.append([x_left, x_right, y_top, y_bot])

        # box = np.zeros((abs(y_bot - y_top + 1), abs(x_right - x_left + 1), 3), dtype=np.uint8)
        # for i in range(x_left, x_right + 1):
        #     for j in range(y_top, y_bot + 1):
        #         for k in range(3):
        #             box[j-y_top][i-x_left][k] = img[j][i][k]
        # return box
        return mask
    def get_box(input):
        p1, p2, p3, p4, p5, p6 = input
        x1, y1 = int(p1[0]), int(p1[1])
        x2, y2 = int(p2[0]), int(p2[1])
        x3, y3 = int(p3[0]), int(p3[1])
        x4, y4 = int(p4[0]), int(p4[1])
        x5, y5 = int(p5[0]), int(p5[1])
        x6, y6 = int(p6[0]), int(p6[1])

        x_left = x1 - 2
        x_right = x4 + 2
        y_top = min(y2,y3) - 5
        y_bot = max(y5, y6) + 5
        return x_left, x_right, y_top, y_bot

    img_batch = tensor_to_np(input)
    res = []
    for i in range(img_batch.shape[0]):
        res.append(gen(img_batch[i]))
    # res = toTensor(res)
    return res

def gen_map(input):
    def gen(img):
        w = img.shape[0]
        h = img.shape[1]
        map = np.ones((w, h, 1), dtype=np.uint8) * 127
        preds = fa.get_landmarks(img)
        if preds is None:
            print('None')
            return np.ones((w, h, 1), dtype=np.uint8) * 255
        # print(preds[0][0])
        preds = preds[0]
        point = []
        # if (len(preds) > 0):
        #     for k, d in enumerate(preds):  # k: index  d: 68 points
        # TODO: Make this nice
        for i in range(36, 48, 1):
            point.append((preds[i, 0], preds[i, 1]))
        res = calculate_center(map, point[:6], 3)
        res = calculate_center(res, point[6:12], 3)

        # overlapping = cv2.addWeighted(img, 0.5, res, 0.5, 0)
        # cv2.imwrite('map.jpg', overlapping)
        # cv2.imwrite('Frame.jpg', img)
        return res

    # fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=False)
    input = tensor_to_np(input)
    res = np.zeros((input.shape[0], input.shape[1], input.shape[2], 1), np.uint8)
    for i in range(len(input)):
        res[i] = gen(input[i])
    res = toTensor(res)
    return res
#
def tensor_to_np(tensor):
    img = tensor.mul(255).byte()
    img = img.cpu().numpy().transpose((0,2,3,1))
    return img

def toTensor(img):
    assert type(img) == np.ndarray,'the img type is {}, but ndarry expected'.format(type(img))
    img = torch.from_numpy(img.transpose((0,3,1,2)))
    return img.float().div(255)  # 255也可以改为256

def main():
    fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=False)

    input = io.imread('maps_cyclegan images-18.jpg')
    w = input.shape[1]
    h = input.shape[0]
    preds = fa.get_landmarks(input)[0]

    point = []
    # if (len(preds) > 0):
    #     for k, d in enumerate(preds):  # k: index  d: 68 points
    # TODO: Make this nice
    plt.figure(figsize=(12,6))  # 图像窗口名称
    plt.imshow(input)
    for i in range(36,48,1):
        plt.plot(preds[i,0], preds[i, 1], marker='.', markersize=2, color='w')
        plt.text(preds[i,0], preds[i, 1], '%d'%i)
        point.append((preds[i,0], preds[i, 1]))
    map = np.ones((h, w, 3), dtype=np.uint8) * 255
    res = calculate_center(input, point[:6], 3)
    res = calculate_center(res, point[6:12], 3)
    plt.axis('off')
    plt.imsave('res.jpg', res)
    plt.savefig('test.jpg')
    plt.show()
    # cv2.waitKey(0)

# detector = dlib.get_frontal_face_detector()
# landmark_predictor = dlib.shape_predictor('shape_predictor_68_face_landmarks.dat')
# img = cv2.imread('test.jpg')
# h = img.shape[0]
# w = img.shape[1]
# faces = detector(img, 1)
# point = []
# if (len(faces) > 0):
#     for k, d in enumerate(faces):
#         cv2.rectangle(img, (d.left(), d.top()), (d.right(), d.bottom()), (255, 255, 255))
#         shape = landmark_predictor(img, d)
#         for i in range(36, 48, 1):
#             cv2.circle(img, (shape.part(i).x, shape.part(i).y), 3, (0, 255, 0), -1, 8)
#             cv2.putText(img, str(i), (shape.part(i).x, shape.part(i).y), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
#                             (255, 255, 255))
#             print(i, (shape.part(i).x, shape.part(i).y))
#             point.append((shape.part(i).x, shape.part(i).y))
#
#     # cv2.imshow('Frame',img)
# map = np.ones((h, w, 3), dtype=np.uint8) * 255
# res = calculate_center(map, point[:6], 3)
# res = calculate_center(res, point[6:12], 3)
#
# overlapping = cv2.addWeighted(img, 0.5, res, 0.5, 0)
# cv2.imwrite('map.jpg', overlapping)
# cv2.imwrite('Frame.jpg', img)


# fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=False)
#
# input = io.imread('test.jpg')
# preds = fa.get_landmarks(input)
if __name__ == '__main__':
    # main()
    input = io.imread('maps_cyclegan images-18.jpg')
    input_path = '/data2/interns/tyzhao/work/work_1/data/face_NIR_1000_madacode/mini_set_HR/img/cycleGAN_data_big/Face/trainB/'
    save_path = '/data2/interns/tyzhao/work/work_1/data/face_NIR_1000_madacode/mini_set_HR/img/eye_data/trainB/'
    gen_eye(input_path, save_path)