From fcb1ed5799bcab9ba0ac4fc441411cafb1d5b753 Mon Sep 17 00:00:00 2001 From: Pozitronik Date: Mon, 30 Sep 2024 13:36:11 +0400 Subject: [PATCH] issue #118: test for new feature --- tests/processors/test_face_swapper.py | 66 ++++++++++++++++++++++++--- 1 file changed, 60 insertions(+), 6 deletions(-) diff --git a/tests/processors/test_face_swapper.py b/tests/processors/test_face_swapper.py index b015182e..627be70d 100644 --- a/tests/processors/test_face_swapper.py +++ b/tests/processors/test_face_swapper.py @@ -1,6 +1,7 @@ import multiprocessing -from argparse import Namespace +import numpy as np +import pytest from colorama import Fore, Back from insightface.app.common import Face @@ -9,13 +10,12 @@ from sinner.helpers.FrameHelper import read_from_image from sinner.processors.frame.FaceSwapper import FaceSwapper from sinner.typing import Frame, FaceSwapperType -from tests.constants import source_jpg, target_png, IMAGE_SHAPE, tmp_dir, no_face_jpg +from tests.constants import source_jpg, target_png, IMAGE_SHAPE, tmp_dir, no_face_jpg, male_face_jpg, female_face_jpg, multiple_faces_jpg, unknown_gender_face_jpg -parameters: Namespace = Parameters(f'--execution-provider=cpu --execution-threads={multiprocessing.cpu_count()} --max-memory=12 --source-path="{source_jpg}" --target-path="{target_png}" --output-path="{tmp_dir}"').parameters - -def get_test_object() -> FaceSwapper: - return FaceSwapper(parameters=parameters) +def get_test_object(additional_params: str = "") -> FaceSwapper: + params = f'--execution-provider=cpu --execution-threads={multiprocessing.cpu_count()} --max-memory=12 --source-path="{source_jpg}" --target-path="{target_png}" --output-path="{tmp_dir}" {additional_params}' + return FaceSwapper(parameters=Parameters(params).parameters) def test_init(): @@ -44,3 +44,57 @@ def test_process_frame(): processed_frame = get_test_object().process_frame(read_from_image(target_png)) assert (processed_frame, Frame) assert processed_frame.shape == IMAGE_SHAPE + + +@pytest.mark.parametrize("target_gender, source_image, target_image, should_swap", [ + ("M", male_face_jpg, male_face_jpg, True), + ("M", male_face_jpg, female_face_jpg, False), + ("F", female_face_jpg, female_face_jpg, True), + ("F", female_face_jpg, male_face_jpg, False), + ("B", male_face_jpg, female_face_jpg, True), + ("B", female_face_jpg, male_face_jpg, True), + ("I", male_face_jpg, male_face_jpg, True), + ("I", male_face_jpg, female_face_jpg, False), + ("I", female_face_jpg, female_face_jpg, True), + ("I", female_face_jpg, male_face_jpg, False), +]) +def test_gender_selection(target_gender, source_image, target_image, should_swap): + test_object = get_test_object(f'--target-gender={target_gender} --source-path="{source_image}"') + target_frame = read_from_image(target_image) + processed_frame = test_object.process_frame(target_frame) + + # Анализируем лица на обработанном кадре + original_face = test_object.face_analyser.get_one_face(target_frame) + processed_face = test_object.face_analyser.get_one_face(processed_frame) + + if should_swap: + if source_image != target_image: + # Проверяем, что лицо изменилось + assert not np.allclose(original_face.embedding, processed_face.embedding, atol=1e-5) + # Проверяем, что пол лица соответствует ожидаемому + if target_gender in ['M', 'F']: + assert processed_face.sex == target_gender + elif target_gender == 'I': + source_face = test_object.source_face + assert processed_face.sex == source_face.sex + else: + # Проверяем, что лицо не изменилось + assert np.allclose(original_face.embedding, processed_face.embedding, atol=1e-5) + assert original_face.sex == processed_face.sex + + +# test is skipped, because I cant find a good image for it +# def test_unknown_gender_warning(capsys): +# test_object = get_test_object(f'--target-gender=M --source-path="{unknown_gender_face_jpg}"') +# target_frame = read_from_image(target_png) +# _ = test_object.process_frame(target_frame) +# captured = capsys.readouterr() +# assert "Unable to determine gender for a face. Skipping this face." in captured.out + + +def test_many_faces_processing(): + test_object = get_test_object("--many-faces") + assert test_object.many_faces is True + target_frame = read_from_image(multiple_faces_jpg) + processed_frame = test_object.process_frame(target_frame) + assert not np.array_equal(target_frame, processed_frame)