Upload gfpgan_enhancer.py

gfpgan_enhancer.py

@@ -0,0 +1,71 @@
+import os
+import cv2
+import torch
+from gfpgan import GFPGANer
+from tqdm import tqdm
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from realesrgan import RealESRGANer
+
+def load_video_to_cv2(input_path):
+    video_stream = cv2.VideoCapture(input_path)
+    fps = video_stream.get(cv2.CAP_PROP_FPS)
+    full_frames = []
+    while True:
+        still_reading, frame = video_stream.read()
+        if not still_reading:
+            video_stream.release()
+            break
+        full_frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+    return full_frames, fps
+
+def save_frames_to_video(frames, output_path, fps):
+    if len(frames) == 0:
+        raise ValueError("No frames to write to video.")
+    
+    height, width, _ = frames[0].shape
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    video_writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+    
+    for frame in frames:
+        video_writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+    
+    video_writer.release()
+
+def process_video_with_gfpgan(input_video_path, output_video_path, model_path='gfpgan/weights/GFPGANv1.4.pth'):
+    # Load video and convert to frames
+    frames, fps = load_video_to_cv2(input_video_path)
+
+
+    realesrgan_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+    bg_upsampler = RealESRGANer(
+        scale=2,
+        model_path="gfpgan/weights/RealESRGAN_x2plus.pth",
+        model=realesrgan_model,
+        tile=400,
+        tile_pad=10,
+        pre_pad=0,
+        half=True)
+    
+    # Set up GFPGAN restorer
+    arch = 'clean'
+    channel_multiplier = 2
+    restorer = GFPGANer(
+        model_path=model_path,
+        upscale=2,
+        arch=arch,
+        channel_multiplier=channel_multiplier,
+        bg_upsampler=bg_upsampler
+    )
+
+    # Enhance each frame
+    enhanced_frames = []
+    print("Enhancing frames...")
+    for frame in tqdm(frames, desc='Processing Frames'):
+        # Enhance face in the frame
+        img = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
+        _, _, enhanced_img = restorer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
+        enhanced_frames.append(cv2.cvtColor(enhanced_img, cv2.COLOR_BGR2RGB))
+
+    # Save the enhanced frames to a video
+    save_frames_to_video(enhanced_frames, output_video_path, fps)
+    print(f'Enhanced video saved at {output_video_path}')