"""polish.py - fix the soft face in a Wav2Lip video with GFPGAN.

Put this file INSIDE your GFPGAN folder (the git clone), with the model file
GFPGANv1.4.pth next to it (the polish card on the page shows where it comes
from). Needs the free ffmpeg tool.

Run:
    python3 polish.py talk.mp4
Out comes talk_hd.mp4 - same lips, same sync, now in focus.
"""
import glob
import os
import subprocess
import sys

import cv2
from gfpgan import GFPGANer

SRC = sys.argv[1] if len(sys.argv) > 1 else "talk.mp4"
FPS = 25                     # Wav2Lip draws its videos at 25 frames per second

# 1) pull the video apart into its frames - a video is just pictures in a row
os.makedirs("frames_in", exist_ok=True)
os.makedirs("frames_out", exist_ok=True)
subprocess.run(["ffmpeg", "-y", "-v", "error", "-i", SRC, "frames_in/%05d.png"], check=True)
frames = sorted(glob.glob("frames_in/*.png"))
print(len(frames), "frames to fix")

# 2) restore every face. upscale=1 keeps the size the same, and GFPGAN only
#    sharpens what is already there - it never MOVES the mouth - so the
#    lip-sync stays perfect.
g = GFPGANer(model_path="GFPGANv1.4.pth", upscale=1, arch="clean",
             channel_multiplier=2, bg_upsampler=None)
for i, f in enumerate(frames):
    _, _, fixed = g.enhance(cv2.imread(f), has_aligned=False,
                            only_center_face=True, paste_back=True)
    cv2.imwrite("frames_out/" + os.path.basename(f), fixed)
    if i % 25 == 0:
        print(f"fixed {i}/{len(frames)}")

# 3) fixed frames + the ORIGINAL sound -> the polished video
#    (yuv420p + faststart = plays on phones)
out = SRC.replace(".mp4", "_hd.mp4")
subprocess.run(["ffmpeg", "-y", "-v", "error",
                "-framerate", str(FPS), "-i", "frames_out/%05d.png",
                "-i", SRC, "-map", "0:v", "-map", "1:a",
                "-c:v", "libx264", "-pix_fmt", "yuv420p", "-movflags", "+faststart",
                "-c:a", "aac", "-shortest", out],
               check=True)
print("done ->", out)