Spaces:
Runtime error
Runtime error
| import sys | |
| import torch | |
| import pickle | |
| import cv2 | |
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image | |
| from collections import defaultdict | |
| from glob import glob | |
| from matplotlib import pyplot as plt | |
| from matplotlib import animation | |
| from easydict import EasyDict as edict | |
| from huggingface_hub import hf_hub_download | |
| sys.path.append("./rome/") | |
| sys.path.append('./DECA') | |
| from rome.infer import Infer | |
| from rome.src.utils.processing import process_black_shape, tensor2image | |
| from rome.src.utils.visuals import mask_errosion | |
| # loading models ---- create model repo | |
| default_modnet_path = hf_hub_download('Pie31415/rome', 'modnet_photographic_portrait_matting.ckpt') | |
| default_model_path = hf_hub_download('Pie31415/rome', 'rome.pth') | |
| # parser configurations | |
| args = edict({ | |
| "save_dir": ".", | |
| "save_render": True, | |
| "model_checkpoint": default_model_path, | |
| "modnet_path": default_modnet_path, | |
| "random_seed": 0, | |
| "debug": False, | |
| "verbose": False, | |
| "model_image_size": 256, | |
| "align_source": True, | |
| "align_target": False, | |
| "align_scale": 1.25, | |
| "use_mesh_deformations": False, | |
| "subdivide_mesh": False, | |
| "renderer_sigma": 1e-08, | |
| "renderer_zfar": 100.0, | |
| "renderer_type": "soft_mesh", | |
| "renderer_texture_type": "texture_uv", | |
| "renderer_normalized_alphas": False, | |
| "deca_path": "DECA", | |
| "rome_data_dir": "rome/data", | |
| "autoenc_cat_alphas": False, | |
| "autoenc_align_inputs": False, | |
| "autoenc_use_warp": False, | |
| "autoenc_num_channels": 64, | |
| "autoenc_max_channels": 512, | |
| "autoenc_num_groups": 4, | |
| "autoenc_num_bottleneck_groups": 0, | |
| "autoenc_num_blocks": 2, | |
| "autoenc_num_layers": 4, | |
| "autoenc_block_type": "bottleneck", | |
| "neural_texture_channels": 8, | |
| "num_harmonic_encoding_funcs": 6, | |
| "unet_num_channels": 64, | |
| "unet_max_channels": 512, | |
| "unet_num_groups": 4, | |
| "unet_num_blocks": 1, | |
| "unet_num_layers": 2, | |
| "unet_block_type": "conv", | |
| "unet_skip_connection_type": "cat", | |
| "unet_use_normals_cond": True, | |
| "unet_use_vertex_cond": False, | |
| "unet_use_uvs_cond": False, | |
| "unet_pred_mask": False, | |
| "use_separate_seg_unet": True, | |
| "norm_layer_type": "gn", | |
| "activation_type": "relu", | |
| "conv_layer_type": "ws_conv", | |
| "deform_norm_layer_type": "gn", | |
| "deform_activation_type": "relu", | |
| "deform_conv_layer_type": "ws_conv", | |
| "unet_seg_weight": 0.0, | |
| "unet_seg_type": "bce_with_logits", | |
| "deform_face_tightness": 0.0001, | |
| "use_whole_segmentation": False, | |
| "mask_hair_for_neck": False, | |
| "use_hair_from_avatar": False, | |
| "use_scalp_deforms": True, | |
| "use_neck_deforms": True, | |
| "use_basis_deformer": False, | |
| "use_unet_deformer": True, | |
| "pretrained_encoder_basis_path": "", | |
| "pretrained_vertex_basis_path": "", | |
| "num_basis": 50, | |
| "basis_init": "pca", | |
| "num_vertex": 5023, | |
| "train_basis": True, | |
| "path_to_deca": "DECA", | |
| "path_to_linear_hair_model": "data/linear_hair.pth", # N/A | |
| "path_to_mobile_model": "data/disp_model.pth", # N/A | |
| "n_scalp": 60, | |
| "use_distill": False, | |
| "use_mobile_version": False, | |
| "deformer_path": "data/rome.pth", | |
| "output_unet_deformer_feats": 32, | |
| "use_deca_details": False, | |
| "use_flametex": False, | |
| "upsample_type": "nearest", | |
| "num_frequencies": 6, | |
| "deform_face_scale_coef": 0.0, | |
| "device": "cuda" | |
| }) | |
| # download FLAME and DECA pretrained | |
| generic_model_path = hf_hub_download('Pie31415/rome', 'generic_model.pkl') | |
| deca_model_path = hf_hub_download('Pie31415/rome', 'deca_model.tar') | |
| with open(generic_model_path, 'rb') as f: | |
| ss = pickle.load(f, encoding='latin1') | |
| with open('./DECA/data/generic_model.pkl', 'wb') as out: | |
| pickle.dump(ss, out) | |
| with open(deca_model_path, "rb") as input: | |
| with open('./DECA/data/deca_model.tar', "wb") as out: | |
| for line in input: | |
| out.write(line) | |
| # load ROME inference model | |
| infer = Infer(args) | |
| def image_inference( | |
| source_img: gr.inputs.Image = None, | |
| driver_img: gr.inputs.Image = None | |
| ): | |
| out = infer.evaluate(source_img, driver_img, crop_center=False) | |
| res = tensor2image(torch.cat([out['source_information']['data_dict']['source_img'][0].cpu(), | |
| out['source_information']['data_dict']['target_img'][0].cpu(), | |
| out['render_masked'].cpu(), out['pred_target_shape_img'][0].cpu()], dim=2)) | |
| return res[..., ::-1] | |
| def extract_frames( | |
| driver_vid: gr.inputs.Video = None | |
| ): | |
| image_frames = [] | |
| vid = cv2.VideoCapture(driver_vid) # path to mp4 | |
| while True: | |
| success, img = vid.read() | |
| if not success: break | |
| img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) | |
| pil_img = Image.fromarray(img) | |
| image_frames.append(pil_img) | |
| return image_frames | |
| def video_inference( | |
| source_img: gr.inputs.Image = None, | |
| driver_vid: gr.inputs.Video = None | |
| ): | |
| image_frames = extract_frames(driver_vid) | |
| resulted_imgs = defaultdict(list) | |
| mask_hard_threshold = 0.5 | |
| N = len(image_frames) | |
| for i in range(0, N, 4): # frame limits | |
| new_out = infer.evaluate(source_img, image_frames[i]) | |
| mask_pred = (new_out['pred_target_unet_mask'].cpu() > mask_hard_threshold).float() | |
| mask_pred = mask_errosion(mask_pred[0].float().numpy() * 255) | |
| render = new_out['pred_target_img'].cpu() * (mask_pred) + (1 - mask_pred) | |
| normals = process_black_shape(((new_out['pred_target_normal'][0].cpu() + 1) / 2 * mask_pred + (1 - mask_pred) ) ) | |
| normals[normals==0.5]=1. | |
| resulted_imgs['res_normal'].append(tensor2image(normals)) | |
| resulted_imgs['res_mesh_images'].append(tensor2image(new_out['pred_target_shape_img'][0])) | |
| resulted_imgs['res_renders'].append(tensor2image(render[0])) | |
| video = np.array(resulted_imgs['res_renders']) | |
| fig = plt.figure() | |
| im = plt.imshow(video[0,:,:,::-1]) | |
| plt.axis('off') | |
| plt.close() # this is required to not display the generated image | |
| def init(): | |
| im.set_data(video[0,:,:,::-1]) | |
| def animate(i): | |
| im.set_data(video[i,:,:,::-1]) | |
| return im | |
| anim = animation.FuncAnimation(fig, animate, init_func=init, frames=video.shape[0], interval=30) | |
| anim.save("avatar.gif", dpi=300, writer = animation.PillowWriter(fps=24)) | |
| return "avatar.gif" | |
| description = """<p style='text-align: center'> Create a personal avatar from just a single image using ROME. <br> <a href='https://arxiv.org/abs/2206.08343' target='_blank'>Paper</a> | <a href='https://samsunglabs.github.io/rome' target='_blank'>Project Page</a> | <a href='https://github.com/SamsungLabs/rome' target='_blank'>Github</a> </p>""" | |
| quote = """ | |
| > <p style='text-align: center'> [The] system creates realistic mesh-based avatars from a single <strong>source</strong> photo. These avatars are rigged, i.e., they can be driven by the animation parameters from a different <strong>driving</strong> frame. </p>""" | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# **<p align='center'>ROME: Realistic one-shot mesh-based head avatars</p>**") | |
| gr.HTML(value="<img src='file/media/tease.gif' alt='Teaser' style='display: block; margin: auto;'>") | |
| gr.Markdown(description) | |
| gr.Markdown(quote) | |
| with gr.Tab("Image Inference"): | |
| with gr.Row(): | |
| source_img = gr.Image(type="pil", label="Source image", show_label=True) | |
| driver_img = gr.Image(type="pil", label="Driver image", show_label=True) | |
| image_output = gr.Image(label="Rendered avatar") | |
| image_button = gr.Button("Predict") | |
| with gr.Tab("Video Inference"): | |
| with gr.Row(): | |
| source_img2 = gr.Image(type="pil", label="Source image", show_label=True) | |
| driver_vid = gr.Video(label="Driver video", source="upload") | |
| video_output = gr.Image(label="Rendered GIF avatar") | |
| video_button = gr.Button("Predict") | |
| with gr.Tab("Webcam Inference"): | |
| with gr.Row(): | |
| source_img3 = gr.Image(type="pil", label="Source image", show_label=True) | |
| driver_cam = gr.Video(label="Driver video", source="webcam") | |
| cam_output = gr.Image(label="Rendered GIF avatar") | |
| cam_button = gr.Button("Predict") | |
| gr.Examples( | |
| examples=[ | |
| ["./examples/lincoln.jpg", "./examples/taras2.jpg"], | |
| ["./examples/lincoln.jpg", "./examples/taras1.jpg"] | |
| ], | |
| inputs=[source_img, driver_img], | |
| outputs=[image_output], | |
| fn=image_inference, | |
| cache_examples=True | |
| ) | |
| image_button.click(image_inference, inputs=[source_img, driver_img], outputs=image_output) | |
| video_button.click(video_inference, inputs=[source_img2, driver_vid], outputs=video_output) | |
| cam_button.click(video_inference, inputs=[source_img3, driver_cam], outputs=cam_output) | |
| demo.launch() |