Spaces:
Runtime error
Runtime error
| ######################################################################################## | |
| import gradio as gr | |
| import cv2 | |
| import matplotlib | |
| import matplotlib.cm | |
| import mediapipe as mp | |
| import numpy as np | |
| import os | |
| import pygltflib | |
| import shutil | |
| import struct | |
| import tempfile | |
| import torch | |
| from mediapipe.framework.formats import landmark_pb2 | |
| from mediapipe.python.solutions.drawing_utils import _normalized_to_pixel_coordinates | |
| from PIL import Image | |
| from quads import QUADS | |
| from typing import List, Mapping, Optional, Tuple, Union | |
| from utils import colorize | |
| class face_image_to_face_mesh: | |
| def __init__(self): | |
| self.zoe_me = True | |
| self.css = (""" | |
| #img-display-container { | |
| max-height: 50vh; | |
| } | |
| #img-display-input { | |
| max-height: 40vh; | |
| } | |
| #img-display-output { | |
| max-height: 55vh; | |
| max-width: 55vh; | |
| width:auto; | |
| height:auto | |
| } | |
| """) | |
| def demo(self): | |
| if self.zoe_me: | |
| DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| self.zoe = torch.hub.load('isl-org/ZoeDepth', "ZoeD_N", pretrained=True).to(DEVICE).eval() | |
| demo = gr.Blocks(css=self.css) | |
| with demo: | |
| gr.Markdown(""" | |
| # Face Image to Face Quad Mesh | |
| Uses MediaPipe to detect a face in an image and convert it to a quad mesh. | |
| Saves to OBJ since gltf does not support quad faces. The 3d viewer has Y pointing the opposite direction from Blender, so ya hafta spin it. | |
| The face depth with Zoe can be a bit much and without it is a bit generic. In blender you can fix this just by snapping to the high poly model. | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| upload_image = gr.Image(label="Input image", type="numpy", source="upload") | |
| gr.Examples( examples=[ | |
| 'examples/blonde-00019-1421846474.png', | |
| 'examples/dude-00110-1227390728.png', | |
| 'examples/granny-00056-1867315302.png', | |
| 'examples/tuffie-00039-499759385.png', | |
| 'examples/character.png', | |
| ], inputs=[upload_image] ) | |
| upload_image_btn = gr.Button(value="Detect faces") | |
| if self.zoe_me: | |
| with gr.Group(): | |
| use_zoe = gr.Checkbox(label="Use ZoeDepth for Z", value=True) | |
| gr.Textbox(show_label=False, value="Override the MediaPipe depth with ZoeDepth.") | |
| zoe_scale = gr.Slider(label="Zoe Scale", value=1.44, minimum=0.0, maximum=3.3, step=0.1) | |
| gr.Textbox(show_label=False, value="How much to scale the ZoeDepth. 2x is pretty dramatic...") | |
| else: | |
| use_zoe = False | |
| zoe_scale = 0 | |
| with gr.Group(): | |
| min_detection_confidence = gr.Slider(label="Min detection confidence", value=0.5, minimum=0.0, maximum=1.0, step=0.01) | |
| gr.Textbox(show_label=False, value="Minimum confidence value ([0.0, 1.0]) from the face detection model for the detection to be considered successful.") | |
| with gr.Group(): | |
| gr.Markdown( | |
| """ | |
| The initial workflow I was imagining was: | |
| 1. sculpt high poly mesh in blender | |
| 2. snapshot the face | |
| 3. generate the mesh using the mediapipe stuff | |
| 4. import the low poly mediapipe face | |
| 5. snap the mesh to the high poly model | |
| 6. model the rest of the low poly model | |
| 7. bake the normal / etc maps to the low poly face model | |
| 8. it's just that easy 😛 | |
| Ideally it would be a plugin... | |
| """) | |
| with gr.Column(): | |
| with gr.Group(): | |
| num_faces_detected = gr.Number(label="Number of faces detected", value=0) | |
| output_image = gr.Image(label="Output image",elem_id='img-display-output') | |
| output_mesh = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="3D Model",elem_id='img-display-output') | |
| depth_image = gr.Image(label="Depth image",elem_id='img-display-output') | |
| upload_image_btn.click( | |
| fn=self.detect, | |
| inputs=[upload_image, min_detection_confidence,use_zoe,zoe_scale], | |
| outputs=[output_mesh, output_image, depth_image, num_faces_detected] | |
| ) | |
| demo.launch() | |
| def detect(self, image, min_detection_confidence, use_zoe,zoe_scale): | |
| width = image.shape[1] | |
| height = image.shape[0] | |
| ratio = width / height | |
| mp_drawing = mp.solutions.drawing_utils | |
| mp_drawing_styles = mp.solutions.drawing_styles | |
| mp_face_mesh = mp.solutions.face_mesh | |
| mesh = "examples/jackiechan.obj" | |
| if self.zoe_me and use_zoe: | |
| depth = self.zoe.infer_pil(image) | |
| print( f'type of depth is {type(depth)}' ) | |
| idepth = colorize(depth, cmap='gray_r') | |
| else: | |
| depth = None | |
| idepth = image | |
| drawing_spec = mp_drawing.DrawingSpec(thickness=1, circle_radius=1) | |
| with mp_face_mesh.FaceMesh( | |
| static_image_mode=True, | |
| max_num_faces=1, | |
| min_detection_confidence=min_detection_confidence) as face_mesh: | |
| results = face_mesh.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) | |
| if not results.multi_face_landmarks: | |
| return mesh, image, idepth, 0 | |
| annotated_image = image.copy() | |
| for face_landmarks in results.multi_face_landmarks: | |
| mesh = self.toObj(image=image, width=width, height=height, ratio=ratio, landmark_list=face_landmarks, depth=depth, zoe_scale=zoe_scale) | |
| mp_drawing.draw_landmarks( | |
| image=annotated_image, | |
| landmark_list=face_landmarks, | |
| connections=mp_face_mesh.FACEMESH_TESSELATION, | |
| landmark_drawing_spec=None, | |
| connection_drawing_spec=mp_drawing_styles | |
| .get_default_face_mesh_tesselation_style()) | |
| mp_drawing.draw_landmarks( | |
| image=annotated_image, | |
| landmark_list=face_landmarks, | |
| connections=mp_face_mesh.FACEMESH_CONTOURS, | |
| landmark_drawing_spec=None, | |
| connection_drawing_spec=mp_drawing_styles | |
| .get_default_face_mesh_contours_style()) | |
| return mesh, annotated_image, idepth, 1 | |
| def toObj( self, image: np.ndarray, width:int, height:int, ratio: float, landmark_list: landmark_pb2.NormalizedLandmarkList, depth: np.ndarray, zoe_scale: float): | |
| print( f'you have such pretty hair' ) | |
| obj_file = tempfile.NamedTemporaryFile(suffix='.obj', delete=False) | |
| mtl_file = tempfile.NamedTemporaryFile(suffix='.mtl', delete=False) | |
| png_file = tempfile.NamedTemporaryFile(suffix='.png', delete=False) | |
| ############################################ | |
| lines = [] | |
| lines.append( f'mtllib {os.path.basename(mtl_file.name)}' ) | |
| (points,coordinates) = self.landmarksToPoints( width, height, ratio, landmark_list, depth, zoe_scale ) | |
| for point in points: | |
| lines.append( "v " + " ".join([str(value) for value in point]) ) | |
| for coordinate in coordinates: | |
| lines.append( "vt " + " ".join([str(value) for value in coordinate]) ) | |
| for quad in QUADS: | |
| normal = self.totallyNormal( points[ quad[ 0 ] -1 ], points[ quad[ 1 ] -1 ], points[ quad[ 2 ] -1 ] ) | |
| lines.append( "vn " + " ".join([str(value) for value in normal]) ) | |
| lines.append( 'usemtl MyMaterial' ) | |
| quadIndex = 0 | |
| for quad in QUADS: | |
| quadIndex = 1 + quadIndex | |
| if True: | |
| lines.append( "f " + " ".join([f'{vertex}/{vertex}/{quadIndex}' for vertex in quad]) ) | |
| else: | |
| lines.append( "f " + " ".join([str(vertex) for vertex in quad]) ) | |
| out = open( obj_file.name, 'w' ) | |
| out.write( '\n'.join( lines ) + '\n' ) | |
| out.close() | |
| shutil.copy(obj_file.name, "/tmp/lol.obj") | |
| ############################################ | |
| lines = [] | |
| lines.append( 'newmtl MyMaterial' ) | |
| lines.append( f'Ka 1.000 1.000 1.000 # white' ) | |
| lines.append( f'Kd 1.000 1.000 1.000 # white' ) | |
| lines.append( f'Ks 0.000 0.000 0.000 # black (off)' ) | |
| lines.append( f'map_Ka {os.path.basename(png_file.name)}' ) | |
| lines.append( f'map_Kd {os.path.basename(png_file.name)}' ) | |
| out = open( mtl_file.name, 'w' ) | |
| out.write( '\n'.join( lines ) + '\n' ) | |
| out.close() | |
| shutil.copy(mtl_file.name, "/tmp/lol.mtl") | |
| ############################################ | |
| cv2.imwrite(png_file.name, cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) | |
| shutil.copy(png_file.name, "/tmp/lol.png") | |
| ############################################ | |
| print( f'I know it is special to you so I saved it to {obj_file.name} since we are friends' ) | |
| return obj_file.name | |
| def landmarksToPoints( self, width: int, height: int, ratio: float, landmark_list: landmark_pb2.NormalizedLandmarkList, depth: np.ndarray, zoe_scale: float ): | |
| points = [] # 3d vertices | |
| coordinates = [] # 2d texture coordinates | |
| mins = [+np.inf] * 3 | |
| maxs = [-np.inf] * 3 | |
| for idx, landmark in enumerate(landmark_list.landmark): | |
| if ((landmark.HasField('visibility') and | |
| landmark.visibility < _VISIBILITY_THRESHOLD) or | |
| (landmark.HasField('presence') and | |
| landmark.presence < _PRESENCE_THRESHOLD)): | |
| idk_what_to_do_for_this = True | |
| x, y = _normalized_to_pixel_coordinates(landmark.x,landmark.y,width,height) | |
| coordinates.append( [x/width,1-y/height] ) | |
| if depth is not None: | |
| landmark.z = depth[y, x] * zoe_scale | |
| point = [landmark.x * ratio, -landmark.y, -landmark.z]; | |
| for pidx,value in enumerate( point ): | |
| mins[pidx] = min(mins[pidx],value) | |
| maxs[pidx] = max(maxs[pidx],value) | |
| points.append( point ) | |
| mids = [(min_val + max_val) / 2 for min_val, max_val in zip(mins, maxs)] | |
| for idx,point in enumerate( points ): | |
| points[idx] = [(val-mid) for val, mid in zip(point,mids)] | |
| print( f'mins: {mins}' ) | |
| print( f'mids: {mids}' ) | |
| print( f'maxs: {maxs}' ) | |
| return (points,coordinates) | |
| def totallyNormal(self, p0, p1, p2): | |
| v1 = np.array(p1) - np.array(p0) | |
| v2 = np.array(p2) - np.array(p0) | |
| normal = np.cross(v1, v2) | |
| normal = normal / np.linalg.norm(normal) | |
| return normal.tolist() | |
| face_image_to_face_mesh().demo() | |
| # EOF | |
| ######################################################################################## | |