59 lines
2.0 KiB
Python
59 lines
2.0 KiB
Python
import os
|
|
import sys
|
|
import argparse
|
|
from PIL import Image
|
|
import numpy as np
|
|
import torch
|
|
|
|
# Add current directory to path
|
|
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
|
|
|
|
# Import the face enhancer
|
|
from face_enhancer import GFPGANFaceEnhancer
|
|
|
|
def main():
|
|
# Parse arguments
|
|
parser = argparse.ArgumentParser(description="Test GFPGAN Face Enhancer")
|
|
parser.add_argument('-i', '--input', type=str, required=True, help='Input image path')
|
|
parser.add_argument('-o', '--output', type=str, default='output', help='Output directory')
|
|
parser.add_argument('-v', '--version', type=str, default='1.4', help='GFPGAN model version')
|
|
parser.add_argument('-s', '--scale', type=int, default=2, help='Upscale factor')
|
|
parser.add_argument('--only_center_face', action='store_true', help='Only process center face')
|
|
parser.add_argument('--bg_upsampler', type=str, default='realesrgan', choices=['realesrgan', 'none'],
|
|
help='Background upsampler')
|
|
args = parser.parse_args()
|
|
|
|
# Check if input exists
|
|
if not os.path.exists(args.input):
|
|
print(f"Error: Input file {args.input} does not exist.")
|
|
return 1
|
|
|
|
# Create output directory
|
|
os.makedirs(args.output, exist_ok=True)
|
|
|
|
# Initialize face enhancer
|
|
face_enhancer = GFPGANFaceEnhancer()
|
|
|
|
# Load and process image
|
|
print(f"Processing {args.input}...")
|
|
img = Image.open(args.input).convert('RGB')
|
|
img_np = np.array(img)
|
|
|
|
# Convert to tensor format (BHWC, float 0-1)
|
|
img_tensor = torch.from_numpy(img_np.astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
# Process the image
|
|
result = face_enhancer.enhance_face(
|
|
img_tensor,
|
|
version=args.version,
|
|
scale=args.scale,
|
|
only_center_face=args.only_center_face,
|
|
bg_upsampler=args.bg_upsampler,
|
|
output_folder=args.output
|
|
)
|
|
|
|
print(f"Done! Output saved to {args.output}/restored_imgs/")
|
|
return 0
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main()) |