Akjava's picture
iniit
58665c8
raw
history blame
4.23 kB
import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
from mediapipe.framework.formats import landmark_pb2
from mediapipe import solutions
import numpy as np
# for X,Y,W,H to x1,y1,x2,y2(Left-top,right-bottom style)
def xywh_to_xyxy(box):
return [box[0],box[1],box[0]+box[2],box[1]+box[3]]
def convert_to_box(face_landmarks_list,indices,w=1024,h=1024):
x1=w
y1=h
x2=0
y2=0
for index in indices:
x=min(w,max(0,(face_landmarks_list[0][index].x*w)))
y=min(h,max(0,(face_landmarks_list[0][index].y*h)))
if x<x1:
x1=x
if y<y1:
y1=y
if x>x2:
x2=x
if y>y2:
y2=y
return [int(x1),int(y1),int(x2-x1),int(y2-y1)]
def box_to_square(bbox):
box=list(bbox)
if box[2]>box[3]:
diff = box[2]-box[3]
box[3]+=diff
box[1]-=diff/2
elif box[3]>box[2]:
diff = box[3]-box[2]
box[2]+=diff
box[0]-=diff/2
return box
def face_landmark_result_to_box(face_landmarker_result,width=1024,height=1024):
face_landmarks_list = face_landmarker_result.face_landmarks
full_indices = list(range(456))
MIDDLE_FOREHEAD = 151
BOTTOM_CHIN_EX = 152
BOTTOM_CHIN = 175
CHIN_TO_MIDDLE_FOREHEAD = [200,14,1,6,18,9]
MOUTH_BOTTOM = [202,200,422]
EYEBROW_CHEEK_LEFT_RIGHT = [46,226,50,1,280,446,276]
LEFT_HEAD_OUTER_EX = 251 #on side face almost same as full
LEFT_HEAD_OUTER = 301
LEFT_EYE_OUTER_EX = 356
LEFT_EYE_OUTER = 264
LEFT_MOUTH_OUTER_EX = 288
LEFT_MOUTH_OUTER = 288
LEFT_CHIN_OUTER = 435
RIGHT_HEAD_OUTER_EX = 21
RIGHT_HEAD_OUTER = 71
RIGHT_EYE_OUTER_EX = 127
RIGHT_EYE_OUTER = 34
RIGHT_MOUTH_OUTER_EX = 58
RIGHT_MOUTH_OUTER = 215
RIGHT_CHIN_OUTER = 150
# TODO naming line
min_indices=CHIN_TO_MIDDLE_FOREHEAD+EYEBROW_CHEEK_LEFT_RIGHT+MOUTH_BOTTOM
chin_to_brow_indices = [LEFT_CHIN_OUTER,LEFT_MOUTH_OUTER,LEFT_EYE_OUTER,LEFT_HEAD_OUTER,MIDDLE_FOREHEAD,RIGHT_HEAD_OUTER,RIGHT_EYE_OUTER,RIGHT_MOUTH_OUTER,RIGHT_CHIN_OUTER,BOTTOM_CHIN]+min_indices
box1 = convert_to_box(face_landmarks_list,min_indices,width,height)
box2 = convert_to_box(face_landmarks_list,chin_to_brow_indices,width,height)
box3 = convert_to_box(face_landmarks_list,full_indices,width,height)
#print(box)
return [box1,box2,box3,box_to_square(box1),box_to_square(box2),box_to_square(box3)]
def draw_landmarks_on_image(detection_result,rgb_image):
face_landmarks_list = detection_result.face_landmarks
annotated_image = np.copy(rgb_image)
# Loop through the detected faces to visualize.
for idx in range(len(face_landmarks_list)):
face_landmarks = face_landmarks_list[idx]
# Draw the face landmarks.
face_landmarks_proto = landmark_pb2.NormalizedLandmarkList()
face_landmarks_proto.landmark.extend([
landmark_pb2.NormalizedLandmark(x=landmark.x, y=landmark.y, z=landmark.z) for landmark in face_landmarks
])
solutions.drawing_utils.draw_landmarks(
image=annotated_image,
landmark_list=face_landmarks_proto,
connections=mp.solutions.face_mesh.FACEMESH_TESSELATION,
landmark_drawing_spec=None,
connection_drawing_spec=mp.solutions.drawing_styles
.get_default_face_mesh_tesselation_style())
return annotated_image
def mediapipe_to_box(image_data,model_path="face_landmarker.task"):
BaseOptions = mp.tasks.BaseOptions
FaceLandmarker = mp.tasks.vision.FaceLandmarker
FaceLandmarkerOptions = mp.tasks.vision.FaceLandmarkerOptions
VisionRunningMode = mp.tasks.vision.RunningMode
options = FaceLandmarkerOptions(
base_options=BaseOptions(model_asset_path=model_path),
running_mode=VisionRunningMode.IMAGE
,min_face_detection_confidence=0, min_face_presence_confidence=0
)
with FaceLandmarker.create_from_options(options) as landmarker:
if isinstance(image_data,str):
mp_image = mp.Image.create_from_file(image_data)
else:
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.asarray(image_data))
face_landmarker_result = landmarker.detect(mp_image)
boxes = face_landmark_result_to_box(face_landmarker_result,mp_image.width,mp_image.height)
return boxes,mp_image,face_landmarker_result