fix(face-mask): guard create_face_mask against None frame

This commit is contained in:
gujishh
2026-04-12 14:19:48 +09:00
parent 8703d394d6
commit fbcea9e135
2 changed files with 111 additions and 1 deletions
+4 -1
View File
@@ -1028,10 +1028,13 @@ def apply_mouth_area(
def create_face_mask(face: Face, frame: Frame) -> np.ndarray:
"""Creates a feathered mask covering the whole face area based on landmarks."""
if frame is None or not hasattr(frame, "shape") or len(frame.shape) < 2:
return np.zeros((0, 0), dtype=np.uint8)
mask = np.zeros(frame.shape[:2], dtype=np.uint8) # Start with uint8
# Validate inputs
if face is None or not hasattr(face, 'landmark_2d_106') or frame is None:
if face is None or not hasattr(face, 'landmark_2d_106'):
# print("Warning: Invalid face or frame for create_face_mask.")
return mask # Return empty mask
+107
View File
@@ -0,0 +1,107 @@
from __future__ import annotations
import importlib
import sys
import types
import unittest
from unittest import mock
def _load_face_swapper_module():
fake_numpy = types.ModuleType("numpy")
class FakeArray:
def __init__(self, shape, dtype=int, fill_value=0) -> None:
self.shape = shape
self.dtype = dtype
self.fill_value = fill_value
size = 1
for dim in shape:
size *= dim
self.size = size
def fake_zeros(shape, dtype=int):
return FakeArray(shape, dtype=dtype, fill_value=0)
setattr(fake_numpy, "ndarray", FakeArray)
setattr(fake_numpy, "zeros", fake_zeros)
setattr(fake_numpy, "uint8", int)
fake_cv2 = types.ModuleType("cv2")
setattr(fake_cv2, "IMREAD_COLOR", 1)
setattr(fake_cv2, "imdecode", lambda *args, **kwargs: None)
setattr(
fake_cv2,
"imencode",
lambda *args, **kwargs: (True, types.SimpleNamespace(tofile=lambda *_: None)),
)
fake_insightface = types.ModuleType("insightface")
fake_globals = types.ModuleType("modules.globals")
fake_globals.execution_providers = []
fake_globals.face_mask_blur = 31
fake_core = types.ModuleType("modules.processors.frame.core")
fake_core.process_video = lambda *args, **kwargs: None
fake_core_module = types.ModuleType("modules.core")
fake_core_module.update_status = lambda *args, **kwargs: None
fake_face_analyser = types.ModuleType("modules.face_analyser")
fake_face_analyser.get_one_face = lambda *args, **kwargs: None
fake_face_analyser.get_many_faces = lambda *args, **kwargs: []
fake_face_analyser.default_source_face = lambda *args, **kwargs: None
fake_typing = types.ModuleType("modules.typing")
fake_typing.Face = object
fake_typing.Frame = object
fake_utilities = types.ModuleType("modules.utilities")
fake_utilities.conditional_download = lambda *args, **kwargs: None
fake_utilities.is_image = lambda *args, **kwargs: False
fake_utilities.is_video = lambda *args, **kwargs: False
fake_cluster = types.ModuleType("modules.cluster_analysis")
fake_cluster.find_closest_centroid = lambda *args, **kwargs: None
fake_gpu = types.ModuleType("modules.gpu_processing")
fake_gpu.gpu_gaussian_blur = lambda img, *args, **kwargs: img
fake_gpu.gpu_sharpen = lambda img, *args, **kwargs: img
fake_gpu.gpu_add_weighted = lambda src1, *args, **kwargs: src1
fake_gpu.gpu_resize = lambda img, *args, **kwargs: img
fake_gpu.gpu_cvt_color = lambda img, *args, **kwargs: img
with mock.patch.dict(
sys.modules,
{
"cv2": fake_cv2,
"insightface": fake_insightface,
"numpy": fake_numpy,
"modules.globals": fake_globals,
"modules.processors.frame.core": fake_core,
"modules.core": fake_core_module,
"modules.face_analyser": fake_face_analyser,
"modules.typing": fake_typing,
"modules.utilities": fake_utilities,
"modules.cluster_analysis": fake_cluster,
"modules.gpu_processing": fake_gpu,
},
clear=False,
):
return importlib.import_module("modules.processors.frame.face_swapper")
class CreateFaceMaskTests(unittest.TestCase):
def test_create_face_mask_returns_empty_mask_when_frame_is_none(self) -> None:
face_swapper = _load_face_swapper_module()
class DummyFace:
landmark_2d_106 = []
result = face_swapper.create_face_mask(DummyFace(), None)
self.assertEqual(result.shape, (0, 0))
self.assertEqual(result.dtype, int)
if __name__ == "__main__":
unittest.main()