inference_for_facefusion.py

#!/usr/bin/env python3

print('\rloading os          ', end='')
import os
import re
import glob
import sys
import inspect

#加载第三方包
sys.path.append('{}/facefusion-wav2lip/'.format(os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))))

print('\rloading facefusion', end='')
#from facefusion_wav.facefusion import core_wav2lip
from facefusion.core_wav2lip import cli
print('\rloading torch       ', end='')
import torch

print('\rloading numpy       ', end='')
import numpy as np

print('\rloading Image       ', end='')
from PIL import Image

print('\rloading argparse    ', end='')
import argparse

print('\rloading math        ', end='')
import math



print('\rloading subprocess  ', end='')
import subprocess

print('\rloading pickle      ', end='')
import pickle

print('\rloading cv2         ', end='')
import cv2

print('\rloading audio       ', end='')
import audio

print('\rloading Wav2Lip     ', end='')
from models import Wav2Lip

print('\rRloading RetinaFace ', end='')
from batch_face import RetinaFace

print('\rloading re          ', end='')
import re

print('\rloading partial     ', end='')
from functools import partial

print('\rloading tqdm        ', end='')
from tqdm import tqdm

print('\rloading warnings    ', end='')
import warnings

warnings.filterwarnings("ignore", category=UserWarning, module='torchvision.transforms.functional_tensor')

print('\rloading load_model  ', end='')
from easy_functions import load_model

print('\rimports loaded!     ')

device = 'cuda'
parser = argparse.ArgumentParser(description='Inference code to lip-sync videos in the wild using Wav2Lip models')

parser.add_argument('--checkpoint_path', type=str,
                    help='Name of saved checkpoint to load weights from', required=True)

parser.add_argument('--segmentation_path', type=str, default="checkpoints/face_segmentation.pth",
                    help='Name of saved checkpoint of segmentation network', required=False)

parser.add_argument('--lianmian', type=str,
                    help='Filepath of video/image that contains faces to use', required=True)
parser.add_argument('--audio', type=str,
                    help='Filepath of video/audio file to use as raw audio source', required=True)

parser.add_argument('--srt', type=str,
                    help='Filepath of srt file to use as video subtitle file', required=True)

parser.add_argument('--outfile', type=str, help='Video path to save result. See default for an e.g.',
                    default='results/result_voice.mp4')

parser.add_argument('--static', type=bool,
                    help='If True, then use only first video frame for inference', default=False)
parser.add_argument('--fps', type=float, help='Can be specified only if input is a static image (default: 25)',
                    default=25., required=False)

parser.add_argument('--pads', nargs='+', type=int, default=[0, 10, 0, 0],
                    help='Padding (top, bottom, left, right). Please adjust to include chin at least')

parser.add_argument('--wav2lip_batch_size', type=int, help='Batch size for Wav2Lip model(s)', default=1)

parser.add_argument('--out_height', default=480, type=int,
                    help='Output video height. Best results are obtained at 480 or 720')

parser.add_argument('--crop', nargs='+', type=int, default=[0, -1, 0, -1],
                    help='Crop video to a smaller region (top, bottom, left, right). Applied after resize_factor and rotate arg. '
                         'Useful if multiple face present. -1 implies the value will be auto-inferred based on height, width')

parser.add_argument('--box', nargs='+', type=int, default=[-1, -1, -1, -1],
                    help='Specify a constant bounding box for the face. Use only as a last resort if the face is not detected.'
                         'Also, might work only if the face is not moving around much. Syntax: (top, bottom, left, right).')

parser.add_argument('--rotate', default=False, action='store_true',
                    help='Sometimes videos taken from a phone can be flipped 90deg. If true, will flip video right by 90deg.'
                         'Use if you get a flipped result, despite feeding a normal looking video')

parser.add_argument('--nosmooth', type=str, default=False,
                    help='Prevent smoothing face detections over a short temporal window')

parser.add_argument('--no_seg', default=False, action='store_true',
                    help='Prevent using face segmentation')

parser.add_argument('--no_sr', default=False, action='store_true',
                    help='Prevent using super resolution')

parser.add_argument('--sr_model', type=str, default='gfpgan',
                    help='Name of upscaler - gfpgan or RestoreFormer', required=False)

parser.add_argument('--fullres', default=3, type=int,
                    help='used only to determine if full res is used so that no resizing needs to be done if so')

parser.add_argument('--debug_mask', type=str, default=False,
                    help='Makes background grayscale to see the mask better')

parser.add_argument('--preview_settings', type=str, default=False,
                    help='Processes only one frame')

parser.add_argument('--mouth_tracking', type=str, default=False,
                    help='Tracks the mouth in every frame for the mask')

parser.add_argument('--mask_dilation', default=150, type=float,
                    help='size of mask around mouth', required=False)

parser.add_argument('--mask_feathering', default=151, type=int,
                    help='amount of feathering of mask around mouth', required=False)

parser.add_argument('--quality', type=str, help='Choose between Fast, Improved',default='Fast')


with open(os.path.join('checkpoints', 'predictor.pkl'), 'rb') as f:
    predictor = pickle.load(f)

with open(os.path.join('checkpoints', 'mouth_detector.pkl'), 'rb') as f:
    mouth_detector = pickle.load(f)

# creating variables to prevent failing when a face isn't detected
kernel = last_mask = x = y = w = h = None


def Experimental(img, original_img, run_params):
    global kernel, last_mask, x, y, w, h  # Add last_mask to global variables

    # Convert color space from BGR to RGB if necessary
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    original_img = cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB)

    if str(args.debug_mask) == 'True':
        original_img = cv2.cvtColor(original_img, cv2.COLOR_RGB2GRAY)
        original_img = cv2.cvtColor(original_img, cv2.COLOR_GRAY2RGB)

    # Detect face
    faces = mouth_detector(img)
    if len(faces) == 0:
        if last_mask is not None:
            last_mask = cv2.resize(last_mask, (img.shape[1], img.shape[0]))
            mask = last_mask  # use the last successful mask
        else:
            cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img)
            return img, None
    else:
        face = faces[0]
        shape = predictor(img, face)

        # Get points for mouth
        mouth_points = np.array([[shape.part(i).x, shape.part(i).y] for i in range(48, 68)])

        # Calculate bounding box dimensions
        x, y, w, h = cv2.boundingRect(mouth_points)

        # Set kernel size as a fraction of bounding box size
        kernel_size = int(max(w, h) * args.mask_dilation)
        upscale_kernel_size = int(max(w, h) * max(args.mask_dilation, 2.5))

        # Create kernels
        kernel = np.ones((kernel_size, kernel_size), np.uint8)
        upscale_kernel = np.ones((upscale_kernel_size, upscale_kernel_size), np.uint8)

        # Create binary mask for mouth
        mask = np.zeros(img.shape[:2], dtype=np.uint8)
        cv2.fillConvexPoly(mask, mouth_points, 255)

        last_mask = mask  # Update last_mask with the new mask

    # Dilate the mask for upscaling
    upscale_dilated_mask = cv2.dilate(mask, upscale_kernel)
    dilated_mask = cv2.dilate(mask, kernel)

    # Find contours in the dilated mask
    contours, _ = cv2.findContours(upscale_dilated_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # Find bounding box coordinates for each contour
    for contour in contours:
        x_dilated, y_dilated, w_dilated, h_dilated = cv2.boundingRect(contour)

        # Crop the image to the bounding box of the dilated mask
        cropped_img = img[y_dilated:y_dilated + h_dilated, x_dilated:x_dilated + w_dilated]

        # Save the cropped image here
        # cv2.imwrite('temp/cp.jpg', cropped_img)

        # Upscale the cropped image
        upscaled_img = upscale(cropped_img, run_params)

        # cv2.imwrite('temp/ucp.jpg', upscaled_img)

        if str(args.debug_mask) == 'True':
            img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)

        # Paste the upscaled image back onto the original image
        img[y_dilated:y_dilated + h_dilated, x_dilated:x_dilated + w_dilated] = upscaled_img

    # Calculate distance transform of dilated mask
    dist_transform = cv2.distanceTransform(dilated_mask, cv2.DIST_L2, 5)

    # Normalize distance transform
    cv2.normalize(dist_transform, dist_transform, 0, 255, cv2.NORM_MINMAX)

    # Convert normalized distance transform to binary mask and convert it to uint8
    _, masked_diff = cv2.threshold(dist_transform, 50, 255, cv2.THRESH_BINARY)
    masked_diff = masked_diff.astype(np.uint8)

    if not args.mask_feathering == 0:
        blur = args.mask_feathering
        # Set blur size as a fraction of bounding box size
        blur = int(max(w, h) * blur)  # 10% of bounding box size
        if blur % 2 == 0:  # Ensure blur size is odd
            blur += 1
        masked_diff = cv2.GaussianBlur(masked_diff, (blur, blur), 0)

    # Convert numpy arrays to PIL Images
    input1 = Image.fromarray(img)
    input2 = Image.fromarray(original_img)

    # Convert mask to single channel where pixel values are from the alpha channel of the current mask
    mask = Image.fromarray(masked_diff)

    # Ensure images are the same size
    assert input1.size == input2.size == mask.size

    # Paste input1 onto input2 using the mask
    input2.paste(input1, (0, 0), mask)

    # Convert the final PIL Image back to a numpy array
    input2 = np.array(input2)

    cv2.cvtColor(input2, cv2.COLOR_BGR2RGB, input2)

    return input2, mask


def create_tracked_mask(img, original_img):
    global kernel, last_mask, x, y, w, h  # Add last_mask to global variables

    # Convert color space from BGR to RGB if necessary
    cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img)
    cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB, original_img)

    # Detect face
    faces = mouth_detector(img)
    if len(faces) == 0:
        if last_mask is not None:
            last_mask = cv2.resize(last_mask, (img.shape[1], img.shape[0]))
            mask = last_mask  # use the last successful mask
        else:
            cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img)
            return img, None
    else:
        face = faces[0]
        shape = predictor(img, face)

        # Get points for mouth
        mouth_points = np.array([[shape.part(i).x, shape.part(i).y] for i in range(48, 68)])

        # Calculate bounding box dimensions
        x, y, w, h = cv2.boundingRect(mouth_points)

        # Set kernel size as a fraction of bounding box size
        kernel_size = int(max(w, h) * args.mask_dilation)
        # if kernel_size % 2 == 0:  # Ensure kernel size is odd
        # kernel_size += 1

        # Create kernel
        kernel = np.ones((kernel_size, kernel_size), np.uint8)

        # Create binary mask for mouth
        mask = np.zeros(img.shape[:2], dtype=np.uint8)
        cv2.fillConvexPoly(mask, mouth_points, 255)

        last_mask = mask  # Update last_mask with the new mask

    # Dilate the mask
    dilated_mask = cv2.dilate(mask, kernel)

    # Calculate distance transform of dilated mask
    dist_transform = cv2.distanceTransform(dilated_mask, cv2.DIST_L2, 5)

    # Normalize distance transform
    cv2.normalize(dist_transform, dist_transform, 0, 255, cv2.NORM_MINMAX)

    # Convert normalized distance transform to binary mask and convert it to uint8
    _, masked_diff = cv2.threshold(dist_transform, 50, 255, cv2.THRESH_BINARY)
    masked_diff = masked_diff.astype(np.uint8)

    # make sure blur is an odd number
    blur = args.mask_feathering
    if blur % 2 == 0:
        blur += 1
    # Set blur size as a fraction of bounding box size
    blur = int(max(w, h) * blur)  # 10% of bounding box size
    if blur % 2 == 0:  # Ensure blur size is odd
        blur += 1
    masked_diff = cv2.GaussianBlur(masked_diff, (blur, blur), 0)

    # Convert numpy arrays to PIL Images
    input1 = Image.fromarray(img)
    input2 = Image.fromarray(original_img)

    # Convert mask to single channel where pixel values are from the alpha channel of the current mask
    mask = Image.fromarray(masked_diff)

    # Ensure images are the same size
    assert input1.size == input2.size == mask.size

    # Paste input1 onto input2 using the mask
    input2.paste(input1, (0, 0), mask)

    # Convert the final PIL Image back to a numpy array
    input2 = np.array(input2)

    # input2 = cv2.cvtColor(input2, cv2.COLOR_BGR2RGB)
    cv2.cvtColor(input2, cv2.COLOR_BGR2RGB, input2)

    return input2, mask


def create_mask(img, original_img):
    global kernel, last_mask, x, y, w, h  # Add last_mask to global variables

    # Convert color space from BGR to RGB if necessary
    cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img)
    cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB, original_img)

    if last_mask is not None:
        last_mask = np.array(last_mask)  # Convert PIL Image to numpy array
        last_mask = cv2.resize(last_mask, (img.shape[1], img.shape[0]))
        mask = last_mask  # use the last successful mask
        mask = Image.fromarray(mask)

    else:
        # Detect face
        faces = mouth_detector(img)
        if len(faces) == 0:
            cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img)
            return img, None
        else:
            face = faces[0]
            shape = predictor(img, face)

            # Get points for mouth
            mouth_points = np.array([[shape.part(i).x, shape.part(i).y] for i in range(48, 68)])

            # Calculate bounding box dimensions
            x, y, w, h = cv2.boundingRect(mouth_points)

            # Set kernel size as a fraction of bounding box size
            kernel_size = int(max(w, h) * args.mask_dilation)
            # if kernel_size % 2 == 0:  # Ensure kernel size is odd
            # kernel_size += 1

            # Create kernel
            kernel = np.ones((kernel_size, kernel_size), np.uint8)

            # Create binary mask for mouth
            mask = np.zeros(img.shape[:2], dtype=np.uint8)
            cv2.fillConvexPoly(mask, mouth_points, 255)

            # Dilate the mask
            dilated_mask = cv2.dilate(mask, kernel)

            # Calculate distance transform of dilated mask
            dist_transform = cv2.distanceTransform(dilated_mask, cv2.DIST_L2, 5)

            # Normalize distance transform
            cv2.normalize(dist_transform, dist_transform, 0, 255, cv2.NORM_MINMAX)

            # Convert normalized distance transform to binary mask and convert it to uint8
            _, masked_diff = cv2.threshold(dist_transform, 50, 255, cv2.THRESH_BINARY)
            masked_diff = masked_diff.astype(np.uint8)

            if not args.mask_feathering == 0:
                blur = args.mask_feathering
                # Set blur size as a fraction of bounding box size
                blur = int(max(w, h) * blur)  # 10% of bounding box size
                if blur % 2 == 0:  # Ensure blur size is odd
                    blur += 1
                masked_diff = cv2.GaussianBlur(masked_diff, (blur, blur), 0)

            # Convert mask to single channel where pixel values are from the alpha channel of the current mask
            mask = Image.fromarray(masked_diff)

            last_mask = mask  # Update last_mask with the final mask after dilation and feathering

    # Convert numpy arrays to PIL Images
    input1 = Image.fromarray(img)
    input2 = Image.fromarray(original_img)

    # Resize mask to match image size
    # mask = Image.fromarray(mask)
    mask = mask.resize(input1.size)

    # Ensure images are the same size
    assert input1.size == input2.size == mask.size

    # Paste input1 onto input2 using the mask
    input2.paste(input1, (0, 0), mask)

    # Convert the final PIL Image back to a numpy array
    input2 = np.array(input2)

    # input2 = cv2.cvtColor(input2, cv2.COLOR_BGR2RGB)
    cv2.cvtColor(input2, cv2.COLOR_BGR2RGB, input2)

    return input2, mask


def get_smoothened_boxes(boxes, T):
    for i in range(len(boxes)):
        if i + T > len(boxes):
            window = boxes[len(boxes) - T:]
        else:
            window = boxes[i: i + T]
        boxes[i] = np.mean(window, axis=0)
    return boxes


def face_detect(images, results_file='last_detected_face.pkl'):
    #发现处理完一个视频，再处理下一个视频会报错，临时关掉加载last_detected_face.pkl文件
    # # If results file exists, load it and return
    # if os.path.exists(results_file):
    #     print('Using face detection data from last input')
    #     with open(results_file, 'rb') as f:
    #         return pickle.load(f)

    results = []
    pady1, pady2, padx1, padx2 = args.pads
    from tqdm import tqdm
    tqdm = partial(tqdm, position=0, leave=True)

    for image, rect in tqdm(zip(images, face_rect(images)), total=len(images), desc="detecting face in every frame",
                            ncols=100):
        if rect is None:
            cv2.imwrite('temp/faulty_frame.jpg', image)  # check this frame where the face was not detected.
            raise ValueError('Face not detected! Ensure the video contains a face in all the frames.')

        y1 = max(0, rect[1] - pady1)
        y2 = min(image.shape[0], rect[3] + pady2)
        x1 = max(0, rect[0] - padx1)
        x2 = min(image.shape[1], rect[2] + padx2)

        results.append([x1, y1, x2, y2])

    boxes = np.array(results)
    if str(args.nosmooth) == 'False': boxes = get_smoothened_boxes(boxes, T=5)
    results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)]

    # #发现处理完一个视频，再处理下一个视频会报错，临时关掉加载last_detected_face.pkl文件
    # # Save results to file
    # with open(results_file, 'wb') as f:
    #     pickle.dump(results, f)

    return results


def datagen(frames, mels):
    img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
    print("\r" + " " * 100, end="\r")
    if args.box[0] == -1:
        if not args.static:
            face_det_results = face_detect(frames)  # BGR2RGB for CNN face detection
        else:
            face_det_results = face_detect([frames[0]])
    else:
        print('Using the specified bounding box instead of face detection...')
        y1, y2, x1, x2 = args.box
        face_det_results = [[f[y1: y2, x1:x2], (y1, y2, x1, x2)] for f in frames]

    for i, m in enumerate(mels):
        idx = 0 if args.static else i % len(frames)
        frame_to_save = frames[idx].copy()
        face, coords = face_det_results[idx].copy()

        face = cv2.resize(face, (args.img_size, args.img_size))

        img_batch.append(face)
        mel_batch.append(m)
        frame_batch.append(frame_to_save)
        coords_batch.append(coords)

        if len(img_batch) >= args.wav2lip_batch_size:
            img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)

            img_masked = img_batch.copy()
            img_masked[:, args.img_size // 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])

            yield img_batch, mel_batch, frame_batch, coords_batch
            img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []

    if len(img_batch) > 0:
        img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)

        img_masked = img_batch.copy()
        img_masked[:, args.img_size // 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])

        yield img_batch, mel_batch, frame_batch, coords_batch


mel_step_size = 16
device = 'cuda'


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


def main():
    args.img_size = 96
    frame_number = 11

    if os.path.isfile(args.lianmian) and args.lianmian.split('.')[1] in ['jpg', 'png', 'jpeg']:
        args.static = True

    if not os.path.isfile(args.lianmian):
        raise ValueError('--face argument must be a valid path to video/image file')

    elif args.lianmian.split('.')[1] in ['jpg', 'png', 'jpeg']:
        full_frames = [cv2.imread(args.lianmian)]
        fps = args.fps

    else:
        if args.fullres != 1:
            print('Resizing video...')
        video_stream = cv2.VideoCapture(args.lianmian)
        fps = video_stream.get(cv2.CAP_PROP_FPS)

        full_frames = []
        while 1:
            still_reading, frame = video_stream.read()
            if not still_reading:
                video_stream.release()
                break

            if args.fullres != 1:
                aspect_ratio = frame.shape[1] / frame.shape[0]
                frame = cv2.resize(frame, (int(args.out_height * aspect_ratio), args.out_height))

            if args.rotate:
                frame = cv2.rotate(frame, cv2.cv2.ROTATE_90_CLOCKWISE)

            y1, y2, x1, x2 = args.crop
            if x2 == -1: x2 = frame.shape[1]
            if y2 == -1: y2 = frame.shape[0]

            frame = frame[y1:y2, x1:x2]

            full_frames.append(frame)

    if not args.audio.endswith('.wav'):
        print('Converting audio to .wav')
        subprocess.check_call([
            "ffmpeg", "-y", "-loglevel", "error",
            "-i", args.audio,
            "temp/temp.wav",
        ])
        args.audio = 'temp/temp.wav'

    print('analysing audio...')
    wav = audio.load_wav(args.audio, 16000)
    mel = audio.melspectrogram(wav)

    if np.isnan(mel.reshape(-1)).sum() > 0:
        raise ValueError('Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again')

    mel_chunks = []

    mel_idx_multiplier = 80. / fps
    i = 0
    while 1:
        start_idx = int(i * mel_idx_multiplier)
        if start_idx + mel_step_size > len(mel[0]):
            mel_chunks.append(mel[:, len(mel[0]) - mel_step_size:])
            break
        mel_chunks.append(mel[:, start_idx: start_idx + mel_step_size])
        i += 1;

    full_frames = full_frames[:len(mel_chunks)]
    if str(args.preview_settings) == 'True':
        full_frames = [full_frames[0]]
        mel_chunks = [mel_chunks[0]]
    print (str(len(full_frames)) + ' frames to process')
    batch_size = args.wav2lip_batch_size
    if str(args.preview_settings) == 'True':
        gen = datagen(full_frames, mel_chunks)
    else:
        gen = datagen(full_frames.copy(), mel_chunks)
    #存储所有帧
    #all_fps = []
    for i, (img_batch, mel_batch, frames, coords) in enumerate(tqdm(
            gen,
            total=int(np.ceil(float(len(mel_chunks)) / batch_size)),
            desc="Processing Wav2Lip", ncols=100
    )):

        if i == 0:
            if not args.quality == 'Fast':
                print(f"mask size: {args.mask_dilation}, feathering: {args.mask_feathering}")
            print("Starting...")
            #frame_h, frame_w = full_frames[0].shape[:-1]
            #fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Be sure to use lower case
            #out = cv2.VideoWriter('temp/result.mp4', fourcc, fps, (frame_w, frame_h))

        img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to('cuda')
        mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to('cuda')

        with torch.no_grad():
            pred = model(mel_batch, img_batch)

        pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.
        inddd = i
        for p, f, c in zip(pred, frames, coords):
            y1, y2, x1, x2 = c

            of = f
            p = cv2.resize(p.astype(np.uint8), (x2 - x1, y2 - y1))
            cf = f[y1:y2, x1:x2]

            if args.quality in ['Improved']:
                if str(args.mouth_tracking) == 'True':
                    for i in range(len(frames)):
                        p, last_mask = create_tracked_mask(p, cf)
                else:
                    for i in range(len(frames)):
                        p, last_mask = create_mask(p, cf)

            f[y1:y2, x1:x2] = p
            
            # all_fps.append(f)
            #列表存储，改为jpg存储
            current_path = os.path.dirname(os.path.abspath(__file__))
            temp_images_path =os.path.join(current_path, r'temp/images')
            if not os.path.exists(temp_images_path):
                os.makedirs(temp_images_path)
            cv2.imwrite(os.path.join(temp_images_path,'{}.jpg'.format(inddd)), f)
            # cv2.imwrite('temp/p.jpg', f)
            

            # if args.quality == 'Experimental':
            #     last_mask = None
            #     for i in range(len(frames)):
            #         f, last_mask = Experimental(f, of, run_params)
            #
            # if str(args.preview_settings) == 'True':
            #     cv2.imwrite('temp/preview.jpg', f)
            #
            # else:
            #     out.write(f)
    #对所有帧频进行处理
    # result_all_frames = cli(all_fps, 12, 1, "cuda", 'face_enhancer')
    # for tuple_item in sorted(result_all_frames, key=lambda x: x[0]):
    #     index, frame = tuple_item
    #     out.write(frame)
    # out.release()
    cli(temp_images_path, 12, 1, "cuda", 'face_enhancer')
    image_paths = from_file_get_images(temp_images_path)
      # 读取第一张图以获取图像尺寸
    img = cv2.imread(image_paths[0])
    height, width, _ = img.shape

    # 设置视频编解码器和输出文件格式
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 你也可以使用其他编解码器，例如 'XVID'
    out = cv2.VideoWriter("temp/result.mp4", fourcc, fps, (width, height))

    # 逐一读取图片并写入视频
    for image_path in image_paths:
        img = cv2.imread(image_path)
        out.write(img)

    # 释放 VideoWriter 对象
    out.release()
  
    if str(args.preview_settings) == 'False':
        print("converting to final video")
	#查看音频文件路径下有没有srt字幕文件，如果有，则合并音频、视频、字幕。如果没有则只合并视频和音频
        audio_path_part = os.path.dirname(args.audio)
        srt_files = glob.glob(os.path.join(audio_path_part, '*.srt'))
        font_path = '思源黑体旧字形 Bold'
        if args.srt:
            ffmpeg_command = (
            	f"ffmpeg -i 'temp/result.mp4' -i {args.audio} -vf \"subtitles={args.srt}:force_style='FontName=\\\"{font_path}\\\",FontSize=20,PrimaryColour=16777215,Alignment=2,OutlineColour=0,Outline=1,Shadow=0,BorderStyle=1,MarginL=10,MarginR=10,MarginV=10,LineSpacing=10'\" "
            	f"-c:v libx264 -c:a aac -b:a 192k -map 0:v -map 1:a {args.outfile}"
            )
            result = subprocess.check_call(ffmpeg_command, shell=True)
        else:
            subprocess.check_call([
            	"ffmpeg", "-y", "-loglevel", "error",
            	"-i", "temp/result.mp4",
            	"-i", args.audio,
            	"-c:v", "h264_nvenc",
            	args.outfile,
            ])


model = detector = detector_model = None


def do_load(checkpoint_path):
    global model, detector, detector_model
    model = load_model(checkpoint_path)
    detector = RetinaFace(gpu_id=0, model_path="checkpoints/mobilenet.pth", network="mobilenet")
    detector_model = detector.model

#获取目录内的图片并按照从小到大排序生成列表
def from_file_get_images(path):
	# Get a list of image files in the folder
	image_files = [f for f in os.listdir(path) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif'))]

	# Sort the image files based on their numeric part in the filename
	image_paths = sorted(
		image_files,
		key=lambda s: [int(text) if text.isdigit() else text.lower() for text in re.split('([0-9]+)', s)]
	)

	# Create a list of full image paths
	image_paths = [os.path.join(path, file) for file in image_paths]
	return image_paths

def face_rect(images):
    face_batch_size = 8
    num_batches = math.ceil(len(images) / face_batch_size)
    prev_ret = None
    for i in range(num_batches):
        batch = images[i * face_batch_size: (i + 1) * face_batch_size]
        all_faces = detector(batch)  # return faces list of all images
        for faces in all_faces:
            if faces:
                box, landmarks, score = faces[0]
                prev_ret = tuple(map(int, box))
            yield prev_ret


if __name__ == '__main__':
    args = parser.parse_args()
    do_load(args.checkpoint_path)
    main()