demo.py

#!/usr/bin/env python

import datetime
import logging
import pathlib
from typing import Optional

import cv2
import numpy as np
import yacs.config

from gaze_estimation import GazeEstimationMethod, GazeEstimator
from gaze_estimation.gaze_estimator.common import (Face, FacePartsName,
                                                   Visualizer)
from gaze_estimation.utils import load_config
from gaze_estimation.gaze_estimator.common.visualizer import * 
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

global pitch

class Demo:
    QUIT_KEYS = {27, ord('q')}

    def __init__(self, config: yacs.config.CfgNode):
        self.config = config
        print(self.config)
        self.gaze_estimator = GazeEstimator(config)
        self.visualizer = Visualizer(self.gaze_estimator.camera)

        self.cap = self._create_capture()
        self.output_dir = self._create_output_dir()
        self.writer = self._create_video_writer()

        self.stop = False
        self.show_bbox = self.config.demo.show_bbox
        self.show_head_pose = self.config.demo.show_head_pose
        self.show_landmarks = self.config.demo.show_landmarks
        self.show_normalized_image = self.config.demo.show_normalized_image
        self.show_template_model = self.config.demo.show_template_model

    def run(self) -> None:
        while True:
            thickness =1
            fontScale = 1
            color=(255,255,255)
            print("------------------while start")
            if self.config.demo.display_on_screen:
                self._wait_key()
                if self.stop:
                    print("---------")
                    break

            ok, frame = self.cap.read()
            if not ok:
                print("-------------------------")
                break

            undistorted = cv2.undistort(
                frame, self.gaze_estimator.camera.camera_matrix,
                self.gaze_estimator.camera.dist_coefficients)
            
            self.visualizer.set_image(frame.copy())
            
            faces = self.gaze_estimator.detect_faces(undistorted)
            print("3") # FACE is not detected thats why we are facing problem in running webcam
            for face in faces:
                self.gaze_estimator.estimate_gaze(undistorted, face)
                print("4")
                self._draw_face_bbox(face)
                self._draw_head_pose(face)
                self._draw_landmarks(face)
                self._draw_face_template_model(face)
                self._draw_gaze_vector(face)
                self._display_normalized_image(face)

                #print("-----------------------------------",self.pitch,self.yaw,self.roll,self.pitch_e,self.yaw_e)
                #print("The Pitch  is: {:.2f}".format(self.pitch))
                cv2.putText(self.visualizer.image, "eye_pitch:{:.2f},eye_yaw:{:.2f}".format(self.pitch_eg,self.yaw_eg), (0,20), cv2.FONT_HERSHEY_PLAIN, fontScale, color, thickness)
                cv2.putText(self.visualizer.image, "pitch:{:.2f},yaw:{:.2f},roll:{:.2f}".format(self.pitch,self.yaw,self.roll), (0,20), cv2.FONT_HERSHEY_PLAIN, fontScale, color, thickness)
                self.image_frame = self.visualizer.image
                cv2.imshow('frame', self.image_frame)


            if self.config.demo.use_camera:
                self.visualizer.image = self.visualizer.image[:, ::-1]
            if self.writer:
                self.writer.write(self.visualizer.image)
            if self.config.demo.display_on_screen:
                thickness =1
                fontScale = 1
                color=(255,255,255)
                #print("-----------------------------------",self.pitch,self.yaw)#,self.roll)
                #cv2.putText(self.visualizer.image, "pitch:{:.2f}".format(self.pitch), (0,20), cv2.FONT_HERSHEY_PLAIN, fontScale, color, thickness)

                #cv2.imshow('frame', self.visualizer.image)
            print("WHILE COMPLETED")
        self.cap.release()
        if self.writer:
            self.writer.release()

    def _create_capture(self) -> cv2.VideoCapture:
        if self.config.demo.use_camera:
            cap = cv2.VideoCapture(0)
        elif self.config.demo.video_path:
            cap = cv2.VideoCapture(self.config.demo.video_path)
        else:
            raise ValueError
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.gaze_estimator.camera.width)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.gaze_estimator.camera.height)
        return cap

    def _create_output_dir(self) -> Optional[pathlib.Path]:
        if not self.config.demo.output_dir:
            return
        output_dir = pathlib.Path(self.config.demo.output_dir)
        output_dir.mkdir(exist_ok=True, parents=True)
        return output_dir

    @staticmethod
    def _create_timestamp() -> str:
        dt = datetime.datetime.now()
        return dt.strftime('%Y%m%d_%H%M%S')

    def _create_video_writer(self) -> Optional[cv2.VideoWriter]:
        if not self.output_dir:
            return None
        ext = self.config.demo.output_file_extension
        if ext == 'mp4':
            fourcc = cv2.VideoWriter_fourcc(*'H264')
        elif ext == 'avi':
            fourcc = cv2.VideoWriter_fourcc(*'PIM1')
        else:
            raise ValueError
        output_path = self.output_dir / f'{self._create_timestamp()}.{ext}'
        writer = cv2.VideoWriter(output_path.as_posix(), fourcc, 30,
                                 (self.gaze_estimator.camera.width,
                                  self.gaze_estimator.camera.height))
        if writer is None:
            raise RuntimeError
        return writer

    def _wait_key(self) -> None:
        key = cv2.waitKey(self.config.demo.wait_time) & 0xff
        if key in self.QUIT_KEYS:
            self.stop = True
        elif key == ord('b'):
            self.show_bbox = not self.show_bbox
        elif key == ord('l'):
            self.show_landmarks = not self.show_landmarks
        elif key == ord('h'):
            self.show_head_pose = not self.show_head_pose
        elif key == ord('n'):
            self.show_normalized_image = not self.show_normalized_image
        elif key == ord('t'):
            self.show_template_model = not self.show_template_model

    def _draw_face_bbox(self, face: Face) -> None:
        if not self.show_bbox:
            return
        #print("==============",self.visualizer.draw_bbox(face.bbox).bbox[0])
        self.visualizer.draw_bbox(face.bbox)
        print("=====FACE RECTANGLE BOX=========",self.visualizer.bbox[0],"---------------",self.visualizer.bbox[1])
        
        self.co_ordinates=self.visualizer.bbox[0],self.visualizer.bbox[1]

    def _draw_head_pose(self, face: Face) -> None:
        if not self.show_head_pose:
            return
        # Draw the axes of the model coordinate system
        length = self.config.demo.head_pose_axis_length
        self.visualizer.draw_model_axes(face, length, lw=2)

        euler_angles = face.head_pose_rot.as_euler('XYZ', degrees=True)
        
        self.pitch, self.yaw, self.roll = face.change_coordinate_system(euler_angles)
        # logger.info(f'[head] pitch: {self.pitch:.2f}, yaw: {self.yaw:.2f}, '
        #             f'roll: {self.roll:.2f}, distance: {face.distance:.2f}')

    def _draw_landmarks(self, face: Face) -> None:
        if not self.show_landmarks:
            return
        self.visualizer.draw_points(face.landmarks,
                                    color=(0, 255, 255),
                                    size=1)

    def _draw_face_template_model(self, face: Face) -> None:
        if not self.show_template_model:
            return
        self.visualizer.draw_3d_points(face.model3d,
                                       color=(255, 0, 525),
                                       size=1)

    def _display_normalized_image(self, face: Face) -> None:
        if not self.config.demo.display_on_screen:
            return
        if not self.show_normalized_image:
            return
        if self.config.mode == GazeEstimationMethod.MPIIGaze.name:
            reye = face.reye.normalized_image
            leye = face.leye.normalized_image
            normalized = np.hstack([reye, leye])
        elif self.config.mode == GazeEstimationMethod.MPIIFaceGaze.name:
            normalized = face.normalized_image
        else: #eth gaze
            normalized = face.normalized_image
            #raise ValueError
        if self.config.demo.use_camera:
            normalized = normalized[:, ::-1]
        cv2.imshow('normalized', normalized)

    def _draw_gaze_vector(self, face: Face) -> None:
        length = self.config.demo.gaze_visualization_length
        if self.config.mode == GazeEstimationMethod.MPIIGaze.name:
            for key in [FacePartsName.REYE, FacePartsName.LEYE]:
                eye = getattr(face, key.name.lower())
                self.visualizer.draw_3d_line(
                    eye.center, eye.center + length * eye.gaze_vector)
                self.pitch_eg, self.yaw_eg = np.rad2deg(eye.vector_to_angle(eye.gaze_vector))
                # logger.info(
                #     f'[{key.name.lower()}] pitch: {self.pitch_e:.2f}, yaw: {self.yaw_e:.2f}')
        elif self.config.mode == GazeEstimationMethod.MPIIFaceGaze.name:
            self.visualizer.draw_3d_line(
                face.center, face.center + length * face.gaze_vector)
            self.pitch_eg, self.yaw_eg = np.rad2deg(face.vector_to_angle(face.gaze_vector))
            #logger.info(f'[face] pitch: {pitch:.2f}, yaw: {yaw:.2f}')
        else: # eth gaze
            self.visualizer.draw_3d_line(
                face.center, face.center + length * face.gaze_vector)
            self.pitch_eg, self.yaw_eg = np.rad2deg(face.vector_to_angle(face.gaze_vector))
            #logger.info(f'[face] pitch: {self.pitch:.2f}, yaw: {self.yaw:.2f}')
            #raise ValueError


def main():
    config = load_config()
    demo = Demo(config)
    demo.run()


if __name__ == '__main__':
    main()