pupil_preprocessing/glue_script.py at main · beresluca/pupil_preprocessing · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
#!/usr/bin/env python3

"""
Script glueing together functions from pupil_preprocessing.py for a given use case / testing.
"""
import csv
import os.path
import sys
from matplotlib import pyplot as plt
from matplotlib import use as mplot_use
#sys.path.append('/home/adamb/PycharmProjects/pupil_processing/pupillabs/')
#sys.path.append('C:\\Users\\Luca\\PycharmProjects\\pupil\\')
sys.path.append('/media/lucab/data_hdd/lucab/PycharmProjects/pupil_preprocessing/preprocessing_code/1_to_common_format/pupillabs')
import pupil_preprocessing as pupil


#mplot_use('TkAgg')

#data_file = 'D:\\pupil_sampleVideos\\pair54_Gondor_freeConv\\tsv\\pupil_data.tsv'
data_file = '/media/lucab/HV620S/pupil_sampleVideos/pair48_Mordor_freeConv/tsv/pupil_data.tsv'
output_folder = '/media/lucab/HV620S/pupil_sampleVideos/pair48_Mordor_freeConv/tsv'

# load data, only 3d pupil estimation values, into pandas dataframe
df = pupil.read_pupiltsv_3d(data_file)
pupil.nancheck(df)

# confidence-based filtering of data
conf_df = pupil.confidence_filter(df, cutoff=0.7)
pupil.nancheck(conf_df)

# valid range filtering
vr_df, ind_below, ind_above = pupil.valid_range_filter(conf_df)
pupil.nancheck(vr_df)

# SD-based filtering of data
sd_df, median_thresh, indices_below, indices_above = pupil.sd_filter(vr_df, sd_const=2)
pupil.nancheck(sd_df)

# Filtering with a MAD-based threshold on the speed of diameter changes
mad_df, mad_threshold, mad_indices = pupil.dilation_speed_filter(sd_df, mad_const=3)
pupil.nancheck(mad_df)

# Filtering with MAD-based threshold for a second time
mad2_df, mad2_threshold, mad2_indices = pupil.dilation_speed_filter(mad_df, mad_const=3)
pupil.nancheck(mad2_df)

# Interpolation
interp_df = pupil.interpolate(mad2_df, max_loner_length=2)
pupil.nancheck(interp_df)

# Median filtering
# med_df = pupil.median_filter(interp_df, order=5)
# pupil.nancheck(med_df)

# Low-pass filter
lp_df = pupil.lowpass_filter_butter_zerophase(interp_df, cutoff=10, fs=120, order=5)
pupil.nancheck(lp_df)


# Plots
time = df.timestamp.to_numpy()-df.timestamp[1]
plt.scatter(time, df.diameter.to_numpy(), c="limegreen", s=3, label="Raw")  # raw data
plt.scatter(time, conf_df.diameter.to_numpy(), c="forestgreen", s=3, label="Confidence thresh")
plt.scatter(time, sd_df.diameter.to_numpy(), c="indianred", s=3, label="SD thresh")
plt.scatter(time, mad_df.diameter.to_numpy(), c="gray", s=3, label="DS thresh")
plt.scatter(time, mad2_df.diameter.to_numpy(), c="royalblue", s=3, label="DS thresh 2")
plt.plot(time, interp_df.diameter.to_numpy(), c="violet", alpha=0.7, label="Interpolated")
# plt.plot(time, med_df.diameter.to_numpy(), c="lightblue", alpha=0.7, label="Median filtered")
plt.plot(time, lp_df.diameter.to_numpy(), c="orange", alpha=0.7, label="LP filtered")
plt.grid(True)
plt.legend(loc="best")
plt.show()

# Save the last dataframe to tsv file, check if it exists already
filename = "prepro_pupil_data.tsv"
tsv_file = os.path.join(output_folder, filename)
if os.path.isfile(tsv_file):
    user_input = input("\n{} already exists, do you want to overwrite it? (y/n)" .format(tsv_file))
    if user_input == "y":
        lp_df.to_csv(tsv_file, sep="\t", float_format='%.4f')
        print("\nSaved preprocessed data to: {}" .format(tsv_file))
    else:
        print("Quitting..")
else:
    lp_df.to_csv(tsv_file, sep="\t", float_format='%.4f')
    print("\nSaved preprocessed data to: {}".format(tsv_file))