Digital_Human/lipreal.py

###############################################################################
#  Copyright (C) 2024 LiveTalking@lipku https://github.com/lipku/LiveTalking
#  email: lipku@foxmail.com
# 
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#  
#       http://www.apache.org/licenses/LICENSE-2.0
# 
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
###############################################################################

import math
import torch
import numpy as np

#from .utils import *
import os
import time
import cv2
import glob
import pickle
import copy

import queue
from queue import Queue
from threading import Thread, Event
import torch.multiprocessing as mp


from lipasr import LipASR
import asyncio
from av import AudioFrame, VideoFrame
from wav2lip.models import Wav2Lip
from basereal import BaseReal

#from imgcache import ImgCache

from tqdm import tqdm
from logger import logger

device = "cuda" if torch.cuda.is_available() else ("mps" if (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()) else "cpu")
print('Using {} for inference.'.format(device))

def _load(checkpoint_path):
    if device == 'cuda':
       checkpoint = torch.load(checkpoint_path) #,weights_only=True
    else:
       checkpoint = torch.load(checkpoint_path,
                         map_location=lambda storage, loc: storage)
    return checkpoint

def load_model(path):
    model = Wav2Lip()
    logger.info("Load checkpoint from: {}".format(path))
    checkpoint = _load(path)
    s = checkpoint["state_dict"]
    new_s = {}
    for k, v in s.items():
       new_s[k.replace('module.', '')] = v
    model.load_state_dict(new_s)

    model = model.to(device)
    return model.eval()

def load_avatar(avatar_id):
    avatar_path = f"./data/avatars/{avatar_id}"
    full_imgs_path = f"{avatar_path}/full_imgs" 
    face_imgs_path = f"{avatar_path}/face_imgs" 
    coords_path = f"{avatar_path}/coords.pkl"
    
    with open(coords_path, 'rb') as f:
        coord_list_cycle = pickle.load(f)
    input_img_list = glob.glob(os.path.join(full_imgs_path, '*.[jpJP][pnPN]*[gG]'))
    input_img_list = sorted(input_img_list, key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
    frame_list_cycle = read_imgs(input_img_list)
    #self.imagecache = ImgCache(len(self.coord_list_cycle),self.full_imgs_path,1000)
    input_face_list = glob.glob(os.path.join(face_imgs_path, '*.[jpJP][pnPN]*[gG]'))
    input_face_list = sorted(input_face_list, key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
    face_list_cycle = read_imgs(input_face_list)

    return frame_list_cycle,face_list_cycle,coord_list_cycle

@torch.no_grad()
def warm_up(batch_size,model,modelres):
    # 预热函数
    logger.info('warmup model...')
    img_batch = torch.ones(batch_size, 6, modelres, modelres).to(device)
    mel_batch = torch.ones(batch_size, 1, 80, 16).to(device)
    model(mel_batch, img_batch)

def read_imgs(img_list):
    frames = []
    logger.info('reading images...')
    for img_path in tqdm(img_list):
        frame = cv2.imread(img_path)
        frames.append(frame)
    return frames

def __mirror_index(size, index):
    #size = len(self.coord_list_cycle)
    turn = index // size
    res = index % size
    if turn % 2 == 0:
        return res
    else:
        return size - res - 1 

def inference(quit_event,batch_size,face_list_cycle,audio_feat_queue,audio_out_queue,res_frame_queue,model):
    
    #model = load_model("./models/wav2lip.pth")
    # input_face_list = glob.glob(os.path.join(face_imgs_path, '*.[jpJP][pnPN]*[gG]'))
    # input_face_list = sorted(input_face_list, key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
    # face_list_cycle = read_imgs(input_face_list)
    
    #input_latent_list_cycle = torch.load(latents_out_path)
    length = len(face_list_cycle)
    index = 0
    count=0
    counttime=0
    logger.info('start inference')
    while not quit_event.is_set():
        starttime=time.perf_counter()
        mel_batch = []
        try:
            mel_batch = audio_feat_queue.get(block=True, timeout=1)
        except queue.Empty:
            continue
            
        is_all_silence=True
        audio_frames = []
        for _ in range(batch_size*2):
            frame,type,eventpoint = audio_out_queue.get()
            audio_frames.append((frame,type,eventpoint))
            if type==0:
                is_all_silence=False

        if is_all_silence:
            for i in range(batch_size):
                res_frame_queue.put((None,__mirror_index(length,index),audio_frames[i*2:i*2+2]))
                index = index + 1
        else:
            # print('infer=======')
            t=time.perf_counter()
            img_batch = []
            for i in range(batch_size):
                idx = __mirror_index(length,index+i)
                face = face_list_cycle[idx]
                img_batch.append(face)
            img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)

            img_masked = img_batch.copy()
            img_masked[:, face.shape[0]//2:] = 0

            img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
            mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
            
            img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)
            mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device)

            with torch.no_grad():
                pred = model(mel_batch, img_batch)
            pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.

            counttime += (time.perf_counter() - t)
            count += batch_size
            #_totalframe += 1
            if count>=100:
                logger.info(f"------actual avg infer fps:{count/counttime:.4f}")
                count=0
                counttime=0
            for i,res_frame in enumerate(pred):
                #self.__pushmedia(res_frame,loop,audio_track,video_track)
                res_frame_queue.put((res_frame,__mirror_index(length,index),audio_frames[i*2:i*2+2]))
                index = index + 1
            #print('total batch time:',time.perf_counter()-starttime)            
    logger.info('lipreal inference processor stop')

class LipReal(BaseReal):
    @torch.no_grad()
    def __init__(self, opt, model, avatar):
        super().__init__(opt,model, avatar)
        #self.opt = opt # shared with the trainer's opt to support in-place modification of rendering parameters.
        # self.W = opt.W
        # self.H = opt.H

        self.fps = opt.fps # 20 ms per frame
        
        self.batch_size = opt.batch_size
        self.idx = 0
        self.res_frame_queue = Queue(self.batch_size*2)  #mp.Queue
        #self.__loadavatar()
        self.model = model
        self.frame_list_cycle,self.face_list_cycle,self.coord_list_cycle = avatar

        self.asr = LipASR(opt,self)
        self.asr.warm_up()
        
        self.render_event = mp.Event()
    
    # def __del__(self):
    #     logger.info(f'lipreal({self.sessionid}) delete')

    import copy
    import cv2
    import numpy as np

    import copy
    import cv2
    import numpy as np
    def paste_back_frame(self, pred_frame, idx: int):
        """
        解决唇形缩小+位置偏移 + 修复脖子旁背景色彩不一致问题
        """
        # ========== 保留调试代码 ==========
        # try:
        #     import os
        #     debug_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'debug')
        #     os.makedirs(debug_dir, exist_ok=True)
        #     bbox = self.coord_list_cycle[idx]
        #     y1, y2, x1, x2 = bbox
        #     debug_frame = copy.deepcopy(self.frame_list_cycle[idx])
        #     cv2.rectangle(debug_frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 2)
        #     cv2.imwrite(os.path.join(debug_dir, f'frame_{idx}_bbox.jpg'), debug_frame)
        #     cv2.imwrite(os.path.join(debug_dir, f'frame_{idx}_pred.jpg'), pred_frame.astype(np.uint8))
        #     logger.info(f"帧{idx}：bbox坐标(y1={y1},y2={y2},x1={x1},x2={x2}) | 生成帧尺寸{pred_frame.shape}")
        # except Exception as e:
        #     logger.error(f"调试代码报错：{str(e)}")
        # ==============================================

        # ========== 基础坐标处理 ==========
        bbox = self.coord_list_cycle[idx]
        combine_frame = copy.deepcopy(self.frame_list_cycle[idx])
        y1, y2, x1, x2 = bbox
        h_ori, w_ori = combine_frame.shape[:2]

        # 精准校准原始嘴巴区域坐标（仅框住嘴唇，不包含脖子/背景）
        y1 = int(y1) if (y1 > 0) else 280  # 仅框嘴唇上沿
        y2 = int(y2) if (y2 < h_ori) else 350  # 仅框嘴唇下沿
        x1 = int(x1) if (x1 > 0) else 180  # 仅框嘴唇左沿
        x2 = int(x2) if (x2 < w_ori) else 320  # 仅框嘴唇右沿

        if (x2 - x1) <= 0 or (y2 - y1) <= 0:
            return combine_frame

        # ========== 1. 强制1:1缩放（保证唇形尺寸） ==========
        target_w = x2 - x1
        target_h = y2 - y1
        res_frame = cv2.resize(
            pred_frame.astype(np.uint8),
            (target_w, target_h),
            interpolation=cv2.INTER_CUBIC
        )

        # ========== 2. 色彩匹配（仅匹配唇部，不影响背景） ==========
        ori_face_region = combine_frame[y1:y2, x1:x2]

        def simple_color_match(source, target):
            source = source.astype(np.float32)
            target = target.astype(np.float32)
            for i in range(3):
                src_mean = np.mean(source[:, :, i])
                trg_mean = np.mean(target[:, :, i])
                # 降低色彩校正强度，减少对背景的影响
                source[:, :, i] = source[:, :, i] * 0.9 + (source[:, :, i] - src_mean + trg_mean) * 0.1
            return np.clip(source, 0, 255).astype(np.uint8)

        color_matched_frame = simple_color_match(res_frame, ori_face_region)

        # ========== 3. 核心优化：唇部专属掩码（仅融合嘴唇，保留背景） ==========
        # 1. 创建椭圆掩码（仅覆盖嘴唇区域，避开脖子/背景）
        mask = np.zeros((target_h, target_w), dtype=np.float32)
        center = (target_w // 2, target_h // 2)
        # 椭圆尺寸：仅覆盖嘴唇核心区域（不超出嘴唇范围）
        axes = (int(target_w * 0.4), int(target_h * 0.35))
        cv2.ellipse(mask, center, axes, 0, 0, 360, 1, -1)

        # 2. 轻量羽化（仅嘴唇边缘，不扩散到背景）
        feather_width = min(4, target_w // 20, target_h // 20)  # 进一步缩小羽化范围
        if feather_width > 0:
            mask = cv2.GaussianBlur(mask, (feather_width * 2 + 1, feather_width * 2 + 1), feather_width)
        mask_3ch = np.repeat(mask[:, :, np.newaxis], 3, axis=2)

        # ========== 4. 精准对齐 + 背景保留 ==========
        paste_x1 = x1
        paste_x2 = x2
        paste_y1 = y1
        paste_y2 = y2

        # ========== 5. 最终融合（仅替换嘴唇，背景完全用原图） ==========
        ori_paste_region = combine_frame[paste_y1:paste_y2, paste_x1:paste_x2].astype(np.float32)
        color_matched_frame = color_matched_frame.astype(np.float32)

        # 核心：掩码只作用于嘴唇，背景区域完全保留原图
        fused_region = mask_3ch * color_matched_frame + (1 - mask_3ch) * ori_paste_region
        fused_region = np.clip(fused_region, 0, 255).astype(np.uint8)

        # 仅替换嘴唇区域，脖子/背景区域完全不动
        combine_frame[paste_y1:paste_y2, paste_x1:paste_x2] = fused_region

        return combine_frame


            
    def render(self,quit_event,loop=None,audio_track=None,video_track=None):
        #if self.opt.asr:
        #     self.asr.warm_up()

        self.init_customindex()
        self.tts.render(quit_event)
        
        infer_quit_event = Event()
        infer_thread = Thread(target=inference, args=(infer_quit_event,self.batch_size,self.face_list_cycle,
                                           self.asr.feat_queue,self.asr.output_queue,self.res_frame_queue,
                                           self.model,))  #mp.Process
        infer_thread.start()
        
        process_quit_event = Event()
        process_thread = Thread(target=self.process_frames, args=(process_quit_event,loop,audio_track,video_track))
        process_thread.start()

        #self.render_event.set() #start infer process render
        count=0
        totaltime=0
        _starttime=time.perf_counter()
        #_totalframe=0
        while not quit_event.is_set(): 
            # update texture every frame
            # audio stream thread...
            t = time.perf_counter()
            self.asr.run_step()

            # if video_track._queue.qsize()>=2*self.opt.batch_size:
            #     print('sleep qsize=',video_track._queue.qsize())
            #     time.sleep(0.04*video_track._queue.qsize()*0.8)
            if video_track and video_track._queue.qsize()>=5:
                logger.debug('sleep qsize=%d',video_track._queue.qsize())
                time.sleep(0.04*video_track._queue.qsize()*0.8)
                
            # delay = _starttime+_totalframe*0.04-time.perf_counter() #40ms
            # if delay > 0:
            #     time.sleep(delay)
        #self.render_event.clear() #end infer process render
        logger.info('lipreal thread stop')

        infer_quit_event.set()
        infer_thread.join()

        process_quit_event.set()
        process_thread.join()