added optional mask to detectfaces

experimental: automatically adds to the mask black pixels adjacent to the edges of crops
This commit is contained in:
David Chatel
2024-03-13 12:02:48 +01:00
parent b125ee5590
commit 49341156b6
2 changed files with 20 additions and 7 deletions
+19 -6
View File
@@ -106,6 +106,9 @@ class DetectFaces:
'threshold': ('FLOAT', {'default': 0.5, 'min': 0.0, 'max': 1.0, 'step': 0.01}),
'min_size': ('INT', {'default': 64, 'max': 512, 'step': 8}),
'max_size': ('INT', {'default': 512, 'min': 512, 'step': 8}),
},
'optional': {
'mask': ('MASK',),
}
}
@@ -114,10 +117,13 @@ class DetectFaces:
FUNCTION = 'run'
CATEGORY = 'facetools'
def run(self, image, threshold, min_size, max_size):
def run(self, image, threshold, min_size, max_size, mask=None):
faces = []
images = (image * 255).type(torch.uint8)
for i, img in enumerate(images):
masked = image
if mask is not None:
masked = image * tv.transforms.functional.resize(1-mask, image.shape[1:3])[..., None]
masked = (masked * 255).type(torch.uint8)
for i, img in enumerate(masked):
unfiltered_faces = detect_faces(img, threshold)
for face in unfiltered_faces:
a, b, c, d = face.bbox
@@ -125,7 +131,7 @@ class DetectFaces:
w = abs(c-a)
if (h <= max_size or w <= max_size) and (min_size <= h or min_size <= w):
face.image_idx = i
face.image = img
face.image = image[i]
faces.append(face)
return (faces,)
@@ -161,8 +167,8 @@ class CropFaces:
warps = []
for face in faces:
M, crop = face.crop(crop_size, crop_factor)
mask = mask_crop(face, M, crop, mask_type)
crops.append(np.array(crop) / 255)
mask = mask_crop(face, M, crop*255, mask_type)
crops.append(np.array(crop))
masks.append(np.array(mask))
warps.append(M)
crops = torch.from_numpy(np.array(crops)).type(torch.float32)
@@ -222,6 +228,13 @@ class WarpFaceBack:
def run(self, images, face, crop, mask, warp):
groups = defaultdict(list)
for f,c,m,w in zip(face, crop, mask, warp):
gray = cv2.cvtColor(c.numpy(), cv2.COLOR_RGB2GRAY)
_, gray = cv2.threshold(gray, 0.01, 1, cv2.THRESH_BINARY)
gray = gray.astype(np.uint8)
cnts, _ = cv2.findContours(gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
largeCnts = list(filter(lambda x: cv2.contourArea(x) > 10000, cnts))
gray = cv2.drawContours(gray, largeCnts, -1, 1, -1)
m *= torch.from_numpy(gray)
groups[f.image_idx].append((f.image,c,m,w))
results = []
+1 -1
View File
@@ -93,7 +93,7 @@ class Face:
T3 = np.array([[1, 0, -cx], [0, 1, -cy], [0, 0, 1]])
T4 = np.array([[1, 0, cx], [0, 1, cy], [0, 0, 1]])
N = N @ T4 @ S @ T3
crop = cv2.warpAffine(self.img.numpy(), N, (size, size))
crop = cv2.warpAffine(self.image.numpy(), N, (size, size))
return N, crop