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 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, get_most_recent_subdirectory | |
| class face_image_to_face_mesh: | |
| def __init__(self): | |
| self.zoe_me = True | |
| self.uvwrap = not True | |
| 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(), cache_examples=True) | |
| with demo: | |
| gr.Markdown(self.header()) | |
| with gr.Row(): | |
| with gr.Column(): | |
| upload_image = gr.Image(label="Input image", type="numpy", source="upload") | |
| self.temp_dir = get_most_recent_subdirectory( upload_image.DEFAULT_TEMP_DIR ) | |
| print( f'The temp_dir is {self.temp_dir}' ) | |
| gr.Examples( examples=[ | |
| 'examples/blonde-00081-399357008.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(): | |
| zoe_scale = gr.Slider(label="Mix the ZoeDepth with the MediaPipe Depth", value=1, minimum=0, maximum=1, step=.01) | |
| flat_scale = gr.Slider(label="Depth scale, smaller is flatter and possibly more flattering", value=1, minimum=0, maximum=1, step=.01) | |
| min_detection_confidence = gr.Slider(label="Mininum face detection confidence", value=.5, minimum=0, maximum=1.0, step=0.01) | |
| else: | |
| use_zoe = False | |
| zoe_scale = 0 | |
| with gr.Group(): | |
| gr.Markdown(self.footer()) | |
| with gr.Column(): | |
| with gr.Group(): | |
| output_mesh = gr.Model3D(clear_color=3*[0], label="3D Model",elem_id='mesh-display-output') | |
| output_image = gr.Image(label="Output image",elem_id='img-display-output') | |
| depth_image = gr.Image(label="Depth image",elem_id='img-display-output') | |
| num_faces_detected = gr.Number(label="Number of faces detected", value=0) | |
| upload_image_btn.click( | |
| fn=self.detect, | |
| inputs=[upload_image, min_detection_confidence,zoe_scale,flat_scale], | |
| outputs=[output_mesh, output_image, depth_image, num_faces_detected] | |
| ) | |
| demo.launch() | |
| def detect(self, image, min_detection_confidence, zoe_scale, flat_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/converted/in-granny.obj" | |
| if self.zoe_me and 0 < zoe_scale: | |
| depth = self.zoe.infer_pil(image) | |
| 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,mtl,png) = self.toObj(image=image, width=width, height=height, ratio=ratio, landmark_list=face_landmarks, depth=depth, zoe_scale=zoe_scale, flat_scale=flat_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, flat_scale: float): | |
| print( f'you have such pretty hair', self.temp_dir ) | |
| hf_hack = True | |
| if hf_hack: | |
| obj_file = tempfile.NamedTemporaryFile(suffix='.obj', delete=False) | |
| mtl_file = tempfile.NamedTemporaryFile(suffix='.mtl', delete=False) | |
| png_file = tempfile.NamedTemporaryFile(suffix='.png', delete=False) | |
| else: | |
| obj_file = tempfile.NamedTemporaryFile(suffix='.obj', dir=self.temp_dir, delete=False) | |
| mtl_file = tempfile.NamedTemporaryFile(suffix='.mtl', dir=self.temp_dir, delete=False) | |
| png_file = tempfile.NamedTemporaryFile(suffix='.png', dir=self.temp_dir, delete=False) | |
| ############################################ | |
| (points,coordinates,colors) = self.landmarksToPoints( image, width, height, ratio, landmark_list, depth, zoe_scale, flat_scale ) | |
| ############################################ | |
| lines = [] | |
| lines.append( f'o MyMesh' ) | |
| if hf_hack: | |
| # the 'file=' is a gradio hack | |
| lines.append( f'#mtllib file={mtl_file.name}' ) | |
| else: | |
| # the 'file=' is a gradio hack | |
| lines.append( f'mtllib file={mtl_file.name}' ) | |
| for index, point in enumerate(points): | |
| color = colors[index] | |
| scaled_color = [value / 255 for value in color] # Scale colors down to 0-1 range | |
| flipped = [-value for value in point] | |
| flipped[ 0 ] = -flipped[ 0 ] | |
| lines.append( "v " + " ".join(map(str, flipped + color)) ) | |
| for coordinate in coordinates: | |
| lines.append( "vt " + " ".join([str(value) for value in coordinate]) ) | |
| for quad in QUADS: | |
| #quad = list(reversed(quad)) | |
| 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 | |
| face_uv = "f " + " ".join([f'{vertex}/{vertex}/{quadIndex}' for vertex in quad]) | |
| face_un = "f " + " ".join([str(vertex) for vertex in quad]) | |
| if self.uvwrap: | |
| lines.append( face_uv ) | |
| else: | |
| lines.append( f'#{face_uv}' ) | |
| lines.append( f'{face_un}' ) | |
| #"f " + " ".join([str(vertex) for vertex in quad]) ) | |
| out = open( obj_file.name, 'w' ) | |
| out.write( '\n'.join( lines ) + '\n' ) | |
| out.close() | |
| ############################################ | |
| lines = [] | |
| lines.append( 'newmtl MyMaterial' ) | |
| lines.append( f'map_Kd file={png_file.name}' ) # the 'file=' is a gradio hack | |
| out = open( mtl_file.name, 'w' ) | |
| out.write( '\n'.join( lines ) + '\n' ) | |
| out.close() | |
| ############################################ | |
| cv2.imwrite(png_file.name, cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) | |
| ############################################ | |
| 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,mtl_file.name,png_file.name) | |
| def landmarksToPoints( self, image:np.ndarray, width: int, height: int, ratio: float, landmark_list: landmark_pb2.NormalizedLandmarkList, depth: np.ndarray, zoe_scale: float, flat_scale: float ): | |
| points = [] # 3d vertices | |
| coordinates = [] # 2d texture coordinates | |
| colors = [] # 3d rgb info | |
| mins = [+np.inf] * 3 | |
| maxs = [-np.inf] * 3 | |
| mp_scale = 1 - zoe_scale | |
| print( f'zoe_scale:{zoe_scale}, mp_scale:{mp_scale}' ) | |
| for idx, landmark in enumerate(landmark_list.landmark): | |
| x, y = _normalized_to_pixel_coordinates(landmark.x,landmark.y,width,height) | |
| color = image[y,x] | |
| colors.append( [value / 255 for value in color ] ) | |
| coordinates.append( [x/width,1-y/height] ) | |
| if depth is not None: | |
| landmark.z = depth[y, x] * zoe_scale + mp_scale * landmark.z | |
| landmark.z = landmark.z * flat_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,colors) | |
| 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() | |
| def header(self): | |
| return (""" | |
| # Image to 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. For photos turning it down to .4 helps, but may still need cleanup... | |
| Highly recommend running it locally. The 3D model has uv values in the faces, but you will have to either use the script or do some manually tomfoolery. | |
| Quick import result in examples/converted/movie-gallery.mp4 under files | |
| """) | |
| def footer(self): | |
| return ( """ | |
| # Using the Textured Mesh in Blender | |
| There a couple of annoying steps atm after you download the obj from the 3d viewer. | |
| You can use the script meshin-around.sh in the files section to do the conversion or manually: | |
| 1. edit the file and change the mtllib line to use fun.mtl | |
| 2. replace / delete all lines that start with 'f', eg :%s,^f.*,, | |
| 3. uncomment all the lines that start with '#f', eg: :%s,^#f,f, | |
| 4. save and exit | |
| 5. create fun.mtl to point to the texture like: | |
| ``` | |
| newmtl MyMaterial | |
| map_Kd fun.png | |
| ``` | |
| Make sure the obj, mtl and png are all in the same directory | |
| Now the import will have the texture data: File -> Import -> Wavefront (obj) -> fun.obj | |
| This is all a work around for a weird hf+gradios+babylonjs bug which seems to be related to the version | |
| of babylonjs being used... It works fine in a local babylonjs sandbox... | |
| # Suggested Workflows | |
| Here are some workflow ideas. | |
| ## retopologize high poly face mesh | |
| 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... | |
| ## stable diffusion integration | |
| 1. generate a face in sd | |
| 2. generate the mesh | |
| 3. repose it and use it for further generation | |
| May need to expanded the generated mesh to cover more, maybe with | |
| <a href="https://github.com/shunsukesaito/PIFu" target="_blank">PIFu model</a>. | |
| """) | |
| def css(self): | |
| return (""" | |
| #mesh-display-output { | |
| max-height: 44vh; | |
| max-width: 44vh; | |
| width:auto; | |
| height:auto | |
| } | |
| #img-display-output { | |
| max-height: 28vh; | |
| max-width: 28vh; | |
| width:auto; | |
| height:auto | |
| } | |
| """) | |
| face_image_to_face_mesh().demo() | |
| # EOF | |
| ######################################################################################## | |