app.py

import gradio as gr
import pandas as pd
import cv2
import mediapipe as mp
import os
from statistics import mean
import numpy as np
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
from mediapipe.framework.formats import landmark_pb2
from mediapipe import solutions


import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

# Record video
# Save video?
# Break video into images
# Run facemesh on all images and save locations
# Run exterme locations
# Run analysis on those compare to the first frame

# Create a FaceLandmarker object.
base_options = python.BaseOptions(model_asset_path='face_landmarker_v2_with_blendshapes.task')
options = vision.FaceLandmarkerOptions(base_options=base_options,
                                       output_face_blendshapes=True,
                                       output_facial_transformation_matrixes=True,
                                       num_faces=1)
detector = vision.FaceLandmarker.create_from_options(options)


global pupilLocation
pupilLocation = pd.DataFrame()
pupil_sizes =  []
ExteremeDistanceLeftEye = pd.DataFrame()
ExteremeDistanceRightEye = pd.DataFrame()

def video_identity(video):
    return video

# To do
# 1. Filter out closed eye frames
# 2. Smooth persuit from video POC

def isEyeOpen(image):
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    hist = cv2.calcHist([image], [0], None, [256], [0, 256])
    colors = np.where(hist > 10)

    #print ("colors=", np.mean(colors) )
    if np.mean(colors) > 15:
        return True
    else:
        return False



#demo = gr.Interface(video_identity,
#                    gr.Video(shape = (1000,1000), source="webcam"),
#                    "playable_video")

def findIrisInFrame(image, counter):
    global pupilLocation, pupil_sizes
    #pupilLocation = pd.DataFrame()  # Make sure it is empty

    image = mp.Image.create_from_file("image.jpg")

    # STEP 4: Detect face landmarks from the input image.
    detection_result = detector.detect(image)

    # STEP 5: Process the detection result. In this case, visualize it.
    #annotated_image = draw_landmarks_on_image(image.numpy_view(), detection_result)

    annotated_image = image.numpy_view().copy()

    face_landmarks_list = detection_result.face_landmarks

    '''
    # Loop through the detected faces to visualize.
    for idx in range(len(face_landmarks_list)):
        face_landmarks = face_landmarks_list[idx]

        # Draw the face landmarks.
        face_landmarks_proto = landmark_pb2.NormalizedLandmarkList()
        face_landmarks_proto.landmark.extend([
            landmark_pb2.NormalizedLandmark(x=landmark.x, y=landmark.y, z=landmark.z) for landmark in face_landmarks
        ])
        solutions.drawing_utils.draw_landmarks(
            image=annotated_image,
            landmark_list=face_landmarks_proto,
            connections=mp.solutions.face_mesh.FACEMESH_IRISES,
            landmark_drawing_spec=None,
            connection_drawing_spec=mp.solutions.drawing_styles
            .get_default_face_mesh_iris_connections_style())
    '''
    #FACEMESH_LEFT_IRIS = (474, 475, 476, 477)
    #FACEMESH_RIGHT_IRIS = (469, 470, 471, 472)
    #(lm_left_iris.x, lm_left_iris.y, lm_left_iris.z) = face_landmarks.landmark[468]
    #(lm_right_iris.x, lm_right_iris.y, lm_right_iris.z) = face_landmarks.landmark[473]

    # Draw the face landmarks.
    face_landmarks = face_landmarks_list[0]
    face_landmarks_proto = landmark_pb2.NormalizedLandmarkList()
    face_landmarks_proto.landmark.extend([
        landmark_pb2.NormalizedLandmark(x=landmark.x, y=landmark.y, z=landmark.z) for landmark in face_landmarks
    ])

    height, width, _ = annotated_image.shape
    nose = [int(face_landmarks_proto.landmark[168].x * width), int(face_landmarks_proto.landmark[168].y * height)]
    cv2.circle(annotated_image, (nose[0], nose[1]), 3, (0, 0, 255), -1)

    leftIrisPoints = [474, 475, 476, 477]
    rightIrisPoints = [469, 470, 471, 472]
    # right, top, left, bottom

    left_iris = []
    for p in leftIrisPoints:
                    point = [int(face_landmarks_proto.landmark[p].x * width), int(face_landmarks_proto.landmark[p].y * height)]
                    left_iris.append(point)
                    cv2.circle(annotated_image, point, 1, (255, 0, 255), 1)

    right_iris = []
    for p in rightIrisPoints:
                    point = [int(face_landmarks_proto.landmark[p].x * width), int(face_landmarks_proto.landmark[p].y * height)]
                    right_iris.append(point)
                    cv2.circle(annotated_image, point, 1, (255, 0, 255), 1)

    leftIris_leftside = (int(left_iris[2][0]), int(left_iris[2][1]))
    leftIris_rightside = (int(left_iris[0][0]), int(left_iris[0][1]))
    leftIris_top = (int(left_iris[1][0]), int(left_iris[1][1]))
    leftIris_bottom = (int(left_iris[3][0]), int(left_iris[3][1]))
    rightIris_leftside = (int(right_iris[2][0]), int(right_iris[2][1]))
    rightIris_rightside = (int(right_iris[0][0]), int(right_iris[0][1]))
    rightIris_top = (int(right_iris[1][0]), int(right_iris[1][1]))
    rightIris_bottom = (int(right_iris[3][0]), int(right_iris[3][1]))

    cv2.imwrite("Images/post-image-%d.jpg" % counter, annotated_image)  # save frame as JPEG file

    '''
                sizeIncrease = 0
                leftEye = annotated_image[leftIris_top[1] - sizeIncrease: leftIris_bottom[1] + sizeIncrease,
                          leftIris_leftside[0] - sizeIncrease: leftIris_rightside[0] + sizeIncrease]
                leftEyeOpen = isEyeOpen(leftEye)

                rightEye = annotated_image[rightIris_top[1] - sizeIncrease: rightIris_bottom[1] + sizeIncrease,
                           rightIris_leftside[0] - sizeIncrease: rightIris_rightside[0] + sizeIncrease]
                rightEyeOpen = isEyeOpen(rightEye)


                cv2.circle(image,
                           (int((leftIris_leftside[0] + leftIris_rightside[0]) / 2),
                            int((leftIris_top[1] + leftIris_bottom[1]) / 2)),
                           # int(abs(leftIris_leftside[0] - leftIris_rightside[0])/2),
                           1,
                           (0, 255, 255), 1)

                cv2.circle(image,
                           (int((rightIris_leftside[0] + rightIris_rightside[0]) / 2),
                            int((rightIris_top[1] + rightIris_bottom[1]) / 2)),
                           # int(abs(rightIris_leftside[0] - rightIris_rightside[0])/2),
                           1,
                           (0, 255, 255), 1)

                cv2.putText(image, str(counter),
                            (rightIris_leftside[0] - 100, leftIris_top[1] - 10), cv2.FONT_HERSHEY_SIMPLEX,
                            1, (255, 0, 0), 1, cv2.LINE_AA)


    '''
    pupil_sizes.append(abs(leftIris_leftside[0] - leftIris_rightside[0]))
    pupil_sizes.append(abs(rightIris_leftside[0] - rightIris_rightside[0]))

    name = "frame%d.jpg" % counter
    newRow = pd.Series([name,
                                    leftIris_leftside[0]   - nose[0],
                                    leftIris_top[1]        - nose[1],
                                    leftIris_rightside[0]  - nose[0],
                                    leftIris_bottom[1]     - nose[1],
                                    rightIris_leftside[0]  - nose[0],
                                    rightIris_top[1]       - nose[1],
                                    rightIris_rightside[0] - nose[0],
                                    rightIris_bottom[1]    - nose[1]
                                    ])
    newRow = newRow.to_frame().T
                #if (leftEyeOpen & rightEyeOpen):
    pupilLocation = pd.concat([pupilLocation, newRow], axis=0, ignore_index=True)
                #else:
                #    print("Ignored frame ", counter, "." , leftEyeOpen , rightEyeOpen)

    return newRow


def handleVideo(input_video):
    global ExteremeDistanceLeftEye, ExteremeDistanceRightEye, pupilLocation, pupil_sizes

    pupilLocation = pd.DataFrame() # Make sure it is empty to begin with
    pupil_sizes =  []
    vidcap = cv2.VideoCapture(input_video)
    success, image = vidcap.read()
    fps = vidcap.get(cv2.CAP_PROP_FPS)
    frame_count = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
    width = vidcap.get(cv2.CAP_PROP_FRAME_WIDTH)  # float `width`
    height = vidcap.get(cv2.CAP_PROP_FRAME_HEIGHT)  # float `height`
    print('FPS =', fps)
    print('Frame count =', frame_count)
    print('Resolution =', width, ' X ', height)

    count = 0
    if not os.path.exists('Images'):
        os.makedirs('Images')
    #os.chdir('Images')

    # Slide video into frames and find iris in each frame
    while success:
        cv2.imwrite("image.jpg", image)  # save frame as JPEG file
        print("Image#", count)
        image = mp.Image.create_from_file("image.jpg")
        findIrisInFrame(image, count)

        #cv2.imwrite("Images/frame%d.jpg" % count, image)  # save frame as JPEG file
        count += 1
        success, image = vidcap.read()


    # Go over all the pupils. If pupil is too small expand it in all directions
    # Find average pupil size
    pupil_average = 11.7
    if (len(pupil_sizes) > 100):
        for i in range(10):
           pupil_sizes.remove(max(pupil_sizes))
           pupil_sizes.remove(min(pupil_sizes))
        pupil_average = mean(pupil_sizes)  # this should be 11.7 mm
    print("pupil_average=", pupil_average)

    # pupil size need to be kept constant in all pictures
    # we find the center of the current pupil and make a circle around it in the size we need

    for index, row in pupilLocation.iterrows():
        currentLeftSize = abs(row[1] - row[3])
        diffFromLeftAverage = pupil_average - currentLeftSize
        currentRightSize = abs(row[5] - row[7])
        diffFromAverage = pupil_average - currentRightSize

        #print("(frame#", index, ") left: ", row[1], " right: ", row[3])
        #if diffFromAverage > 2:
        if (currentRightSize - currentLeftSize > 200):
            print("Fixed Left pupil")
            row[1] = int(row[1] - diffFromAverage / 2)
            row[2] = int(row[2] - diffFromAverage / 2)
            row[3] = int(row[3] + diffFromAverage / 2)
            row[4] = int(row[4] + diffFromAverage / 2)

        if (currentLeftSize - currentRightSize > 200):
            print("Fixed Right pupil")
            row[5] = int(row[5] - diffFromAverage / 2)
            row[6] = int(row[6] - diffFromAverage / 2)
            row[7] = int(row[7] + diffFromAverage / 2)
            row[8] = int(row[8] + diffFromAverage / 2)


    print("file counter=", count)
    # Convert pupilLocation to pupilDiff
    pupilDiff = pupilLocation.copy()
    pupilDiff = pupilDiff.drop(pupilDiff.columns[0], axis=1)  # Remove file name
    for i in range(pupilDiff.shape[0] - 1):  # Calculate deltas
        pupilDiff.loc[i + 1] = (pupilDiff.loc[i + 1] - pupilDiff.loc[0])
    pupilDiff = pupilDiff.drop(0, axis=0) # Remove the first row

    # Find extreme iris locations (images and measurements)

    #pupilDiff.columns = ['LL', 'LT', 'LR', 'LB', 'RR', 'RT', 'RR', 'RB']

    # Take just the relevant libmus
    #if [1] is positive, take it, otherwise take 3

    #pupilDiff[9]  = [pupilDiff[1] if x >= 0 else pupilDiff[3] for x in pupilDiff[1]]

    pupilDiff[9] = 0
    pupilDiff[10] = 0

    pupilDiff[9]  = np.where(pupilDiff[1] >= 0, pupilDiff[1], pupilDiff[3])
                             #pupilDiff[[1,3]].max(axis=1), pupilDiff[[1,3]].min(axis=1))
    pupilDiff[10] = np.where(pupilDiff[1] >= 0, pupilDiff[5], pupilDiff[7])
                             #pupilDiff[[5,7]].max(axis=1), pupilDiff[[5,7]].min(axis=1))
    pupilDiff[11] = np.where(pupilDiff[2] >= 0, pupilDiff[2], pupilDiff[4])
                             #pupilDiff[[2,4]].max(axis=1), pupilDiff[[2,4]].min(axis=1))
    pupilDiff[12] = np.where(pupilDiff[2] >= 0, pupilDiff[6], pupilDiff[8])
                             #pupilDiff[[6,8]].max(axis=1), pupilDiff[[6,8]].min(axis=1))
    print(pupilDiff[[1,3,5,7,9,10]])

    # slope
    x1 = (pupilLocation[1] + pupilLocation[3]) / 2
    y1 = (pupilLocation[2] + pupilLocation[4]) / 2
    x2 = (pupilLocation[5] + pupilLocation[7]) / 2
    y2 = (pupilLocation[6] + pupilLocation[8]) / 2
    pupilDiff[13] = ((y2 - y1) / (0.001 + x2 - x1))


    pupilDiff.to_csv('pupil_diff.csv')


    pixels = 11.7 / pupil_average
    print("pixels (In MM) = ", pixels)


    pupilDiff = round(pupilDiff * pixels,3)

    fig1 = plt.figure()
    plt.plot(pupilDiff[[9,10]]) #1,3,5,7
    plt.title("Pupil movement X axis")
    plt.ylabel("MM of movement")
    plt.xlabel("Frame")
    plt.ylim(-10, 10)
    plt.legend(['Left', 'Right']) #'LL', 'LR', 'RL', 'RR'])

    fig2 = plt.figure()
    plt.plot(pupilDiff[[11,12]]) #, df[countries].to_numpy())
    plt.ylim(-10, 10)
    plt.title("Pupil movement Y axis")
    plt.ylabel("MM of movement")
    plt.xlabel("Frame")
    plt.legend(['Left', 'Right'])


    # Left eye
    LeftEyeLookingRight = pd.to_numeric(pupilDiff[1]).idxmax()
    LeftEyeLookingDown  = pd.to_numeric(pupilDiff[2]).idxmax()
    LeftEyeLookingLeft  = pd.to_numeric(pupilDiff[3]).idxmin()
    LeftEyeLookingUp    = pd.to_numeric(pupilDiff[4]).idxmin()


    # Right eye
    RightEyeLookingRight = pd.to_numeric(pupilDiff[5]).idxmax()
    RightEyeLookingDown  = pd.to_numeric(pupilDiff[6]).idxmax()
    RightEyeLookingLeft  = pd.to_numeric(pupilDiff[7]).idxmin()
    RightEyeLookingUp    = pd.to_numeric(pupilDiff[8]).idxmin()


    print("Left eye images = ", LeftEyeLookingRight, LeftEyeLookingDown, LeftEyeLookingLeft, LeftEyeLookingUp)
    print("Right eye images = ", RightEyeLookingRight, RightEyeLookingDown, RightEyeLookingLeft, RightEyeLookingUp)

    ExtermeImageLeftEye = list([cv2.cvtColor(cv2.imread("Images/post-image-%d.jpg" % LeftEyeLookingRight), cv2.COLOR_BGR2RGB),
                                cv2.cvtColor(cv2.imread("Images/post-image-%d.jpg" % LeftEyeLookingLeft), cv2.COLOR_BGR2RGB),
                                cv2.cvtColor(cv2.imread("Images/post-image-%d.jpg" % LeftEyeLookingUp), cv2.COLOR_BGR2RGB),
                                cv2.cvtColor(cv2.imread("Images/post-image-%d.jpg" % LeftEyeLookingDown), cv2.COLOR_BGR2RGB)])

    ExtermeImageRightEye = list([cv2.cvtColor(cv2.imread("Images/post-image-%d.jpg" % RightEyeLookingRight), cv2.COLOR_BGR2RGB),
                                 cv2.cvtColor(cv2.imread("Images/post-image-%d.jpg" % RightEyeLookingLeft), cv2.COLOR_BGR2RGB),
                                 cv2.cvtColor(cv2.imread("Images/post-image-%d.jpg" % RightEyeLookingUp), cv2.COLOR_BGR2RGB),
                                 cv2.cvtColor(cv2.imread("Images/post-image-%d.jpg" % RightEyeLookingDown), cv2.COLOR_BGR2RGB)])


    # return the distances


    d = { 'direction': ['Right', 'Left', 'Up', 'Down'] ,
          'mm' : [abs(round(pd.to_numeric(pupilDiff[1]).max(),1)),
                  abs(round(pd.to_numeric(pupilDiff[3]).min(),1)),
                  abs(round(pd.to_numeric(pupilDiff[4]).min(),1)),
                  abs(round(pd.to_numeric(pupilDiff[2]).max(),1))
                  ]}
    ExteremeDistanceLeftEye = pd.DataFrame(data=d)

    d = {'direction': ['Right', 'Left', 'Up', 'Down'],
         'mm': [abs(round(pd.to_numeric(pupilDiff[5]).max(), 1)),
                abs(round(pd.to_numeric(pupilDiff[7]).min(), 1)),
                abs(round(pd.to_numeric(pupilDiff[8]).min(), 1)),
                abs(round(pd.to_numeric(pupilDiff[6]).max(), 1))
                ]}
    ExteremeDistanceRightEye = pd.DataFrame(data=d)



    print() #.idxmax(axis=0))
    # Upmost buttom limbus
    #
    return ExteremeDistanceLeftEye, ExteremeDistanceRightEye, ExtermeImageLeftEye, ExtermeImageRightEye, fig1, fig2 # lines


with gr.Blocks() as demo:
            gr.Markdown(
            """
            # Range of Motion Video Analysis
            Capture a video of the following looks: stright, left, right, up & down
            """)
            video1 = gr.Video()#source="webcam")

            b = gr.Button("Analyze Video")

            gr.Markdown(
                """
                # Left eye results (in mm):
                """)
            LeftEyeGallery = gr.Gallery(
                label="Left eye", show_label=False, elem_id="left_eye_gallery"
            ).style(grid=[4], height="auto")
            movementDataLeft = gr.Dataframe(ExteremeDistanceLeftEye)


            gr.Markdown(
                """
                # Right eye results (in mm):
                """)
            RightEyeGallery = gr.Gallery(
                label="Right eye", show_label=False, elem_id="right_eye_gallery"
            ).style(grid=[4], height="auto")
            movementDataRight = gr.Dataframe(ExteremeDistanceRightEye)

            plot1 = gr.Plot(label="Plot1")

            plot2 = gr.Plot(label="Plot2")

            out = [movementDataLeft, movementDataRight, LeftEyeGallery, RightEyeGallery, plot1, plot2]
            b.click(fn=handleVideo, inputs=video1, outputs=out)

demo.launch()