File size: 1,308 Bytes
cd0d132 f24cb6a 1826277 f24cb6a 1826277 65f86ae f24cb6a 65f86ae 1826277 f24cb6a 65f86ae f24cb6a 65f86ae f24cb6a 65f86ae f24cb6a cd0d132 1826277 6ffae60 1826277 f24cb6a |
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 |
import os
import torch
from gfpgan import GFPGANer
from PIL import Image
import cv2
class EndpointHandler:
def __init__(self, model_path='GFPGANv1.4.pth'):
# Load the GFPGAN model
self.model_path = model_path
self.bg_upsampler = None # You can set this to RealESRGANer if needed
self.face_enhancer = GFPGANer(
model_path=self.model_path, upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=self.bg_upsampler)
# Ensure the output directory exists
os.makedirs('output', exist_ok=True)
def enhance_image(self, image_path):
# Load the image
img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
# Perform face enhancement
_, _, output = self.face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
# Save the enhanced image
save_path = "output/enhanced_image.png"
cv2.imwrite(save_path, output)
return output, save_path
# Example usage for testing
if __name__ == "__main__":
handler = EndpointHandler()
test_image_path = 'path_to_test_image.jpg' # Replace with your test image path
enhanced_image, save_path = handler.enhance_image(test_image_path)
print(f"Enhanced image saved at {save_path}")
|