conan/processing/process_RTGene.py
2021-10-17 21:32:30 +02:00

364 lines
16 KiB
Python

import tensorflow as tf
import sys
import os
import argparse
import cv2
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from tqdm.notebook import tqdm
from rt_gene.gaze_tools import get_phi_theta_from_euler, limit_yaw
from rt_gene.extract_landmarks_method_base import LandmarkMethodBase
from rt_gene.estimate_gaze_base import GazeEstimatorBase
from rt_gene.estimate_gaze_tensorflow import GazeEstimator
from rt_gene.gaze_tools_standalone import euler_from_matrix
import itertools
import pandas as pd
#os.environ["CUDA_VISIBLE_DEVICES"]="1"
def getCenter(box):
return np.array([box[2]+box[0], box[3]+box[1]])/2
def load_camera_calibration(calibration_file):
fileType = calibration_file.split(".")[-1]
if fileType == "pkl":
import pickle
infile = open(calibration_file,'rb')
data = pickle.load(infile)
return data["distortion_coef"], data["camera_matrix"]
elif fileType == "yaml":
import yaml
with open(calibration_file, 'r') as f:
cal = yaml.safe_load(f)
dist_coefficients = np.array(cal['distortion_coefficients']['data'], dtype='float32').reshape(1, 5)
camera_matrix = np.array(cal['camera_matrix']['data'], dtype='float32').reshape(3, 3)
return dist_coefficients, camera_matrix
def extract_eye_image_patches(subjects, landmark_estimator):
for subject in subjects:
le_c, re_c, _, _ = subject.get_eye_image_from_landmarks(subject, landmark_estimator.eye_image_size)
subject.left_eye_color = le_c
subject.right_eye_color = re_c
def estimate_gaze(base_name, color_img, landmark_estimator, gaze_estimator, dist_coefficients, camera_matrix, args):
faceboxes = landmark_estimator.get_face_bb(color_img)
if len(faceboxes) == 0:
tqdm.write('Could not find faces in the image')
return
subjects = landmark_estimator.get_subjects_from_faceboxes(color_img, faceboxes)
extract_eye_image_patches(subjects, landmark_estimator)
input_r_list = []
input_l_list = []
input_head_list = []
valid_subject_list = []
roll_pitch_yaw_list = []
for idx, subject in enumerate(subjects):
if subject.left_eye_color is None or subject.right_eye_color is None:
#tqdm.write('Failed to extract eye image patches')
continue
success, rotation_vector, _ = cv2.solvePnP(landmark_estimator.model_points,
subject.landmarks.reshape(len(subject.landmarks), 1, 2),
cameraMatrix=camera_matrix,
distCoeffs=dist_coefficients, flags=cv2.SOLVEPNP_DLS)
if not success:
tqdm.write('Not able to extract head pose for subject {}'.format(idx))
continue
_rotation_matrix, _ = cv2.Rodrigues(rotation_vector)
_rotation_matrix = np.matmul(_rotation_matrix, np.array([[0, 1, 0], [0, 0, -1], [-1, 0, 0]]))
_m = np.zeros((4, 4))
_m[:3, :3] = _rotation_matrix
_m[3, 3] = 1
# Go from camera space to ROS space
_camera_to_ros = [[0.0, 0.0, 1.0, 0.0],
[-1.0, 0.0, 0.0, 0.0],
[0.0, -1.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 1.0]]
roll_pitch_yaw = list(euler_from_matrix(np.dot(_camera_to_ros, _m)))
roll_pitch_yaw = limit_yaw(roll_pitch_yaw)
roll_pitch_yaw_list.append(roll_pitch_yaw)
phi_head, theta_head = get_phi_theta_from_euler(roll_pitch_yaw)
face_image_resized = cv2.resize(subject.face_color, dsize=(224, 224), interpolation=cv2.INTER_CUBIC)
head_pose_image = landmark_estimator.visualize_headpose_result(face_image_resized, (phi_head, theta_head))
if args['vis_headpose']:
plt.axis("off")
plt.imshow(cv2.cvtColor(head_pose_image, cv2.COLOR_BGR2RGB))
plt.show()
if args['save_headpose']:
cv2.imwrite(os.path.join(args['output_path'], os.path.splitext(base_name)[0] + '_headpose.jpg'), head_pose_image)
input_r_list.append(gaze_estimator.input_from_image(subject.right_eye_color))
input_l_list.append(gaze_estimator.input_from_image(subject.left_eye_color))
input_head_list.append([theta_head, phi_head])
valid_subject_list.append(idx)
if len(valid_subject_list) == 0:
return
gaze_est = gaze_estimator.estimate_gaze_twoeyes(inference_input_left_list=input_l_list,
inference_input_right_list=input_r_list,
inference_headpose_list=input_head_list)
file_base = os.path.splitext(base_name)[0]
file = "_".join(file_base.split("_")[:-1])
frame = int(file_base.split("_")[-1])
ret = []
for subject_id, gaze, headpose, roll_pitch_yaw in zip(valid_subject_list, gaze_est.tolist(), input_head_list, roll_pitch_yaw_list):
subject = subjects[subject_id]
#print(roll_pitch_yaw)
# Build visualizations
r_gaze_img = gaze_estimator.visualize_eye_result(subject.right_eye_color, gaze)
l_gaze_img = gaze_estimator.visualize_eye_result(subject.left_eye_color, gaze)
s_gaze_img = np.concatenate((r_gaze_img, l_gaze_img), axis=1)
if args['vis_gaze']:
plt.axis("off")
plt.imshow(cv2.cvtColor(s_gaze_img, cv2.COLOR_BGR2RGB))
plt.show()
if args['save_gaze']:
cv2.imwrite(os.path.join(args['output_path'], os.path.splitext(base_name)[0] + '_gaze.jpg'), s_gaze_img)
# cv2.imwrite(os.path.join(args.output_path, os.path.splitext(base_name)[0] + '_left.jpg'), subject.left_eye_color)
# cv2.imwrite(os.path.join(args.output_path, os.path.splitext(base_name)[0] + '_right.jpg'), subject.right_eye_color)
if args['save_estimate']:
with open(os.path.join(args['output_path'], os.path.splitext(base_name)[0] + '_output.txt'), 'w+') as f:
f.write(os.path.splitext(base_name)[0] + ', [' + str(headpose[1]) + ', ' + str(headpose[0]) + ']' +
', [' + str(gaze[1]) + ', ' + str(gaze[0]) + ']' + '\n')
# Phi: pos - look down, neg - look up
# Theta: pos - rotate left, neg - rotate right
d = {"File":file, "Frame": frame, "SubjectId":subject_id, "HeadBox":subject.box, "Landmarks": subject.landmarks, "GazeTheta":gaze[0], "GazePhi":gaze[1], "HeadPoseTheta":headpose[0], "HeadPosePhi":headpose[1], "HeadPoseRoll":roll_pitch_yaw[0], "HeadPosePitch":roll_pitch_yaw[1], "HeadPoseYaw":roll_pitch_yaw[2]}
ret.append(d)
return ret
def visualize(df, FRAMES):
path_list = [f for f in os.listdir(FRAMES) if '.jpg' in f]
path_list.sort()
image = cv2.imread(os.path.join(FRAMES, path_list[0]))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
fig, ax = plt.subplots(1, figsize=(18,10))
for i in range(len(df.SubjectId.unique())):
bbox = df.loc[(df.Frame == 0) & (df.SubjectId == i)]['HeadBox'].values
print(bbox)
if not np.any(pd.isna(bbox)) and len(bbox) > 0:
bbox = np.array(bbox[0])
rect = patches.Rectangle((bbox[0],bbox[1]),bbox[2]-bbox[0],bbox[3]-bbox[1],linewidth=1,edgecolor='c',facecolor='none')
plt.text(bbox[0], bbox[1], 'ID%i' % i, color='c' ,fontsize=20)
ax.add_patch(rect)
ax.imshow(image)
plt.show()
def visualize_sorting(df_sorted):
subs = sorted(df_sorted[~df_sorted.PId.isna()].PId.unique())
for sid in subs:
x = df_sorted[df_sorted.PId==sid].HeadCenter.apply(lambda x: x[0])
y = df_sorted[df_sorted.PId==sid].HeadCenter.apply(lambda x: x[1])
frames = df_sorted[df_sorted.PId==sid].Frame.to_list()
plt.scatter(frames, x, alpha=.2, label = "Sub %i" % sid)
plt.legend()
plt.show()
def process(file, maxPeople, cameraRes = [5760, 2880]):
VIDEO = file
VIDEOOUT = VIDEO.split("/")[-1].split(".")[0]
ROOT = "/".join(VIDEO.split("/")[:-1]) + "/"
TMP_DIR = "/".join(VIDEO.split("/")[:-2]) + "/temp/"
FRAMES = "%s%s_frames" % (TMP_DIR, VIDEOOUT)
if not os.path.exists(VIDEO):
print('WARNING: Could not find video file')
return
script_path = "./"
args = {}
args["calib_file"] = "./calib_insta.pkl"
args["vis_headpose"] = False # store_false
args["save_headpose"] = False # store_false
args["vis_gaze"] = False # store_false
args["save_gaze"] = False # store_false
args["save_estimate"] = False # store_false
args["device_id_facedetection"] = "cuda:0" # store_false
args["im_path"] = os.path.join(script_path, './samples_gaze/')
args["output_path"] = os.path.join(script_path, './samples_gaze/')
args["models"] = [os.path.join(script_path, '../model_nets/Model_allsubjects1.h5')]
args['gaze_backend'] = 'tensorflow'
tqdm.write('Loading networks')
landmark_estimator = LandmarkMethodBase(device_id_facedetection=args["device_id_facedetection"],
checkpoint_path_face=os.path.join(script_path,
"../model_nets/SFD/s3fd_facedetector.pth"),
checkpoint_path_landmark=os.path.join(script_path,
"../model_nets/phase1_wpdc_vdc.pth.tar"),
model_points_file=os.path.join(script_path,
"../model_nets/face_model_68.txt"))
#gaze_estimator = GazeEstimator("/gpu:0", args['models'])
if args['gaze_backend'] == "tensorflow":
from rt_gene.estimate_gaze_tensorflow import GazeEstimator
gaze_estimator = GazeEstimator("/gpu:0", args['models'])
elif args['gaze_backend'] == "pytorch":
from rt_gene.estimate_gaze_pytorch import GazeEstimator
gaze_estimator = GazeEstimator("cuda:0", args['models'])
else:
raise ValueError("Incorrect gaze_base backend, choices are: tensorflow or pytorch")
if not os.path.isdir(args["output_path"]):
os.makedirs(args["output_path"])
video = cv2.VideoCapture(VIDEO)
print('Video frame count: ', video.get(cv2.CAP_PROP_FRAME_COUNT))
if args["calib_file"] is not None and os.path.exists(args["calib_file"]):
_dist_coefficients, _camera_matrix = load_camera_calibration(args["calib_file"])
else:
im_width = video.get(cv2.CAP_PROP_FRAME_WIDTH)
im_height = video.get(cv2.CAP_PROP_FRAME_HEIGHT)
print('WARNING!!! You should provide the camera calibration file, otherwise you might get bad results. \n\
Using a crude approximation!')
_dist_coefficients, _camera_matrix = np.zeros((1, 5)), np.array(
[[im_height, 0.0, im_width / 2.0], [0.0, im_height, im_height / 2.0], [0.0, 0.0, 1.0]])
lstRet = []
for i in tqdm(list(range(int(video.get(cv2.CAP_PROP_FRAME_COUNT))))):
image_file_name = "%s_%i.XXX" % (os.path.splitext(VIDEO)[0], i)
ret, image = video.read()
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
lstRet.append(estimate_gaze(image_file_name, image, landmark_estimator, gaze_estimator, _dist_coefficients, _camera_matrix, args))
lst = list(itertools.chain.from_iterable(lstRet))
df = pd.DataFrame(lst)
df["HeadCenter"] = df.HeadBox.apply(lambda x: getCenter(x))
df["Phi"] = df.GazePhi + df.HeadPosePhi # gaze yaw
df["Theta"] = df.GazeTheta + df.HeadPoseTheta # gaze pitch
df['Yaw'] = df.GazePhi + df.HeadPoseYaw
df['Pitch'] = df.GazeTheta + df.HeadPosePitch
# path = '%s%s_raw.pkl' % (TMP_DIR, VIDEOOUT)
# df.to_pickle(path)
# print('Saved raw detections to: ', path)
visualize(df, FRAMES)
# Sort ID detections
###############################################################################################################
# Find first frame where all are detected
for frame in sorted(df.Frame.unique()):
frame_df = df.loc[df.Frame == frame]
if len(frame_df['SubjectId'].unique()) == maxPeople:
first_frame = frame
print('First frame where all are detected: ', first_frame)
break
empty_rows = pd.DataFrame()
empty_rows['Frame'] = np.zeros(maxPeople).astype(int)
for col in df.columns:
if not col == 'Frame':
empty_rows[col] = df.loc[df.Frame == first_frame, [col]].values
df = df.loc[df.Frame != 0]
df = df.append(empty_rows).sort_values(by=['Frame'])
df.head()
df_sorted = df.copy()
df_sorted["PId"] = None
df_sorted.loc[df_sorted.Frame == 0, "PId"] = list(range(maxPeople))
df_sorted = df_sorted.sort_values("Frame")
df_sorted.index = list(range(len(df_sorted)))
for frameId in tqdm(sorted(df_sorted.Frame.unique())[1:]):
pidAssignement = []
for frameIdBefore in range(frameId - 1, -1, -1):
allFramesBefore = df_sorted[(df_sorted.Frame == frameIdBefore) & (~df_sorted.PId.isna())]
if (np.array_equal(sorted(allFramesBefore.PId.to_list()), np.arange(maxPeople))):
dfFramesCurrent = df_sorted[df_sorted.Frame == frameId]
for indexCurrentFrame, frameCurrent in dfFramesCurrent.iterrows():
lst = []
for indexBeforeFrame, frameBefore in allFramesBefore.iterrows():
if (frameBefore.HeadCenter[0] > frameCurrent.HeadCenter[0]):
p1 = np.array(frameCurrent.HeadCenter)
p2 = np.array(frameBefore.HeadCenter)
else:
p1 = np.array(frameBefore.HeadCenter)
p2 = np.array(frameCurrent.HeadCenter)
v1 = p1 - p2
dist1 = np.linalg.norm(v1)
p1[0] = p1[0] + cameraRes[0]
v2 = p1 - p2
dist2 = np.linalg.norm(v2)
dist = min([dist1, dist2])
lst.append([dist, frameCurrent.name, indexBeforeFrame, frameBefore])
lst.sort(key=lambda x: x[0])
pidAssignement.append([indexCurrentFrame, lst[0][-1].PId])
break
for index, pid in pidAssignement:
df_sorted.loc[df_sorted.index == index, "PId"] = pid
visualize_sorting(df_sorted)
del df_sorted["SubjectId"]
# Rearrange DataFrame: each ID has specific columns
###############################################################################################################
df_sorted = df_sorted[~df_sorted.PId.isna()].drop_duplicates(subset=['Frame', 'PId'])
FACE_COUNT = len(df_sorted[~df_sorted.PId.isna()].PId.unique())
df2 = df_sorted.pivot(index='Frame', columns="PId",
values=["Landmarks", "GazeTheta", "GazePhi", "HeadCenter", "HeadPoseTheta", "HeadPosePhi",
"HeadPoseYaw", "HeadPosePitch", "HeadPoseRoll", "Phi", "Theta"])
lst = []
for label in ["Landmarks", "GazeTheta", "GazePhi", "Head", "HeadPoseTheta", "HeadPosePhi", "HeadPoseYaw",
"HeadPosePitch", "HeadPoseRoll", "Phi", "Theta"]:
for head_id in range(FACE_COUNT):
lst.append("ID%i_%s" % (head_id, label))
df2.columns = lst
df2 = df2.reset_index()
path = "%s%s_RTGene.pkl" % (TMP_DIR, VIDEOOUT)
df2.to_pickle(path)
print("Saved RT-Gene detections to %s" % path)