92 lines
2.6 KiB
Python
92 lines
2.6 KiB
Python
import cv2
|
|
import yaml
|
|
import numpy as np
|
|
import mediapipe as mp
|
|
import pyrealsense2 as rs
|
|
from collections import deque
|
|
from biohead.algo import *
|
|
from biohead.utils import calc_feature
|
|
|
|
|
|
|
|
def init_interval_dict():
|
|
interval = {}
|
|
for key in HeadJoints().to_dict().keys():
|
|
interval[key] = [np.inf, -np.inf]
|
|
return interval
|
|
|
|
|
|
def update_interval(interval, result):
|
|
for key, val in result.to_dict().items():
|
|
if key in interval:
|
|
pre_min, pre_max = interval[key]
|
|
cur_min = min(getattr(result, key), pre_min)
|
|
cur_max = max(getattr(result, key), pre_max)
|
|
interval[key] = [cur_min, cur_max]
|
|
|
|
|
|
def save_interval_yaml(interval, filename="./config/calibrated_interval.yaml"):
|
|
# Wrap in AngleInterval
|
|
out = {}
|
|
for k, v in interval.items():
|
|
out[k] = [float(v[0]), float(v[1])]
|
|
with open(filename, "w") as f:
|
|
yaml.dump(out, f)
|
|
|
|
|
|
def calibrate():
|
|
with open(r"./config/config.yaml", "r") as f:
|
|
config = yaml.load(f, Loader=yaml.FullLoader)
|
|
|
|
w, h, fps = config['Camera']['image_width'], config['Camera']['image_height'], config['Camera']['fps']
|
|
pipe = rs.pipeline()
|
|
cfg = rs.config()
|
|
cfg.enable_stream(rs.stream.color, w, h, rs.format.bgr8, fps)
|
|
pipe.start(cfg)
|
|
align = rs.align(rs.stream.color)
|
|
|
|
mp_mesh = mp.solutions.face_mesh
|
|
mesh = mp_mesh.FaceMesh(
|
|
max_num_faces=1,
|
|
refine_landmarks=True,
|
|
min_detection_confidence=config['MediaPipe']['min_detection_confidence'],
|
|
min_tracking_confidence=config['MediaPipe']['min_tracking_confidence']
|
|
)
|
|
|
|
ray_origins = deque(maxlen=config['Smooth'])
|
|
ray_directions = deque(maxlen=config['Smooth'])
|
|
|
|
interval = init_interval_dict()
|
|
while True:
|
|
frames = align.process(pipe.wait_for_frames())
|
|
color_f = frames.get_color_frame()
|
|
if not color_f: continue
|
|
|
|
color = np.asanyarray(color_f.get_data())
|
|
h, w, _ = color.shape
|
|
rgb = cv2.cvtColor(color, cv2.COLOR_BGR2RGB)
|
|
|
|
res = mesh.process(rgb)
|
|
if not res.multi_face_landmarks: continue
|
|
|
|
uv = calc_feature(color, res, ray_origins, ray_directions)
|
|
result = HeadJoints()
|
|
result = calc_eyebrow(uv, result)
|
|
result = calc_eyelid(uv, result)
|
|
result = calc_eyeball(uv, result)
|
|
result = calc_mouth(uv, result)
|
|
result = calc_jaw(uv, result)
|
|
|
|
update_interval(interval, result)
|
|
|
|
cv2.imshow("Calibration", color)
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
break
|
|
|
|
pipe.stop()
|
|
cv2.destroyAllWindows()
|
|
save_interval_yaml(interval)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
calibrate() |