#coding:utf-8 import os import math import shutil import time from step05_detect_fa_hairmatting_inplace import pkl_process from PIL import Image, ImageFont, ImageDraw import cv2 import numpy as np import torch import pickle import json import glob import uuid from gpt4v_caption import caption_image from utils.call_hair_train import call_hair_train import os import os.path as osp from common.logger import LogFactory from core.process_modules import Process_Data, localtranslationwarpfastwithstrength,\ Generator_Hair, chinClass, Generator_Fusion_Res, Change_Hair_Color, GenderClassifyProcessor, BodySeg, \ localtranslationwarpfastwithstrength_v2, localtranslationwarpfastwithstrength_v2_soft, updateEndPosition from process_modules import Get_Landmark as Get_Landmark_mtcnn from core.face_enhance.face_enhancement import FaceEnhancement from datetime import datetime import random from utils import landmark_processor # from core.cos_module import COS_object as OSS_object from core.oss_module import OSS_object from core.faceseg.face_seg import FaceSeg from common.logger import config from process_modules import PersonProcessor_yolov5,KeypointsProcessor,Human_Keypoints,pt_conv_25_to_17 from gen_super_image import get_high_train_img from core.MMCVFaceRecognitionServer import MomocvFaceRecognitionServer import process_modules from common.logger import config version = config.get('default', 'version') if version == "local": train_gpu_nums = 2 else: train_gpu_nums = 1 # webui_services = ["http://hairservice.tslead.net:32678/", "http://hairservice.tslead.net:32679/"] class HairInit(object): __instance = None __first_init = False def __new__(cls, gpu=True, use_enhance=False, infer_use_enhance=False, color_user_enhance=False): if not cls.__instance: cls.__instance = object.__new__(cls) return cls.__instance def __init__(self, gpu=True, use_enhance=False, infer_use_enhance=False, color_user_enhance=False): if not self.__first_init: worker_id = int(os.environ.get('APP_WORKER_ID', 1)) rand_max = 9527 self.gpu_index = (worker_id + rand_max) % train_gpu_nums os.environ['CUDA_VISIBLE_DEVICES'] = str(self.gpu_index) print('current worker id {} set the gpu id :{}'.format(worker_id, self.gpu_index)) device_id = self.gpu_index self.get_landmark = Get_Landmark_mtcnn(gpu_id=device_id) self.get_landmark_mtcnn = Get_Landmark_mtcnn(gpu_id=device_id) self.face_recognition = MomocvFaceRecognitionServer(gpu_id=device_id) self.hair_size = 768 self.color_output_size = 768 self.process_data = Process_Data(gpu, device_id) self.process_data_infer = process_modules.Process_Data(gpu, device_id) self.generator_hair = Generator_Hair(gpu, device_id) self.hair_fusion = Generator_Fusion_Res(gpu, device_id) self.use_enhance = use_enhance self.face_enhance = FaceEnhancement(512, device_id) self.change_haircolor = Change_Hair_Color(gpu, device_id) self.logger_init = LogFactory.getLogger("init") self.logger_process = LogFactory.getLogger("process") self.logger_call = LogFactory.getLogger("call") self.oss2 = OSS_object() self.face_seg = FaceSeg(device_id) self.chin_cls = chinClass(device_id) model_path = "./weights/gender_models" self.output_img_size = 128 if not os.path.exists(model_path): print("GenderClassifyProcessor don't have model!") self.gender_model = GenderClassifyProcessor(gpu_id=device_id) print("Load model finish ... ") self.baseColor_dir = os.path.join(config.get('default', "haircolorDir"), config.get('default', "baseColor_ID")) for i in range(2): image = cv2.imread('data/front.jpg') with torch.no_grad(): tmp_dir = config.get('default', "tmp_dir") if not os.path.exists(tmp_dir): os.makedirs(tmp_dir) self.infer_hairstyle(image, 'data/template', tmp_dir, 'test.jpg') self.effect_prepare_mask_fc32 = cv2.imread("./data/mask.png").astype(np.float32) / 255 # body process self.person_processor = PersonProcessor_yolov5(gpu_id=device_id) self.keypoints_processor = KeypointsProcessor(gpu_id=device_id) self.human_keypoint = Human_Keypoints(gpu=True, device_id=device_id) print("Load model finish ... ") # infer init # self.table_enlight = cv2.imread('./data/convert_enlight.png') self.change_haircolor = Change_Hair_Color(gpu, device_id) self.gender_classify = GenderClassifyProcessor(gpu_id=device_id) HairInit.__first_init = True print("[HAIR_INIT] All init done. ") def infer_hairstyle(self, origin_img, hairstyle_dir, userinfo_dir, mask_newname): config_path = os.path.join(hairstyle_dir, "config.json") if not os.path.exists(config_path): return None, None, 10001 config_info = json.load(open(config_path, "rb")) gender = config_info["gender"] ratio = int(config_info["ratio"]) # print("gender: ", gender, " ratio: ", ratio) user_rgb_8uc3_orisize = origin_img.copy() tt = time.time() landmark1k_dir = osp.join(userinfo_dir, 'kpt_1k.txt') if not osp.exists(landmark1k_dir): with torch.no_grad(): landmarks_origin_img_1k = self.get_landmark.forward(origin_img) if landmarks_origin_img_1k is None: return None, None, 10001 np.savetxt(landmark1k_dir, landmarks_origin_img_1k) else: landmarks_origin_img_1k = np.loadtxt(landmark1k_dir) user_bald_res_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_res_ori_raw.png') # res_matting_8uc3_bald_orisize_dir = osp.join(userinfo_dir, 'res_matting_mask_ori.png') user_baldseg_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_seg_ori.png') user_baldseg_8uc3_768_dir = osp.join(userinfo_dir, 'user_baldseg_768.png') user_bald_8uc3_768_dir = osp.join(userinfo_dir, 'bald_seg_768.png') user_landmark_f1k2_768_dir = osp.join(userinfo_dir, 'landmark_f1k2_768.txt') user_hairstyle_M_dir = osp.join(userinfo_dir, 'hairstyle_M.txt') pre_list = [user_bald_res_8uc3_orisize_dir, user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_768_dir, user_bald_8uc3_768_dir, user_landmark_f1k2_768_dir, user_hairstyle_M_dir] condition_exist = True for tmp_file in pre_list: if not osp.exists(tmp_file): condition_exist = False if not condition_exist: user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, user_baldseg_8uc3_768, user_bald_8uc3_768, \ user_landmark_f1k2_768, user_hairstyle_M,_ = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) cv2.imwrite(user_bald_res_8uc3_orisize_dir, user_bald_res_8uc3_orisize) cv2.imwrite(user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_orisize) cv2.imwrite(user_baldseg_8uc3_768_dir, user_baldseg_8uc3_768) cv2.imwrite(user_bald_8uc3_768_dir, user_bald_8uc3_768) np.savetxt(user_landmark_f1k2_768_dir, user_landmark_f1k2_768) np.savetxt(user_hairstyle_M_dir, user_hairstyle_M) else: user_bald_res_8uc3_orisize = cv2.imread(user_bald_res_8uc3_orisize_dir) user_baldseg_8uc3_orisize = cv2.imread(user_baldseg_8uc3_orisize_dir) # user_baldseg_8uc3_768 = cv2.imread(user_baldseg_8uc3_768_dir) # user_bald_8uc3_768 = cv2.imread(user_bald_8uc3_768_dir) # user_landmark_f1k2_768 = np.loadtxt(user_landmark_f1k2_768_dir) # user_hairstyle_M = np.loadtxt(user_hairstyle_M_dir) if ratio == 0: user_hairstyle_M = self.process_data.get_hair_M_boy_v1(landmarks_origin_img_1k) elif ratio == 1: user_hairstyle_M = self.process_data.get_hair_M_girl_v1(landmarks_origin_img_1k) elif ratio == 2: user_hairstyle_M = self.process_data.get_hair_M_girl_v2(landmarks_origin_img_1k) else: user_hairstyle_M = self.process_data.get_hair_M_girl_v1(landmarks_origin_img_1k) user_landmark_f1k2_768 = landmark_processor.transform_points(landmarks_origin_img_1k, user_hairstyle_M) user_bald_8uc3_768 = cv2.warpAffine(user_bald_res_8uc3_orisize, user_hairstyle_M, (self.hair_size, self.hair_size)) user_baldseg_8uc3_768 = cv2.warpAffine(user_baldseg_8uc3_orisize, user_hairstyle_M, (self.hair_size, self.hair_size)) print('condition cosst:', time.time() - tt ) # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) # ratio = 1536. / max(show_concat.shape[:2]) # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) # cv2.imshow("show_concat", show_concat) # cv2.waitKey() another_pose_hair_img_path = os.path.join(hairstyle_dir, "input_another_pose_hair_image.npy") if not os.path.exists(another_pose_hair_img_path): return None, None, 10002 another_pose_hair_image = np.load(another_pose_hair_img_path) # cv2.imshow('user_baldseg_8uc3_768', user_baldseg_8uc3_768) # cv2.imshow('user_bald_8uc3_768', user_bald_8uc3_768) t0 = time.time() # 换发型 hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, user_baldseg_8uc3_768, user_bald_8uc3_768, user_landmark_f1k2_768, gender) self.logger_process.info('Generator_Hair_inference_use_pref costs:{}'.format(time.time() - t0)) t1 = time.time() hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) res_matting_mask_ori = osp.join(userinfo_dir, mask_newname) res_matting_mask_ori_raw = osp.join(userinfo_dir, 'res_matting_mask_ori_raw.png') if osp.exists(res_matting_mask_ori): os.remove(res_matting_mask_ori) if osp.exists(res_matting_mask_ori_raw): os.remove(res_matting_mask_ori_raw) self.logger_process.info('get_fusion_res_hairpaste costs:{}'.format(time.time() - t1)) user_res_8uc3_orisize = self.hair_fusion.inference(hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize, landmarks_origin_img_1k, user_baldseg_8uc3_orisize, ratio) hair_gene_matte_8uc3_orisize_cp = hair_gene_matte_8uc3_orisize.copy() # show_concat = np.concatenate((hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize, user_baldseg_8uc3_orisize, user_res_8uc3_orisize), axis=1) # ratio = 1536. / max(show_concat.shape[:2]) # show_concat = cv2.resize(show_concat, (0, 0), fx=ratio, fy=ratio) # cv2.imshow("mid_concat", show_concat) # cv2.waitKey() t2 = time.time() user_res_8uc3_orisize_for_haircolor = user_res_8uc3_orisize.copy() # if self.use_enhance: user_res_8uc3_orisize_enhance = self.face_enhance.process(user_res_8uc3_orisize, landmarks_origin_img_1k) t3 = time.time() self.logger_process.info('use_enhance costs:{}'.format(t3 - t2)) user_bald_res_8uc3_orisize_enhance, user_baldseg_8uc3_orisize_enhance, user_baldseg_8uc3_768_enhance, user_bald_8uc3_768_enhance, \ user_landmark_f1k2_768_enhance, user_hairstyle_M_enhance, user_matting_8uc3_bald_orisize = self.process_data.get_prepare_user_768_data( user_res_8uc3_orisize_enhance, landmarks_origin_img_1k, ratio=ratio) t4 = time.time() self.logger_process.info('gen bald costs:{}'.format(t4 - t3)) # 重新提取matting _, hair_gene_matte_8uC0_orisize, _ = self.process_data.generator_matte.matte_inference(user_res_8uc3_orisize_enhance, landmarks_origin_img_1k) t5 = time.time() self.logger_process.info('gen matte_inference costs:{}'.format(t5 - t4)) hair_gene_matte_8uc3_orisize = np.repeat(hair_gene_matte_8uC0_orisize[:, :, np.newaxis], 3, axis=2) hair_gene_matte_8uc3_orisize = cv2.blur(hair_gene_matte_8uc3_orisize, (3, 3)) # user_res_8uc3_orisize = user_bald_res_8uc3_orisize_enhance.astype(np.float32)/255. * (1. - hair_gene_matte_8uc3_orisize.astype(np.float32) / 255.) + user_res_8uc3_orisize_enhance.astype(np.float32)/255. * hair_gene_matte_8uc3_orisize.astype(np.float32) / 255. # user_res_8uc3_orisize = (user_res_8uc3_orisize*255).astype(np.uint8) user_res_fc32_orisize_enhance = user_res_8uc3_orisize_enhance.astype(np.float32)/255. user_res_fc32_orisize = user_res_8uc3_orisize.astype(np.float32)/255. hair_gene_matte_fc32_orisize = hair_gene_matte_8uc3_orisize.astype(np.float32)/255. hair_gene_matte_fc32_orisize = cv2.GaussianBlur(hair_gene_matte_fc32_orisize, (11, 11), 0, 0) user_res_fc32_orisize_enhance = user_res_fc32_orisize_enhance * hair_gene_matte_fc32_orisize + user_res_fc32_orisize * (1 - hair_gene_matte_fc32_orisize) user_res_8uc3_orisize_enhance = (user_res_fc32_orisize_enhance * 255).astype(np.uint8) # mid_show = np.concatenate((user_res_fc32_orisize_enhance, user_res_8uc3_orisize_enhance, hair_gene_matte_8uc3_orisize, user_res_8uc3_orisize), axis=1) # ratio = 1536. / max(mid_show.shape[:2]) # mid_show = cv2.resize(mid_show, (0, 0), fx=ratio, fy=ratio) # cv2.imshow("mid_show", mid_show) # cv2.imshow("user_res_8uc3_orisize_enhance_fg", user_res_8uc3_orisize_enhance.astype(np.float32)/255. * hair_gene_matte_8uc3_orisize.astype(np.float32) / 255) # cv2.waitKey() user_bald_res_8uc3_orisize_for_fusion_dir = osp.join(userinfo_dir, 'bald_res_ori.png') # if osp.exists(user_bald_res_8uc3_orisize_for_fusion_dir): # os.remove(user_bald_res_8uc3_orisize_for_fusion_dir) cv2.imwrite(user_bald_res_8uc3_orisize_for_fusion_dir, user_bald_res_8uc3_orisize_enhance) # if osp.exists(res_matting_mask_ori_raw): # os.remove(res_matting_mask_ori_raw) cv2.imwrite(res_matting_mask_ori_raw, hair_gene_matte_8uc3_orisize) # user_bald_res_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_res_ori.png') # cv2.imwrite(user_bald_res_8uc3_orisize_dir, user_res_8uc3_orisize) res_fix_img_mask_8uc4_orisize = np.concatenate((user_res_8uc3_orisize_enhance, hair_gene_matte_8uc3_orisize_cp[:, :, :1]), axis=2) # cv2.imshow('') cv2.imwrite(res_matting_mask_ori, res_fix_img_mask_8uc4_orisize) print('costs:', time.time() - t5) return user_res_8uc3_orisize_enhance, user_res_8uc3_orisize_for_haircolor, 0