| import cv2 |
| import numpy as np |
| import onnxruntime as ort |
| import gradio as gr |
| from PIL import Image |
|
|
| |
| onnx_path = "models/gfpgan_1.4.onnx" |
| session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"]) |
|
|
| def preprocess_face(img, box, size=512): |
| x1, y1, x2, y2 = [int(v) for v in box] |
| face = img[y1:y2, x1:x2] |
| face = cv2.resize(face, (size, size), interpolation=cv2.INTER_LINEAR) |
| face = face.astype(np.float32) / 255.0 |
| face = face.transpose(2, 0, 1)[np.newaxis, :] |
| return face, (x1, y1, x2, y2) |
|
|
| def postprocess_face(output, face_box, original_shape): |
| restored = output.squeeze().transpose(1, 2, 0) |
| restored = np.clip(restored * 255, 0, 255).astype(np.uint8) |
| restored = cv2.resize(restored, (face_box[2] - face_box[0], face_box[3] - face_box[1])) |
| return restored |
|
|
| def detect_faces(img): |
| gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) |
| face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml') |
| boxes = face_cascade.detectMultiScale(gray, 1.3, 5) |
| return [[x, y, x + w, y + h] for (x, y, w, h) in boxes] |
|
|
| def enhance_faces(image_pil): |
| image = np.array(image_pil.convert("RGB")) |
| img = image.copy() |
| faces = detect_faces(img) |
| |
| if not faces: |
| return image_pil, "No faces detected." |
|
|
| for box in faces: |
| face_input, face_coords = preprocess_face(img, box) |
| output = session.run(None, {session.get_inputs()[0].name: face_input})[0] |
| restored = postprocess_face(output, box, img.shape) |
|
|
| |
| x1, y1, x2, y2 = box |
| img[y1:y2, x1:x2] = restored |
|
|
| return Image.fromarray(img), f"{len(faces)} face(s) enhanced." |
|
|
| |
| iface = gr.Interface( |
| fn=enhance_faces, |
| inputs=gr.Image(type="pil", label="Input Image"), |
| outputs=[gr.Image(type="pil", label="Restored Image"), gr.Textbox(label="Log")], |
| title="GFPGAN v1.4 (ONNX) - Face Restoration", |
| description="Run GFPGAN face restoration using the ONNX model (no PyTorch dependencies)." |
| ) |
|
|
| if __name__ == "__main__": |
| iface.launch() |
