Head stabilizer

cam.py

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()