Skip to content

Instantly share code, notes, and snippets.

@karel1980
Created November 6, 2021 19:49
Show Gist options
  • Save karel1980/eb94625dbb09ff31b4c5b179be006302 to your computer and use it in GitHub Desktop.
Save karel1980/eb94625dbb09ff31b4c5b179be006302 to your computer and use it in GitHub Desktop.
Head stabilizer
import cv2
import mediapipe as mp
import numpy as np
import math
import sys
LEFT_EYE = 33
RIGHT_EYE = 263
def main():
cap = cv2.VideoCapture(0)
# Check if the webcam is opened correctly
if not cap.isOpened():
raise IOError("Cannot open webcam")
show_mesh = False
mp_face_mesh = mp.solutions.face_mesh
face_mesh = mp_face_mesh.FaceMesh(static_image_mode=False, max_num_faces=3)
mp_drawing = mp.solutions.drawing_utils
mp_drawing_styles = mp.solutions.drawing_styles
w, h = (640, 360)
fourcc = cv2.VideoWriter_fourcc(*'XVID')
print(sys.argv[1])
vid = cv2.VideoWriter(sys.argv[1],fourcc, 20.0, (640,360))
drawing_spec = mp_drawing.DrawingSpec(thickness=1, circle_radius=1)
font = cv2.FONT_HERSHEY_SIMPLEX
fontScale = 1
color = (0, 255, 0)
thickness = 2
while True:
ret, frame = cap.read()
frame = cv2.resize(frame, None, fx=0.5, fy=0.5, interpolation=cv2.INTER_AREA)
mesh_result = face_mesh.process(frame)
face_landmarks = mesh_result.multi_face_landmarks
if face_landmarks:
left_eye = np.array([face_landmarks[0].landmark[LEFT_EYE].x * frame.shape[1], face_landmarks[0].landmark[LEFT_EYE].y * frame.shape[0]])
right_eye = np.array([face_landmarks[0].landmark[RIGHT_EYE].x * frame.shape[1], face_landmarks[0].landmark[RIGHT_EYE].y * frame.shape[0]])
if show_mesh:
mp_drawing.draw_landmarks(
image=frame,
landmark_list=face_landmarks[0],
connections=mp_face_mesh.FACE_CONNECTIONS,
landmark_drawing_spec=drawing_spec,
connection_drawing_spec=drawing_spec)
target_left = np.array([frame.shape[1] * 0.45, frame.shape[0] * 0.33])
target_right = np.array([frame.shape[1] * 0.55, frame.shape[0] * 0.33])
frame = orient_image(frame, left_eye, right_eye, target_left, target_right, (frame.shape[1], frame.shape[0]))
frame = cv2.putText(frame, 'show_mesh = %s'%show_mesh, (10, frame.shape[0]-20), font, fontScale, color, thickness, cv2.LINE_AA)
# Why can't I just vide.write(frame)?
frame2 = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frame3 = cv2.cvtColor(frame2, cv2.COLOR_RGB2BGR)
vid.write(frame3)
cv2.imshow('Input', frame3)
c = cv2.waitKey(1)
if c == ord('q'):
break
if c == ord('m'):
show_mesh = not show_mesh
print("releasing video writer")
vid.release()
cap.release()
cv2.destroyAllWindows()
def orient_image(img, left_eye, right_eye, target_left_eye, target_right_eye, target_shape):
scale = 1.0 / (np.linalg.norm(right_eye - left_eye) / np.linalg.norm(target_right_eye - target_left_eye))
eyes_angle = - calculate_angle(target_right_eye-target_left_eye, right_eye - left_eye)
img_with_alpha = add_alpha_channel(img)
affine_transformation_matrix = cv2.getRotationMatrix2D(left_eye, eyes_angle * 180 / math.pi,
scale)
affine_transformation_matrix[:, 2] += target_left_eye - left_eye
result = np.zeros((target_shape[0], target_shape[1], 4))
# Note: cv2.BORDER_TRANSPARENT didn't work as expected,
# so manually set cv2.BORDER_CONSTANT with a transparent borderValue
result = cv2.warpAffine(img_with_alpha, affine_transformation_matrix, target_shape, result, flags=cv2.INTER_LINEAR,
borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0))
return result
def calculate_angle(a, b):
# calculates angle between 2 vectors in radians
left_unit = a / np.linalg.norm(a)
right_unit = b / np.linalg.norm(b)
radians = - np.arctan2(np.cross(left_unit, right_unit), np.dot(left_unit, right_unit))
return radians
def add_alpha_channel(img):
b_channel, g_channel, r_channel = cv2.split(img)
alpha_channel = np.ones(b_channel.shape, dtype=b_channel.dtype) * 255
return cv2.merge((b_channel, g_channel, r_channel, alpha_channel))
if __name__=="__main__":
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment