Spaces:
Running
Running
File size: 1,138 Bytes
9340724 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
import sys
import os
from face_alignment import mtcnn
import argparse
from PIL import Image
from tqdm import tqdm
import random
from datetime import datetime
import torch
device = "cuda:0" if torch.cuda.is_available() else "cpu"
mtcnn_model = mtcnn.MTCNN(device=device, crop_size=(112, 112))
def add_padding(pil_img, top, right, bottom, left, color=(0,0,0)):
width, height = pil_img.size
new_width = width + right + left
new_height = height + top + bottom
result = Image.new(pil_img.mode, (new_width, new_height), color)
result.paste(pil_img, (left, top))
return result
def get_aligned_face(image_path, rgb_pil_image=None):
if rgb_pil_image is None:
img = Image.open(image_path).convert('RGB')
else:
assert isinstance(rgb_pil_image, Image.Image), 'Face alignment module requires PIL image or path to the image'
img = rgb_pil_image
# find face
try:
bboxes, faces = mtcnn_model.align_multi(img, limit=1)
face = faces[0]
except Exception as e:
print('Face detection Failed due to error.')
print(e)
face = None
return face
|