Spaces:
Running
Running
File size: 2,856 Bytes
523fb10 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
import os
import cv2
import numpy as np
import glob
import tqdm
import shutil
import argparse
from face_enhancement import FaceEnhancement
def process_video(target_path, out_path, faceenhancer):
fps = 25.0
os.makedirs(out_path, exist_ok=True)
original_vid_path = target_path
vid_name = "out.mp4"
if not os.path.isdir(target_path):
vid_name = target_path.split("/")[-1]
vidcap = cv2.VideoCapture(target_path)
fps = vidcap.get(cv2.CAP_PROP_FPS)
try:
for match in glob.glob(os.path.join("./tmp/", "*.png")):
os.remove(match)
for match in glob.glob(os.path.join(out_path, "*.png")):
os.remove(match)
except Exception as e:
print(e)
os.makedirs("./tmp/", exist_ok=True)
os.system(
f"ffmpeg -i {target_path} -qscale:v 1 -qmin 1 -qmax 1 -vsync 0 ./tmp/frame_%d.png"
)
target_path = "./tmp/"
else:
print("folder not implemented.")
exit()
globbed_images = sorted(glob.glob(os.path.join(target_path, "*.png")))
for image in tqdm.tqdm(globbed_images):
name = image.split("/")[-1]
filename = os.path.join(out_path, name)
im = cv2.imread(image, cv2.IMREAD_COLOR) # BGR
h, w, _ = im.shape
# im = cv2.resize(im, (0,0), fx=2, fy=2) #optional
img, orig_faces, enhanced_faces = faceenhancer.process(im)
img = cv2.resize(img, (w, h))
cv2.imwrite(filename, img)
# merge frames to video
video_save_path = os.path.join(out_path, vid_name)
os.system(
f"ffmpeg -y -r {fps} -i {out_path}/frame_%d.png -i {original_vid_path}"
f" -map 0:v:0 -map 1:a? -c:a copy -c:v libx264 -r {fps} -pix_fmt yuv420p {video_save_path}"
)
# delete tmp file
shutil.rmtree("./tmp/")
for match in glob.glob(os.path.join(out_path, "*.png")):
os.remove(match)
if __name__ == "__main__":
model = {
"name": "GPEN-BFR-512",
"in_size": 512,
"out_size": 512,
"channel_multiplier": 2,
"narrow": 1,
}
parser = argparse.ArgumentParser()
parser.add_argument("--indir", type=str, required=True, help="input file")
parser.add_argument(
"--outdir",
type=str,
required=True,
help="Please provide output folder which has no more than one parent dir that has not been created.",
)
args = parser.parse_args()
os.makedirs(args.outdir, exist_ok=True)
faceenhancer = FaceEnhancement(
use_sr=True,
in_size=model["in_size"],
out_size=model["out_size"],
model=model["name"],
channel_multiplier=model["channel_multiplier"],
narrow=model["narrow"],
)
process_video(
args.indir,
args.outdir,
faceenhancer,
)
|