Ex3-고양이 수염달기_회전에 강건하게(2)

1편에서 알고리즘을 정했으니, 이것을 함수로 만들면 좋을 것 같습니다.

0. import dependencies

import cv2
import matplotlib.pyplot as plt
import numpy as np
import os
import dlib

1. 얼굴 인식 모듈과 스티커 준비

model_path = '/content/drive/MyDrive/AIffel2021/Ex03-FaceRecognition/models/shape_predictor_68_face_landmarks.dat'
landmark_predictor = dlib.shape_predictor(model_path)

sticker_path = '/content/drive/MyDrive/AIffel2021/Ex03-FaceRecognition/images/cat-whiskers.png'
img_sticker = cv2.imread(sticker_path)

2. 함수로 만들기

사진 안에 여러 사람이 있는 경우와, 특징점 29와 특징점 27의 x좌표가 같아 ZeroDivisionError가 나는 경우의 예외처리를 추가하고, 저장하려는 위치를 인자로 받아 이미지 파일을 생성하는 함수로 구현했습니다. 이미지 처리의 순서는 각 사람마다 코, 왼쪽 수염, 오른쪽 수염을 붙이는 것으로, 모든 과정이 끝난 후의 이미지를 표시하고, 저장합니다. 세부 내용에 대해서는 1편을 참고해 주십시오.

def add_cat_whisker(my_img_path, dir_to_write):
    # dir_to_write에 이미지 파일을 생성합니다. 
    img_bgr = cv2.imread(my_image_path)  # BGR 기반
    img_show = img_bgr.copy()  # 이 함수는 im_show를 수정합니다
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    detector_hog = dlib.get_frontal_face_detector()
    dlib_rects = detector_hog(img_rgb, 2)
    
    list_landmarks = []
    for dlib_rect in dlib_rects:
        points = landmark_predictor(img_rgb, dlib_rect)
        list_points = list(map(lambda p: (p.x, p.y), points.parts()))
        list_landmarks.append(list_points)
    # list_landmarks[0]은 첫번째 사람의 특징점의 리스트입니다

    ORIGINAL_WIDTH = 512


    # - 코 좌표:
    nose_left = 193
    nose_width = 128
    nose_top = 173
    nose_height = 128

    # -왼쪽 수염 좌표:
    l_whisker_left = 0
    l_whisker_width = 193
    l_whisker_top = 166
    l_whisker_height = 193

    # -오른쪽 수염 좌표:
    r_whisker_left = 319
    r_whisker_width = 193
    r_whisker_top = 166
    r_whisker_height = 193

    # 스티커 자르기
    sticker_nose = img_sticker[nose_top:nose_top+nose_height, nose_left:nose_left+nose_width]
    
    sticker_l_whisker = img_sticker[l_whisker_top:l_whisker_top+l_whisker_height,\
    				l_whisker_left:l_whisker_left+l_whisker_width]
    
    sticker_r_whisker = img_sticker[r_whisker_top:r_whisker_top+r_whisker_height,\
    				r_whisker_left:r_whisker_left+r_whisker_width]
    

    # 여러 명인 경우(bbox가 여러개)
    for i, _ in enumerate(list_landmarks):
        landmarks = list_landmarks[i]  # (i+1)번째 사람의 특징점 리스트


        # 코 붙이기
        # 각도
        dy = landmarks[29][1] - landmarks[27][1]
        dx = landmarks[29][0] - landmarks[27][0]
        try:
            angle = math.atan(dy/dx)
        except ZeroDivisionError:  # dx == 0 인 경우 각도를 직접 지정합니다
            angle = (math.pi/2 - 1e-5)
        angle = angle * 180 / math.pi
        sign = angle / abs(angle)

        if angle > 0:  # 오른쪽으로 기울어진 경우와 왼쪽으로 기울어진 경우를 나눕니다.
            nose_angle = 90 - angle
        else:
            nose_angle = -90 - angle

        # 크기 & 중심점
        perp_nose_vector = [landmarks[35][0] - landmarks[31][0], landmarks[35][1] - landmarks[31][1]] # 입술의 양단을 잇습니다
        nose_length = int(np.linalg.norm(perp_nose_vector)) # 슬라이싱에 사용되므로 좌표의 경우 int형으로 변환해 줍니다
        nose_size = (nose_length, nose_length)
        offset_x, offset_y = 0.1 * perp_nose_vector[0], 0.1 * perp_nose_vector[1] # 미세조정, 경험적으로 이렇게 하는 게 조금 더 나은 것 같습니다.

        nose_x, nose_y = landmarks[30] # 중심점
        nose_x, nose_y = int(nose_x + sign * offset_x), int(nose_y + sign * offset_y)
	    
        # 스티커 붙이기
        sticker_nose = cv2.resize(sticker_nose, nose_size)
        rows, cols = sticker_nose.shape[:2]
        # 이미지에서 스티커를 붙일 시작점의 좌표를 pin_nose_start로 합니다.
        pin_nose_start = (nose_x - rows // 2, nose_y - cols //2)
        M= cv2.getRotationMatrix2D((cols//2, rows//2), nose_angle, 1)
        dst = cv2.warpAffine(sticker_nose, M,(cols, rows), borderValue=(255,255,255)) # 경계는 흰색

        sticker_area = \
        img_show[pin_nose_start[1] : pin_nose_start[1] + cols, pin_nose_start[0] : pin_nose_start[0] + rows] 
        # img_show에서 원본으로 사용할 부분을 변수로 저장합니다.

        img_show[pin_nose_start[1] : pin_nose_start[1] + cols, pin_nose_start[0] : pin_nose_start[0] + rows]  =\
        np.where(dst==255, sticker_area, dst).astype(np.uint8) 
        # 스티커 붙이기. 스티커가 완전한 백색인 부분은 붙이지 않습니다. (sticker_area 사용)


        # 왼쪽 수염 붙이기
        # 각도
        left_whisker_vector = [landmarks[29][0] - landmarks[2][0] ,landmarks[29][1] - landmarks[2][1]]
        l_angle = math.atan(left_whisker_vector[1]/left_whisker_vector[0])
        l_angle = l_angle * 180 / math.pi
        l_angle = sign * abs(l_angle)

        # 크기 & 중심점
        left_whisker_length = np.linalg.norm(left_whisker_vector)
        left_whisker_size = (int(left_whisker_length * 0.6), int(left_whisker_length * l_whisker_height / l_whisker_width * 0.6)) # scale factor = 0.6, 볼 밖으로 나오지 않게 하기 위해 사이즈를 줄였습니다
        l_whisker_x, l_whisker_y =(landmarks[29][0] + landmarks[2][0]) // 2, (landmarks[29][1] + landmarks[2][1]) // 2 # 몫 연산 후 int가 됩니다

	# 스티커 붙이기
        sticker_l_whisker = cv2.resize(sticker_l_whisker, left_whisker_size)
        rows, cols = sticker_l_whisker.shape[:2]
        # 이미지에서 스티커를 붙일 시작점의 좌표를 pin_l_whisker_start로 합니다.
        pin_l_whisker_start = (l_whisker_x - rows // 2,	l_whisker_y - cols // 2)
        M= cv2.getRotationMatrix2D((cols//2, rows//2),l_angle, 1)
        dst = cv2.warpAffine(sticker_l_whisker, M,(cols, rows), borderValue=(255,255,255))

        sticker_area = img_show[pin_l_whisker_start[1]: pin_l_whisker_start[1] + cols,\
                pin_l_whisker_start[0]: pin_l_whisker_start[0] + rows]
                    # img_show에서 원본으로 사용할 부분을 변수로 저장합니다.

        img_show[pin_l_whisker_start[1]: pin_l_whisker_start[1] + cols,\
            pin_l_whisker_start[0]: pin_l_whisker_start[0] + rows] = \
                np.where(dst==255, sticker_area,dst).astype(np.uint8)
                # 스티커 붙이기. 스티커가 완전한 백색인 부분은 붙이지 않습니다. (sticker_area 사용)
                

        # 오른쪽 수염 붙이기
        # 각도
        right_whisker_vector = [landmarks[29][0] - landmarks[14][0], landmarks[29][1] - landmarks[14][1]]
        r_angle = math.atan(right_whisker_vector[1]/right_whisker_vector[0])
        r_angle = r_angle * 180 / math.pi
        r_angle = sign * abs(r_angle)

        # 크기 & 중심점
        right_whisker_length = np.linalg.norm(right_whisker_vector)
        right_whisker_size = (int(right_whisker_length * 0.6), int(right_whisker_length * r_whisker_height / r_whisker_width * 0.6)) # scale factor = 0.6, 볼 밖으로 나오지 않게 하기 위해 사이즈를 줄였습니다
        r_whisker_x, r_whisker_y = (landmarks[14][0] + landmarks[29][0]) // 2, (landmarks[14][1] + landmarks[29][1]) // 2

        # 스티커 붙이기
        sticker_r_whisker = cv2.resize(sticker_r_whisker, right_whisker_size)
        rows, cols = sticker_r_whisker.shape[:2]
        # 이미지에서 스티커를 붙일 시작점의 좌표를 pin_r_whisker_start로 합니다.
        pin_r_whisker_start = (r_whisker_x - rows // 2,	r_whisker_y - cols // 2)
        M= cv2.getRotationMatrix2D((cols//2, rows//2),r_angle, 1)
        dst = cv2.warpAffine(sticker_r_whisker, M,(cols, rows), borderValue=(255,255,255))

        sticker_area = img_show[pin_r_whisker_start[1]: pin_r_whisker_start[1] + cols,\
                pin_r_whisker_start[0]: pin_r_whisker_start[0] + rows] 
                # img_show에서 원본으로 사용할 부분을 변수로 저장합니다.

        img_show[pin_r_whisker_start[1]: pin_r_whisker_start[1] + cols,\
            pin_r_whisker_start[0]: pin_r_whisker_start[0] + rows] = \
                np.where(dst==255, sticker_area,dst).astype(np.uint8) 
                # 스티커 붙이기. 스티커가 완전한 백색인 부분은 붙이지 않습니다. (sticker_area 사용)	
    cv2.imwrite(dir_to_write ,img_show)  # 이미지 파일을 생성합니다 
    plt.imshow(cv2.cvtColor(img_show, cv2.COLOR_BGR2RGB))
    plt.show()

3. 실행 결과

여러 사람에 대해서도 잘 붙은 것을 확인할 수 있습니다.