mirror of
https://github.com/Haoming02/sd-webui-old-photo-restoration.git
synced 2026-01-26 19:29:52 +00:00
154 lines
4.0 KiB
Python
154 lines
4.0 KiB
Python
# Copyright (c) Microsoft Corporation
|
|
|
|
from skimage.transform import SimilarityTransform
|
|
from matplotlib.patches import Rectangle
|
|
from skimage.transform import warp
|
|
from skimage import img_as_ubyte
|
|
import matplotlib.pyplot as plt
|
|
from PIL import Image
|
|
import numpy as np
|
|
import dlib
|
|
import os
|
|
|
|
|
|
def _standard_face_pts():
|
|
pts = (
|
|
np.array(
|
|
[196.0, 226.0, 316.0, 226.0, 256.0, 286.0, 220.0, 360.4, 292.0, 360.4],
|
|
np.float32,
|
|
)
|
|
/ 256.0
|
|
- 1.0
|
|
)
|
|
|
|
return np.reshape(pts, (5, 2))
|
|
|
|
|
|
def get_landmark(face_landmarks, id):
|
|
part = face_landmarks.part(id)
|
|
x = part.x
|
|
y = part.y
|
|
|
|
return (x, y)
|
|
|
|
|
|
def search(face_landmarks):
|
|
|
|
x1, y1 = get_landmark(face_landmarks, 36)
|
|
x2, y2 = get_landmark(face_landmarks, 39)
|
|
x3, y3 = get_landmark(face_landmarks, 42)
|
|
x4, y4 = get_landmark(face_landmarks, 45)
|
|
|
|
x_nose, y_nose = get_landmark(face_landmarks, 30)
|
|
|
|
x_left_mouth, y_left_mouth = get_landmark(face_landmarks, 48)
|
|
x_right_mouth, y_right_mouth = get_landmark(face_landmarks, 54)
|
|
|
|
x_left_eye = int((x1 + x2) / 2)
|
|
y_left_eye = int((y1 + y2) / 2)
|
|
x_right_eye = int((x3 + x4) / 2)
|
|
y_right_eye = int((y3 + y4) / 2)
|
|
|
|
results = np.array(
|
|
[
|
|
[x_left_eye, y_left_eye],
|
|
[x_right_eye, y_right_eye],
|
|
[x_nose, y_nose],
|
|
[x_left_mouth, y_left_mouth],
|
|
[x_right_mouth, y_right_mouth],
|
|
]
|
|
)
|
|
|
|
return results
|
|
|
|
|
|
def compute_transformation_matrix(img, landmark, normalize, target_face_scale=1.0):
|
|
|
|
std_pts = _standard_face_pts() # [-1,1]
|
|
target_pts = (std_pts * target_face_scale + 1) / 2 * 512.0
|
|
|
|
# print(target_pts)
|
|
|
|
h, w, c = img.shape
|
|
if normalize == True:
|
|
landmark[:, 0] = landmark[:, 0] / h * 2 - 1.0
|
|
landmark[:, 1] = landmark[:, 1] / w * 2 - 1.0
|
|
|
|
# print(landmark)
|
|
|
|
affine = SimilarityTransform()
|
|
|
|
affine.estimate(target_pts, landmark)
|
|
|
|
return affine.params
|
|
|
|
|
|
def show_detection(image, box, landmark):
|
|
plt.imshow(image)
|
|
print(box[2] - box[0])
|
|
plt.gca().add_patch(
|
|
Rectangle(
|
|
(box[1], box[0]),
|
|
box[2] - box[0],
|
|
box[3] - box[1],
|
|
linewidth=1,
|
|
edgecolor="r",
|
|
facecolor="none",
|
|
)
|
|
)
|
|
plt.scatter(landmark[0][0], landmark[0][1])
|
|
plt.scatter(landmark[1][0], landmark[1][1])
|
|
plt.scatter(landmark[2][0], landmark[2][1])
|
|
plt.scatter(landmark[3][0], landmark[3][1])
|
|
plt.scatter(landmark[4][0], landmark[4][1])
|
|
plt.show()
|
|
|
|
|
|
def affine2theta(affine, input_w, input_h, target_w, target_h):
|
|
# param = np.linalg.inv(affine)
|
|
param = affine
|
|
theta = np.zeros([2, 3])
|
|
theta[0, 0] = param[0, 0] * input_h / target_h
|
|
theta[0, 1] = param[0, 1] * input_w / target_h
|
|
theta[0, 2] = (
|
|
2 * param[0, 2] + param[0, 0] * input_h + param[0, 1] * input_w
|
|
) / target_h - 1
|
|
theta[1, 0] = param[1, 0] * input_h / target_w
|
|
theta[1, 1] = param[1, 1] * input_w / target_w
|
|
theta[1, 2] = (
|
|
2 * param[1, 2] + param[1, 0] * input_h + param[1, 1] * input_w
|
|
) / target_w - 1
|
|
return theta
|
|
|
|
|
|
def detect_hr(input_image: Image) -> list:
|
|
face_detector = dlib.get_frontal_face_detector()
|
|
detected_faces = []
|
|
|
|
landmark = os.path.join(
|
|
os.path.dirname(os.path.abspath(__file__)),
|
|
"shape_predictor_68_face_landmarks.dat",
|
|
)
|
|
landmark_locator = dlib.shape_predictor(landmark)
|
|
|
|
image = np.array(input_image)
|
|
faces = face_detector(image)
|
|
|
|
if len(faces) == 0:
|
|
return detected_faces
|
|
|
|
for face_id in range(len(faces)):
|
|
|
|
current_face = faces[face_id]
|
|
face_landmarks = landmark_locator(image, current_face)
|
|
current_fl = search(face_landmarks)
|
|
|
|
affine = compute_transformation_matrix(
|
|
image, current_fl, False, target_face_scale=1.3
|
|
)
|
|
|
|
aligned_face = warp(image, affine, output_shape=(512, 512, 3))
|
|
detected_faces.append(Image.fromarray(img_as_ubyte(aligned_face)))
|
|
|
|
return detected_faces
|