Akjava/mediapipe-68-points-facial-mask
11
1import spaces2import gradio as gr3import subprocess4from PIL import Image5import json6import os7import time8 9import mp_box10import draw_landmarks6811import landmarks68_utils12import io13import numpy as np14'''15Face landmark detection based Face Detection.16https://ai.google.dev/edge/mediapipe/solutions/vision/face_landmarker17from model card18https://storage.googleapis.com/mediapipe-assets/MediaPipe%20BlazeFace%20Model%20Card%20(Short%20Range).pdf19Licensed Apache License, Version 2.020Train with google's dataset(more detail see model card)21 22'''23 24dir_name ="files"25passed_time = 60*60 26def clear_old_files(dir,passed_time):27 try:28 files = os.listdir(dir)29 current_time = time.time()30 for file in files:31 file_path = os.path.join(dir,file)32 33 ctime = os.stat(file_path).st_ctime34 diff = current_time - ctime35 #print(f"ctime={ctime},current_time={current_time},passed_time={passed_time},diff={diff}")36 if diff > passed_time:37 os.remove(file_path)38 except:39 print("maybe still gallery using error")40 41def get_image_id(image,length=32):42 buffer = io.BytesIO()43 image.save(buffer, format='PNG')44 hash_object = hashlib.sha256(buffer.getvalue())45 hex_dig = hash_object.hexdigest()46 unique_id = hex_dig[:length]47 return unique_id48 49def save_image(image,extension="jpg"):50 id = get_image_id(image)51 os.makedirs(dir_name,exist_ok=True)52 file_path = f"{dir_name}/{id}.{extension}"53 54 image.save(file_path)55 return file_path56 57def picker_color_to_rgba(picker_color):58 color_value = picker_color.strip("rgba()").split(",")59 color_value[0] = int(float(color_value[0]))60 color_value[1] = int(float(color_value[1]))61 color_value[2] = int(float(color_value[2]))62 color_value[3] = int(float(color_value[3]))63 return color_value64 65#@spaces.GPU(duration=120)66def process_images(image,progress=gr.Progress(track_tqdm=True)):67 if image == None:68 raise gr.Error("Need Image")69 70 progress(0, desc="Start Mediapipe")71 72 boxes,mp_image,face_landmarker_result = mp_box.mediapipe_to_box(image)73 annotated_image,bbox,landmark_points = draw_landmarks68.draw_landmarks_on_image(image,face_landmarker_result)74 landmark_list = draw_landmarks68.convert_to_landmark_group_json(landmark_points)75 76 annotations = []77 galleries = []78 79 def append(mask,label):80 file_path = save_image(mask)81 galleries.append((file_path,label))82 annotations.append((np.array(mask.convert("1")),label))83 84 def fill_points(points,base_image=None):85 if base_image == None:86 base_image = landmarks68_utils.create_color_image(image.width,image.height,(0,0,0))87 landmarks68_utils.fill_points(base_image,points)88 return base_image89 90 # TODO support type91 left_eye_points = landmarks68_utils.get_landmark_points(landmark_list,landmarks68_utils.PARTS_LEFT_EYE)92 right_eye_points = landmarks68_utils.get_landmark_points(landmark_list,landmarks68_utils.PARTS_RIGHT_EYE)93 eyes_mask = fill_points(left_eye_points)94 eyes_mask = fill_points(right_eye_points,eyes_mask)95 append(eyes_mask,"eyes")96 97 98 upper_lip_points = landmarks68_utils.get_landmark_points(landmark_list,landmarks68_utils.PARTS_UPPER_LIP)99 upper_lip_mask = fill_points(upper_lip_points)100 append(upper_lip_mask,"upper-lip")101 102 lower_lip_points = landmarks68_utils.get_landmark_points(landmark_list,landmarks68_utils.PARTS_LOWER_LIP)103 lower_lip_mask = fill_points(lower_lip_points)104 append(lower_lip_mask,"lower-lip")105 106 inner_mouth_points = landmarks68_utils.get_innner_mouth_points(landmark_list)107 inner_mouth_mask = fill_points(inner_mouth_points)108 append(inner_mouth_mask,"inner-mouth")109 110 111 # TODO support type112 contour_points = landmarks68_utils.get_landmark_points(landmark_list,landmarks68_utils.PARTS_CONTOUR)113 114 contour_points=landmarks68_utils.get_face_points(landmark_list)115 116 contour_mask = fill_points(contour_points)117 append(contour_mask,"contour")118 119 mixed = Image.composite(eyes_mask,upper_lip_mask,eyes_mask.convert("L"))120 mixed = Image.composite(mixed,lower_lip_mask,mixed.convert("L"))121 mixed = Image.composite(mixed,inner_mouth_mask,mixed.convert("L"))122 append(mixed,"mixed")123 124 return [image,annotations],galleries125 126 127def write_file(file_path,text):128 with open(file_path, 'w', encoding='utf-8') as f:129 f.write(text)130 131def read_file(file_path):132 """read the text of target file133 """134 with open(file_path, 'r', encoding='utf-8') as f:135 content = f.read()136 137 return content138 139css="""140#col-left {141 margin: 0 auto;142 max-width: 640px;143}144#col-right {145 margin: 0 auto;146 max-width: 640px;147}148.grid-container {149 display: flex;150 align-items: center;151 justify-content: center;152 gap:10px153}154 155.image {156 width: 128px; 157 height: 128px; 158 object-fit: cover; 159}160 161.text {162 font-size: 16px;163}164"""165 166#css=css,167 168import hashlib169 170def text_to_sha256(text):171 text_bytes = text.encode('utf-8')172 hash_object = hashlib.sha256()173 hash_object.update(text_bytes)174 sha256_hex = hash_object.hexdigest()175 return sha256_hex176 177 178def create_json_download(text):179 file_id = f"{dir_name}/landmark_{text_to_sha256(text)[:32]}.json"180 write_file(file_id,text)181 # try to save182 return file_id183 184with gr.Blocks(css=css, elem_id="demo-container") as demo:185 with gr.Column():186 gr.HTML(read_file("demo_header.html"))187 gr.HTML(read_file("demo_tools.html"))188 with gr.Row():189 with gr.Column():190 image = gr.Image(height=800,sources=['upload','clipboard'],image_mode='RGB',elem_id="image_upload", type="pil", label="Upload")191 with gr.Row(elem_id="prompt-container", equal_height=False):192 with gr.Row():193 btn = gr.Button("Create Landmark 68 Mask", elem_id="run_button",variant="primary")194 195 with gr.Accordion(label="Advanced Settings", open=False):196 with gr.Row( equal_height=True):197 draw_number = gr.Checkbox(label="draw Number")198 199 font_scale = gr.Slider(200 label="Font Scale",201 minimum=0.1,202 maximum=2,203 step=0.1,204 value=0.5)205 206 text_color = gr.ColorPicker(value="rgba(200,200,200,1)",label="text color")207 #square_shape = gr.Checkbox(label="Square shape")208 with gr.Row( equal_height=True):209 210 line_color = gr.ColorPicker(value="rgba(0,0,255,1)",label="line color")211 line_size = gr.Slider(212 label="Line Size",213 minimum=0,214 maximum=20,215 step=1,216 value=1)217 with gr.Row( equal_height=True):218 dot_color = gr.ColorPicker(value="rgba(255,0,0,1)",label="dot color")219 dot_size = gr.Slider(220 label="Dot Size",221 minimum=0,222 maximum=40,223 step=1,224 value=3)225 with gr.Row( equal_height=True):226 box_color = gr.ColorPicker(value="rgba(200,200,200,1)",label="box color")227 box_size = gr.Slider(228 label="Box Size",229 minimum=0,230 maximum=20,231 step=1,232 value=1)233 with gr.Row( equal_height=True):234 json_format = gr.Radio(choices=["raw","face-detection"],value="face-detection",label="json-output format")235 236 with gr.Column():237 image_out = gr.AnnotatedImage(label="Output", elem_id="output-img")238 image_gallery = gr.Gallery(label="masks",preview=True)239 #download_button.click(fn=json_download,inputs=text_out,outputs=download_button)240 241 242 btn.click(fn=process_images, inputs=[image],outputs=[image_out,image_gallery] ,api_name='infer')243 gr.Examples(244 examples =["examples/00003245_00.jpg","examples/00004200.jpg","examples/00002200.jpg","examples/00005259.jpg","examples/00018022.jpg","examples/img-above.jpg","examples/img-below.jpg","examples/img-side.jpg"],245 inputs=[image]246 )247 gr.HTML(read_file("demo_footer.html"))248 249 if __name__ == "__main__":250 demo.launch()251