eye_heatmap.py

# -*- coding: utf-8 -*-
"""eye_heatmap.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1WfA5D8ZlcyOYZdtyc-mZvXm3LiqPMHmC
"""

import pandas as pd
import cv2
import numpy as np
from tqdm import tqdm
import argparse
import os
import json
import matplotlib.pyplot as plt

parser = argparse.ArgumentParser()
parser.add_argument("-t", "--threshold", type=float, default=0.5, help="the threshold value")
parser.add_argument("-img", "--image_path", type=str, default="./eye_tracker/what-are-stock-images.jpg", help="the path of image for which heatmap to be generated")
parser.add_argument("-txt", "--txt_path", type=str, default="./exp.txt", help="the path of txt file obtained from eye tracker")
# parser.add_argument("-txt", "--txt_path", type=str, default="./exp.txt", help="the path of csv file obtainde from eye_tracker")
args = parser.parse_args()
print(args)

def create_gaussian_filter(sizex,sizey, sigma=33, center=None,fix=1):
    x = np.arange(0, sizex, 1, float)
    y = np.arange(0, sizey, 1, float)
    x, y = np.meshgrid(x,y)
    
    if center is None:
        x0 = sizex // 2
        y0 = sizey // 2
    else:
        if np.isnan(center[0])==False and np.isnan(center[1])==False:            
            x0 = center[0]
            y0 = center[1]        
        else:
            return np.zeros((sizey,sizex))

    return fix*np.exp(-4*np.log(2) * ((x-x0)**2 + (y-y0)**2) / sigma**2)


# import pandas as pd
# import cv2
# import numpy as np
# from tqdm import tqdm

# def create_gaussian_filter(sizex,sizey, sigma=33, center=None,fix=1):
#     x = np.arange(0, sizex, 1, float)
#     y = np.arange(0, sizey, 1, float)
#     x, y = np.meshgrid(x,y)
    
#     if center is None:
#         x0 = sizex // 2
#         y0 = sizey // 2
#     else:
#         if np.isnan(center[0])==False and np.isnan(center[1])==False:            
#             x0 = center[0]
#             y0 = center[1]        
#         else:
#             return np.zeros((sizey,sizex))

#     return fix*np.exp(-4*np.log(2) * ((x-x0)**2 + (y-y0)**2) / sigma**2)


def create_heatmap(fix_arr, width, height, imgfile, alpha=0.5, threshold=10):
    
    heatmap = np.zeros((H,W), np.float32)
    for n_subject in tqdm(range(fix_arr.shape[0])):
        heatmap += create_gaussian_filter(W, H, 33, (fix_arr[n_subject,0],fix_arr[n_subject,1]),
                                fix_arr[n_subject,2])

    # Normalization
    heatmap = heatmap/np.amax(heatmap)
    heatmap = heatmap*255
    heatmap = heatmap.astype("uint8")
    
    if imgfile.any():
        # Resize heatmap to imgfile shape 
        h, w, _ = imgfile.shape
        heatmap = cv2.resize(heatmap, (w, h))
        heatmap_color = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
        
        # Create mask
        mask = np.where(heatmap<=threshold, 1, 0)
        mask = np.reshape(mask, (h, w, 1))
        mask = np.repeat(mask, 3, axis=2)

        # Marge images
        marge = imgfile*mask + heatmap_color*(1-mask)
        marge = marge.astype("uint8")
        marge = cv2.addWeighted(imgfile, 1-alpha, marge,alpha,0)
        return marge

    else:
        heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
        return heatmap

objectsList = []
print("Started Reading JSON file which contains multiple JSON document")
with open(args.txt_path) as f:
    for jsonObj in f:
        objectDict = json.loads(jsonObj)
        if(objectDict["category"]=="tracker"):
            objectsList.append(objectDict)

temp = pd.DataFrame()


print("Adding each JSON Decoded Object in Dataframe")
for object in objectsList:
    obj=object["values"]['frame']
    dict={
    
    'timestamp':[obj['timestamp']],

    'is_fix':obj['fix'],
    
    'gaze_smooth_x':obj['avg']['x'],
    'gaze_smooth_y':obj['avg']['y'],
    'gaze_raw_x':obj['raw']['x'],
    'gaze_raw_y':obj['raw']['y'],

    'lefteye_smooth_x':obj['lefteye']['avg']['x'],
    'lefteye_smooth_y':obj['lefteye']['avg']['y'],
    'lefteye_raw_x':obj['lefteye']['raw']['x'],
    'lefteye_raw_y':obj['lefteye']['raw']['y'],
    'lefteye_psize':obj['lefteye']['psize'],
    'lefteye_p_x':obj['lefteye']['pcenter']['x'],
    'lefteye_p_y':obj['lefteye']['pcenter']['y'],

    'righteye_smooth_x':obj['righteye']['avg']['x'],
    'righteye_smooth_y':obj['righteye']['avg']['y'],
    'righteye_raw_x':obj['righteye']['raw']['x'],
    'righteye_raw_y':obj['righteye']['raw']['y'],
    'righteye_psize':obj['righteye']['psize'],
    'righteye_p_x':obj['righteye']['pcenter']['x'],
    'righteye_p_y':obj['righteye']['pcenter']['y'],

    }
    df = pd.DataFrame(dict)
    temp=pd.concat([temp,df],ignore_index=True)
    
temp.to_csv('exp.csv')

# img1 = cv2.imread("./eye_tracker/what-are-stock-images.jpg")
# plt.imshow(img1)
# print(args.image_path)
if args.image_path is None:
    image = cv2.imread('sample.png') #Default image name
else:
    image = cv2.imread(args.image_path)
    # cv2.imshow(image)

# Read the dataset as a pandas dataframe

df = temp

# Filter out the rows where is_fix is False
df = df[df["is_fix"] == True]
fixations = []

df['timestamp'] = pd.to_datetime(df['timestamp'], format='%Y-%m-%d %H:%M:%S.%f')

# df.dtypes

# initialize the list of fixations
fixations = []

for index, row in df.iterrows():
    timestamp = row["timestamp"]
    gaze_x = row["gaze_smooth_x"]
    gaze_y = row["gaze_smooth_y"]
    righteye = row["righteye_psize"]
    lefteye = row["lefteye_psize"]
    if len(fixations) == 0:
        fixations.append([timestamp, gaze_x, gaze_y, pd.Timedelta(seconds=1), righteye, lefteye])
    else:
        # Get the last fixation in the list
        last_fixation = fixations[-1]

        # Calculate the distance between the current gaze point and the last fixation center
        distance = ((gaze_x - last_fixation[1]) ** 2 + (gaze_y - last_fixation[2]) ** 2) ** 0.5

        # Define a threshold for the maximum distance to consider two gaze points as part of the same fixation
        threshold = 5

        # If the distance is less than the threshold, update the last fixation with the current values
        if distance < threshold:
            # Update the fixation center as the weighted average of the gaze points
            last_fixation[1] = (last_fixation[1] * last_fixation[3].total_seconds() + gaze_x) / (last_fixation[3].total_seconds() + 1)
            last_fixation[2] = (last_fixation[2] * last_fixation[3].total_seconds() + gaze_y) / (last_fixation[3].total_seconds() + 1)

            # Update the fixation duration as the difference between the current timestamp and the first timestamp
            last_fixation[3] = timestamp - last_fixation[0]
        else:
            # Otherwise, append a new fixation with the current values
            fixations.append([timestamp, gaze_x, gaze_y, pd.Timedelta(seconds=1), righteye, lefteye])

# where n is the number of fixations and each row is (timestamp, x, y, duration)
fix_arr = np.array(fixations)

# fix_arr[0]

x = fix_arr[:, 1]
y = fix_arr[:, 2]
duration = fix_arr[:, 3]
# print(duration)
# Convert the duration column to float using total_seconds()
duration = np.array([d.total_seconds() for d in duration])

# Stack the x, y, and duration columns horizontally
fix_arr_float = np.column_stack((x, y, duration))

fix_arr_float -= fix_arr_float.min()
fix_arr_float /= fix_arr_float.max()
fix_arr_float

# img = cv2.imread('eye_tracker/what-are-stock-images.jpg')

# Generate toy fixation data
# when you use, replace here with your data
H, W, _ = image.shape

fix_arr_float[:,0] *= W
fix_arr_float[:,1] *= H

# Create heatmap
heatmap = create_heatmap(fix_arr_float, W, H, image, 0.7, 5)
cv2.imwrite("waldo_heatmap.png",heatmap)

# df.head()

'''
Pupil dilations code
'''



def plot_pupil_dilations(fixations):
    fig, ax = plt.subplots()
    time =[]
    diameter = []
    for i in fixations:
        time.append(i[0])
        diameter.append((i[4]+i[5])/2)
    ax.scatter(time, diameter, s=5)
    # plt.plot(diameter)
    ax.set_xlabel('Time')
    ax.set_ylabel('Pupil diameter (in mm)')
    ax.set_title('Pupil Dilations')
    
    plt.show()
    # ax2 = plt.sub

# # Example usage
# x = [1, 2, 3, 4]
# y = [5, 6, 7, 8]
# diameter = [2.5, 3.0, 3.5, 4.0]

plot_pupil_dilations(fixations)