add face_stabilizer.py

facefusion/face_stabilizer.py

@@ -0,0 +1,82 @@
+from collections import deque
+from time import monotonic
+from typing import List, Optional
+
+import numpy
+
+from facefusion.types import Face, FaceStabilizerStore, FaceStabilizerTrail, FaceStabilizerTrailEntry, Points, Resolution
+
+FACE_STABILIZER_STORE : FaceStabilizerStore = {}
+FACE_STABILIZER_TRAIL : FaceStabilizerTrail = {}
+FACE_STABILIZER_TRAIL_LENGTH = 64
+
+
+def clear_face_stabilizer() -> None:
+	FACE_STABILIZER_STORE.clear()
+	FACE_STABILIZER_TRAIL.clear()
+
+
+def get_face_stabilizer_trail(face_key : str) -> Optional[FaceStabilizerTrailEntry]:
+	return FACE_STABILIZER_TRAIL.get(face_key)
+
+
+def append_face_stabilizer_trail(face_key : str, raw_point : Points, smoothed_point : Points) -> None:
+	face_stabilizer_trail = FACE_STABILIZER_TRAIL.get(face_key, deque(maxlen = FACE_STABILIZER_TRAIL_LENGTH))
+	face_stabilizer_trail.append((raw_point.copy(), smoothed_point.copy()))
+	FACE_STABILIZER_TRAIL[face_key] = face_stabilizer_trail
+
+
+def stabilize_faces(faces : List[Face], resolution : Resolution, smoothness : float) -> List[Face]:
+	stabilized_faces = []
+
+	if smoothness > 0:
+		timestamp = monotonic()
+
+		for index, face in enumerate(faces):
+			stabilized_faces.append(stabilize_face(face, str(index), timestamp, resolution, smoothness))
+		return stabilized_faces
+
+	return faces
+
+
+def stabilize_face(face : Face, face_key : str, timestamp : float, resolution : Resolution, smoothness : float) -> Face:
+	width, height = resolution
+	bounding_box_scale = numpy.array([ width, height, width, height ])
+	landmark_scale = numpy.array([ width, height ])
+	bounding_box_min_cutoff = 0.1 / smoothness
+	bounding_box_beta = 40.0 * (1 - smoothness)
+	landmark_min_cutoff = 0.05 / smoothness
+	landmark_beta = 80.0 * (1 - smoothness)
+	landmark_5 = face.landmark_set.get('5') / landmark_scale
+	landmark_5_smoothed = apply_one_euro_filter(face_key + '.landmark_5', landmark_5, timestamp, landmark_min_cutoff, landmark_beta)
+
+	bounding_box = apply_one_euro_filter(face_key + '.bounding_box', face.bounding_box / bounding_box_scale, timestamp, bounding_box_min_cutoff, bounding_box_beta) * bounding_box_scale
+	landmark_set = face.landmark_set.copy()
+	landmark_set['5'] = landmark_5_smoothed * landmark_scale
+	landmark_set['5/68'] = apply_one_euro_filter(face_key + '.landmark_5_68', face.landmark_set.get('5/68') / landmark_scale, timestamp, landmark_min_cutoff, landmark_beta) * landmark_scale
+	landmark_set['68'] = apply_one_euro_filter(face_key + '.landmark_68', face.landmark_set.get('68') / landmark_scale, timestamp, landmark_min_cutoff, landmark_beta) * landmark_scale
+	landmark_set['68/5'] = apply_one_euro_filter(face_key + '.landmark_68_5', face.landmark_set.get('68/5') / landmark_scale, timestamp, landmark_min_cutoff, landmark_beta) * landmark_scale
+	append_face_stabilizer_trail(face_key, landmark_5[2], landmark_5_smoothed[2])
+
+	return face._replace(bounding_box = bounding_box, landmark_set = landmark_set)
+
+
+def apply_one_euro_filter(filter_key : str, points : Points, timestamp : float, min_cutoff : float, beta : float) -> Points:
+	points = points.astype(numpy.float64)
+
+	if filter_key in FACE_STABILIZER_STORE:
+		points_prev, velocity_prev, timestamp_prev = FACE_STABILIZER_STORE.get(filter_key)
+		timestamp_delta = timestamp - timestamp_prev
+
+		if timestamp_delta > 0:
+			velocity_alpha = 1.0 / (1.0 + 1.0 / (2 * numpy.pi * timestamp_delta))
+			velocity = velocity_alpha * (points - points_prev) / timestamp_delta + (1 - velocity_alpha) * velocity_prev
+			cutoff = min_cutoff + beta * numpy.abs(velocity)
+			points_alpha = 1.0 / (1.0 + 1.0 / (2 * numpy.pi * cutoff * timestamp_delta))
+			points = points_alpha * points + (1 - points_alpha) * points_prev
+			FACE_STABILIZER_STORE[filter_key] = (points, velocity, timestamp)
+			return points
+		return points_prev
+
+	FACE_STABILIZER_STORE[filter_key] = (points, numpy.zeros_like(points), timestamp)
+	return points