Spaces:
Runtime error
Runtime error
File size: 2,204 Bytes
0d4381f | 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 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | """Face Parsing module for MuseTalk - Simplified implementation"""
import torch
import cv2
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
class FaceParsing:
def __init__(self, left_cheek_width=80, right_cheek_width=80):
self.left_cheek_width = left_cheek_width
self.right_cheek_width = right_cheek_width
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def __call__(self, image, size=(512, 512), mode="jaw"):
if isinstance(image, str):
image = Image.open(image)
width, height = image.size
with torch.no_grad():
image = image.resize(size, Image.BILINEAR)
face_mask = self._get_face_mask(image, mode=mode)
return face_mask
def _get_face_mask(self, image, mode="jaw"):
width, height = image.size
if mode == "jaw":
mask = np.zeros((height, width), dtype=np.uint8)
mask[:, int(height * 0.5) :] = 255
kernel = np.ones((21, 21), dtype=np.uint8)
kernel = cv2.erode(kernel, np.ones((3, 3), dtype=np.uint8), iterations=2)
cheek_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (35, 3))
mask = cv2.erode(mask, cheek_kernel, iterations=2)
cheek_mask = np.zeros((height, width), dtype=np.uint8)
center = width // 2
cv2.rectangle(
cheek_mask, (0, 0), (center - self.left_cheek_width, height), 255, -1
)
cv2.rectangle(
cheek_mask,
(center + self.right_cheek_width, 0),
(width, height),
255,
-1,
)
mask = cv2.bitwise_and(mask, cheek_mask)
return Image.fromarray(mask)
else:
mask = np.zeros((height, width), dtype=np.uint8)
mask[:, int(height * 0.5) :] = 255
return Image.fromarray(mask)
if __name__ == "__main__":
fp = FaceParsing()
test_img = Image.new("RGB", (512, 512), (100, 150, 200))
mask = fp(test_img, mode="jaw")
mask.save("test_mask.png")
print("Face parsing test complete")
|