Update facetools.py

Facetools nodes were cutting off faces that extended beyond chunk boundaries in some cases. All faces found are considered and are coalesced rather than pruned, meaning that you should not see half a face any more.
This commit is contained in:
Jonathan 2023-10-13 15:07:08 -05:00 committed by psychedelicious
parent 6994783c17
commit ffb01f1345

View File

@ -46,6 +46,8 @@ class FaceResultData(TypedDict):
y_center: float y_center: float
mesh_width: int mesh_width: int
mesh_height: int mesh_height: int
chunk_x_offset: int
chunk_y_offset: int
class FaceResultDataWithId(FaceResultData): class FaceResultDataWithId(FaceResultData):
@ -78,6 +80,48 @@ FONT_SIZE = 32
FONT_STROKE_WIDTH = 4 FONT_STROKE_WIDTH = 4
def coalesce_faces(face1: FaceResultData, face2: FaceResultData) -> FaceResultData:
face1_x_offset = face1["chunk_x_offset"] - min(face1["chunk_x_offset"], face2["chunk_x_offset"])
face2_x_offset = face2["chunk_x_offset"] - min(face1["chunk_x_offset"], face2["chunk_x_offset"])
face1_y_offset = face1["chunk_y_offset"] - min(face1["chunk_y_offset"], face2["chunk_y_offset"])
face2_y_offset = face2["chunk_y_offset"] - min(face1["chunk_y_offset"], face2["chunk_y_offset"])
new_im_width = (
max(face1["image"].width, face2["image"].width)
+ max(face1["chunk_x_offset"], face2["chunk_x_offset"])
- min(face1["chunk_x_offset"], face2["chunk_x_offset"])
)
new_im_height = (
max(face1["image"].height, face2["image"].height)
+ max(face1["chunk_y_offset"], face2["chunk_y_offset"])
- min(face1["chunk_y_offset"], face2["chunk_y_offset"])
)
pil_image = Image.new(mode=face1["image"].mode, size=(new_im_width, new_im_height))
pil_image.paste(face1["image"], (face1_x_offset, face1_y_offset))
pil_image.paste(face2["image"], (face2_x_offset, face2_y_offset))
# Mask images are always from the origin
new_mask_im_width = max(face1["mask"].width, face2["mask"].width)
new_mask_im_height = max(face1["mask"].height, face2["mask"].height)
mask_pil = create_white_image(new_mask_im_width, new_mask_im_height)
black_image = create_black_image(face1["mask"].width, face1["mask"].height)
mask_pil.paste(black_image, (0, 0), ImageOps.invert(face1["mask"]))
black_image = create_black_image(face2["mask"].width, face2["mask"].height)
mask_pil.paste(black_image, (0, 0), ImageOps.invert(face2["mask"]))
new_face = FaceResultData(
image=pil_image,
mask=mask_pil,
x_center=max(face1["x_center"], face2["x_center"]),
y_center=max(face1["y_center"], face2["y_center"]),
mesh_width=max(face1["mesh_width"], face2["mesh_width"]),
mesh_height=max(face1["mesh_height"], face2["mesh_height"]),
chunk_x_offset=max(face1["chunk_x_offset"], face2["chunk_x_offset"]),
chunk_y_offset=max(face2["chunk_y_offset"], face2["chunk_y_offset"]),
)
return new_face
def prepare_faces_list( def prepare_faces_list(
face_result_list: list[FaceResultData], face_result_list: list[FaceResultData],
) -> list[FaceResultDataWithId]: ) -> list[FaceResultDataWithId]:
@ -91,7 +135,7 @@ def prepare_faces_list(
should_add = True should_add = True
candidate_x_center = candidate["x_center"] candidate_x_center = candidate["x_center"]
candidate_y_center = candidate["y_center"] candidate_y_center = candidate["y_center"]
for face in deduped_faces: for idx, face in enumerate(deduped_faces):
face_center_x = face["x_center"] face_center_x = face["x_center"]
face_center_y = face["y_center"] face_center_y = face["y_center"]
face_radius_w = face["mesh_width"] / 2 face_radius_w = face["mesh_width"] / 2
@ -105,6 +149,7 @@ def prepare_faces_list(
) )
if p < 1: # Inside of the already-added face's radius if p < 1: # Inside of the already-added face's radius
deduped_faces[idx] = coalesce_faces(face, candidate)
should_add = False should_add = False
break break
@ -138,7 +183,6 @@ def generate_face_box_mask(
chunk_x_offset: int = 0, chunk_x_offset: int = 0,
chunk_y_offset: int = 0, chunk_y_offset: int = 0,
draw_mesh: bool = True, draw_mesh: bool = True,
check_bounds: bool = True,
) -> list[FaceResultData]: ) -> list[FaceResultData]:
result = [] result = []
mask_pil = None mask_pil = None
@ -211,19 +255,6 @@ def generate_face_box_mask(
mask_pil = create_white_image(w + chunk_x_offset, h + chunk_y_offset) mask_pil = create_white_image(w + chunk_x_offset, h + chunk_y_offset)
mask_pil.paste(init_mask_pil, (chunk_x_offset, chunk_y_offset)) mask_pil.paste(init_mask_pil, (chunk_x_offset, chunk_y_offset))
left_side = x_center - mesh_width
right_side = x_center + mesh_width
top_side = y_center - mesh_height
bottom_side = y_center + mesh_height
im_width, im_height = pil_image.size
over_w = im_width * 0.1
over_h = im_height * 0.1
if not check_bounds or (
(left_side >= -over_w)
and (right_side < im_width + over_w)
and (top_side >= -over_h)
and (bottom_side < im_height + over_h)
):
x_center = float(x_center) x_center = float(x_center)
y_center = float(y_center) y_center = float(y_center)
face = FaceResultData( face = FaceResultData(
@ -233,11 +264,11 @@ def generate_face_box_mask(
y_center=y_center + chunk_y_offset, y_center=y_center + chunk_y_offset,
mesh_width=mesh_width, mesh_width=mesh_width,
mesh_height=mesh_height, mesh_height=mesh_height,
chunk_x_offset=chunk_x_offset,
chunk_y_offset=chunk_y_offset,
) )
result.append(face) result.append(face)
else:
context.services.logger.info("FaceTools --> Face out of bounds, ignoring.")
return result return result
@ -346,7 +377,6 @@ def get_faces_list(
chunk_x_offset=0, chunk_x_offset=0,
chunk_y_offset=0, chunk_y_offset=0,
draw_mesh=draw_mesh, draw_mesh=draw_mesh,
check_bounds=False,
) )
if should_chunk or len(result) == 0: if should_chunk or len(result) == 0:
context.services.logger.info("FaceTools --> Chunking image (chunk toggled on, or no face found in full image).") context.services.logger.info("FaceTools --> Chunking image (chunk toggled on, or no face found in full image).")
@ -360,24 +390,26 @@ def get_faces_list(
if width > height: if width > height:
# Landscape - slice the image horizontally # Landscape - slice the image horizontally
fx = 0.0 fx = 0.0
steps = int(width * 2 / height) steps = int(width * 2 / height) + 1
increment = (width - height) / (steps - 1)
while fx <= (width - height): while fx <= (width - height):
x = int(fx) x = int(fx)
image_chunks.append(image.crop((x, 0, x + height - 1, height - 1))) image_chunks.append(image.crop((x, 0, x + height, height)))
x_offsets.append(x) x_offsets.append(x)
y_offsets.append(0) y_offsets.append(0)
fx += (width - height) / steps fx += increment
context.services.logger.info(f"FaceTools --> Chunk starting at x = {x}") context.services.logger.info(f"FaceTools --> Chunk starting at x = {x}")
elif height > width: elif height > width:
# Portrait - slice the image vertically # Portrait - slice the image vertically
fy = 0.0 fy = 0.0
steps = int(height * 2 / width) steps = int(height * 2 / width) + 1
increment = (height - width) / (steps - 1)
while fy <= (height - width): while fy <= (height - width):
y = int(fy) y = int(fy)
image_chunks.append(image.crop((0, y, width - 1, y + width - 1))) image_chunks.append(image.crop((0, y, width, y + width)))
x_offsets.append(0) x_offsets.append(0)
y_offsets.append(y) y_offsets.append(y)
fy += (height - width) / steps fy += increment
context.services.logger.info(f"FaceTools --> Chunk starting at y = {y}") context.services.logger.info(f"FaceTools --> Chunk starting at y = {y}")
for idx in range(len(image_chunks)): for idx in range(len(image_chunks)):
@ -404,7 +436,7 @@ def get_faces_list(
return all_faces return all_faces
@invocation("face_off", title="FaceOff", tags=["image", "faceoff", "face", "mask"], category="image", version="1.0.1") @invocation("face_off", title="FaceOff", tags=["image", "faceoff", "face", "mask"], category="image", version="1.0.2")
class FaceOffInvocation(BaseInvocation): class FaceOffInvocation(BaseInvocation):
"""Bound, extract, and mask a face from an image using MediaPipe detection""" """Bound, extract, and mask a face from an image using MediaPipe detection"""
@ -498,7 +530,7 @@ class FaceOffInvocation(BaseInvocation):
return output return output
@invocation("face_mask_detection", title="FaceMask", tags=["image", "face", "mask"], category="image", version="1.0.1") @invocation("face_mask_detection", title="FaceMask", tags=["image", "face", "mask"], category="image", version="1.0.2")
class FaceMaskInvocation(BaseInvocation): class FaceMaskInvocation(BaseInvocation):
"""Face mask creation using mediapipe face detection""" """Face mask creation using mediapipe face detection"""
@ -616,7 +648,7 @@ class FaceMaskInvocation(BaseInvocation):
@invocation( @invocation(
"face_identifier", title="FaceIdentifier", tags=["image", "face", "identifier"], category="image", version="1.0.1" "face_identifier", title="FaceIdentifier", tags=["image", "face", "identifier"], category="image", version="1.0.2"
) )
class FaceIdentifierInvocation(BaseInvocation): class FaceIdentifierInvocation(BaseInvocation):
"""Outputs an image with detected face IDs printed on each face. For use with other FaceTools.""" """Outputs an image with detected face IDs printed on each face. For use with other FaceTools."""