-
Notifications
You must be signed in to change notification settings - Fork 0
/
preprocess.py
120 lines (73 loc) · 3.6 KB
/
preprocess.py
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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
import numpy as np
import os
import cv2
import pandas as pd
from tqdm import tqdm
def center_gaze_direction(image, degree, gaze_org):
cols_per_degree = image.shape[1] // 360
north_col = - cols_per_degree * gaze_org + image.shape[1]//2
center_col = degree * cols_per_degree + north_col
if center_col >= image.shape[1]:
center_col -= image.shape[1]
left_col = center_col - image.shape[1]//2
output_image = np.concatenate((image[:, left_col:], image[:, :left_col]), axis=1)
return output_image
def get_home_direction(pos_x, pos_y):
if pos_x<0:
return np.rad2deg(np.arctan(pos_y/pos_x))
else:
return np.rad2deg(min(np.pi + np.arctan(pos_y/pos_x), -np.pi + np.arctan(pos_y/pos_x), key=abs))
def deg_to_unit_vector(deg):
return np.array([np.cos(np.deg2rad(deg)), np.sin(np.deg2rad(deg))])
def get_relative_home_direction(home_direction, gaze):
gaze_matrix = np.array([[np.cos(np.deg2rad(gaze)), np.sin(np.deg2rad(gaze))], [-np.sin(np.deg2rad(gaze)), np.cos(np.deg2rad(gaze))]])
relative_home_direction = np.matmul(gaze_matrix, home_direction)
return relative_home_direction
def preprocess(rectilinear_path, suffix=''):
data_name = rectilinear_path.split('/')[-1]
data_csv = f'{rectilinear_path}/data.csv'
image_info = pd.read_csv(data_csv)
output_csv = f'{rectilinear_path}/label_training_{data_name}{suffix}.csv'
output_location = f'{rectilinear_path}/training_{data_name}{suffix}/'
if not os.path.exists(output_location):
os.makedirs(output_location)
with open(output_csv, 'w') as f:
f.write('filename, gaze, pos_x, pos_y, label_1, label_2\n')
for file in tqdm(os.listdir(rectilinear_path)):
if not file.endswith('.jpg'):
continue
img = cv2.imread(os.path.join(rectilinear_path, file))
i = file.split('_')[0]
date = file.split('_')[2]
time = file.split('_')[3].split('.')[0]
i = i + "_" + date + "_" + time
#from image_info, find the row that path_idx == i
pos_x = image_info[image_info['path_idx'] == i]['pos_x'].values[0]
pos_y = image_info[image_info['path_idx'] == i]['pos_y'].values[0]
heading_org = image_info[image_info['path_idx'] == i]['heading'].values[0]
gaze_org = (heading_org*180/np.pi)%360
gaze_org = int(gaze_org)
home_angle_deg = get_home_direction(pos_x,pos_y)
home_vector = deg_to_unit_vector(home_angle_deg)
for k in range(1, 361):
gaze = k
if gaze > 360:
gaze -= 360
elif gaze < 1:
gaze += 360
image_gaze = center_gaze_direction(img, gaze, gaze_org)
if suffix == '_ds':
image_gaze = cv2.resize(image_gaze, (450, 48))
label = get_relative_home_direction(home_vector, gaze)
###### for distance #####
distance = np.sqrt(pos_x**2 + pos_y**2)
label = label * distance
if gaze < 10:
gaze = "00" + str(gaze)
elif gaze < 100:
gaze = "0" + str(gaze)
else:
gaze = str(gaze)
filename_gaze = str(i) + "_" + gaze + ".jpg"
cv2.imwrite(os.path.join(output_location, filename_gaze), image_gaze)
f.write(filename_gaze + "," + str(gaze) + "," + str(pos_x) + "," + str(pos_y) + "," + str(label[0]) + "," + str(label[1]) + "\n")