close_lip.py

"""
close_lip.py


唇を閉じた画像を作ります。

唇の開いている範囲をinpaintします。
そこに、hole部分を影として戻します。



著者: Akihito Miyazaki
作成日: 2024-04-23
更新履歴:
  - 2024-04-23: 最初のリリース
  - 1024-11-16:coverted to huggingface-space (but args broken)
"""
import os
import cv2
import numpy as np
from PIL import Image
import lip_utils


from glibvision.cv2_utils import blend_rgb_images
from glibvision.numpy_utils import apply_binary_mask_to_color,create_2d_image

import argparse





def create_top_lip_low_mask(image,face_landmarks_list,line_thick = 1):
    black = create_2d_image(image.shape)
    lip_utils.fill_top_lower(black,face_landmarks_list,line_thick,lip_utils.COLOR_WHITE)
    return  black

def create_lip_hole_mask(image,face_landmarks_list,line_thick = 1):
    black = create_2d_image(image.shape)
    lip_utils.fill_lip_hole(black,face_landmarks_list,line_thick,lip_utils.COLOR_WHITE)

    return  black

def process_close_lip_image(img,landmarks_list):
    img_h, img_w = lip_utils.get_image_size(img)

    hole_mask = create_lip_hole_mask(img,landmarks_list,0)

    lower_lip_mask = create_top_lip_low_mask(img,landmarks_list)

    #these make dirty
    #kernel = np.ones((3, 3), np.uint8)
    #lower_lip_mask = cv2.erode(lower_lip_mask, kernel, iterations=1)


    if lip_utils.DEBUG:
        cv2.imwrite("close_lip_01_mask.jpg",lower_lip_mask)
    
    mixed_mask = cv2.bitwise_or(hole_mask,lower_lip_mask)
    if lip_utils.DEBUG:
        cv2.imwrite("close_lip_01_mask_mixed.jpg",mixed_mask)

    img_inpainted = cv2.inpaint(img, mixed_mask,3, cv2.INPAINT_NS)
    if lip_utils.DEBUG:
        cv2.imwrite("close_lip_02_inpaint.jpg",img_inpainted)
   

    copy_impainted=img_inpainted.copy()
    apply_binary_mask_to_color(copy_impainted,(0,8,50),hole_mask)
    #lip_utils.fill_lip_hole(img_inpainted,landmarks_list,0,(0,8,50)) # BGR
    if lip_utils.DEBUG:
        cv2.imwrite("close_lip_03_hole.jpg",copy_impainted)
   

    

    ## Inpaintした画像をぼかす。
    blurred_image = cv2.GaussianBlur(copy_impainted, (9, 9), 0) #場合によっては奇数じゃないとエラーが出ることがある
    if lip_utils.DEBUG:
        cv2.imwrite("close_lip_04_burred.jpg",blurred_image)


     # Inpaintの境界線から少し広げている
    kernel = np.ones((3, 3), np.uint8)
    shrink_mask = cv2.erode(hole_mask, kernel, iterations=1)


    shrink_burred_mask = cv2.GaussianBlur(shrink_mask, (3, 3), 0)
    if lip_utils.DEBUG:
        cv2.imwrite("close_lip_05_final_hole_mask.jpg",shrink_burred_mask)

    img_inpainted = blend_rgb_images(img_inpainted,blurred_image,shrink_burred_mask)
    if lip_utils.DEBUG:
        cv2.imwrite("close_lip_05_final_hole.jpg",img_inpainted)
    # Inpaintの境界線から少し広げている
    kernel = np.ones((3, 3), np.uint8)
    extend_mask = cv2.dilate(lower_lip_mask, kernel, iterations=1)


    extend_burred_mask = cv2.GaussianBlur(extend_mask, (3, 3), 0)
    if lip_utils.DEBUG:
        cv2.imwrite("close_lip_05_final_lip_mask.jpg",extend_burred_mask)
    img_inpainted = blend_rgb_images(img_inpainted,blurred_image,extend_burred_mask)
    
    if lip_utils.DEBUG:
        cv2.imwrite("close_lip_05_final_lip.jpg",img_inpainted)

    mixed_mask = cv2.bitwise_or(shrink_burred_mask,extend_burred_mask)
    mixed_mask[mixed_mask>0] = 255
    mixed_mask = cv2.dilate(mixed_mask, (1,1), iterations=1)
    #    mixed_mask = cv2.GaussianBlur(mixed_mask, (3, 3), 0)

    if lip_utils.DEBUG:
        cv2.imwrite("close_lip_05_final_mixed_mask.jpg",mixed_mask)

    return img_inpainted,mixed_mask

if __name__ == "__main__":
  parser = argparse.ArgumentParser(description='Open Mouth')
  parser.add_argument('--input',"-i",help='変換する画像の元(必須) 口を閉じていること',required=True)
  parser.add_argument('--output',"-o",help='画像の保存先(別途一時的なレイヤーファイルも作られる)')
  parser.add_argument('--landmark',"-l",help='landmarkdata')
  parser.add_argument('--scale',"-sc",help='スケール精度が上がる',default=4,type=int)

  args = parser.parse_args()
  # 画像ファイルのパス
  img_path = args.input
  img = cv2.imread(img_path)
  #landmarks_list = landmark_utils.load_landmarks(img,args.scale,args.landmark)
  landmarks_list = None
  eye_closed_image,mask = process_close_lip_image(img,landmarks_list)

  output_path = args.output
  if output_path == None:
    parent_path,file = os.path.split(img_path)
    name,ext = os.path.splitext(file)
    output_path = os.path.join(parent_path,f"{name}_lipclose{ext}")
      

  

  parent_path,file = os.path.split(output_path)
  name,ext = os.path.splitext(file)

  mask_path = os.path.join(parent_path,f"{name}_mask{ext}")
  cv2.imwrite(mask_path,mask)
  cv2.imwrite(output_path,eye_closed_image)
  print(f"complete image {output_path} and mask {mask_path}")