convert face store (#1241)

tests/test_face_creator.py

@@ -2,10 +2,9 @@
 import numpy
 import pytest
 
-from facefusion import face_aligner, face_classifier, face_detector, face_recognizer, ffmpeg, ffmpeg_builder, inference_manager, process_manager, state_manager
+from facefusion import face_aligner, face_classifier, face_detector, face_recognizer, face_store, ffmpeg, ffmpeg_builder, inference_manager, process_manager, state_manager
 from facefusion.download import conditional_download
 from facefusion.face_creator import average_face_geometry, get_many_faces, get_one_face, refill_faces
-from facefusion.face_store import clear_faces
 from facefusion.vision import read_static_image
 from .assert_helper import get_test_example_file, get_test_examples_directory
 
@@ -15,6 +14,8 @@ def before_all() -> None:
 	state_manager.init()
 	inference_manager.init()
 
+	face_store.init()
+
 	process_manager.start()
 	conditional_download(get_test_examples_directory(),
 	[
@@ -56,7 +57,8 @@ def before_each() -> None:
 	face_detector.clear_inference_pool()
 	face_aligner.clear_inference_pool()
 	face_recognizer.clear_inference_pool()
-	clear_faces()
+
+	face_store.clear()
 
 
 def test_get_one_face() -> None:

tests/test_face_store.py

@@ -0,0 +1,74 @@
+from typing import Iterator
+
+import numpy
+import pytest
+
+from facefusion import session_context, session_manager
+from facefusion.face_store import clear, get_faces, init, resolve_lock, set_faces
+
+
+@pytest.fixture(scope = 'function', autouse = True)
+def before_each() -> Iterator[None]:
+	local_id = session_context.resolve_local_id()
+
+	session_context.set_session_id(local_id)
+	init()
+
+	yield
+
+	session_context.set_session_id(local_id)
+
+
+def test_init() -> None:
+	local_id = session_context.resolve_local_id()
+	vision_frame = numpy.zeros((16, 16, 3), numpy.uint8)
+
+	set_faces(vision_frame, [])
+	session_context.set_session_id('session-a')
+	init()
+
+	assert get_faces(vision_frame) is None
+
+	set_faces(vision_frame, [])
+	session_manager.fork_session()
+
+	assert get_faces(vision_frame) == []
+
+	session_manager.join_session()
+	session_context.set_session_id(local_id)
+
+	assert get_faces(vision_frame) == []
+
+
+def test_get_faces() -> None:
+	vision_frame = numpy.zeros((16, 16, 3), numpy.uint8)
+
+	assert get_faces(vision_frame) is None
+
+	set_faces(vision_frame, [])
+
+	assert get_faces(vision_frame) == []
+
+
+def test_set_faces() -> None:
+	vision_frame = numpy.zeros((16, 16, 3), numpy.uint8)
+
+	set_faces(vision_frame, [])
+	set_faces(vision_frame, [])
+
+	assert get_faces(vision_frame) == []
+
+
+def test_resolve_lock() -> None:
+	vision_frame = numpy.zeros((16, 16, 3), numpy.uint8)
+
+	assert resolve_lock(vision_frame) is resolve_lock(vision_frame)
+
+
+def test_clear() -> None:
+	vision_frame = numpy.zeros((16, 16, 3), numpy.uint8)
+
+	set_faces(vision_frame, [])
+	clear()
+
+	assert get_faces(vision_frame) is None

tests/test_face_tracker.py

@@ -1,11 +1,10 @@
 import numpy
 import pytest
 
-from facefusion import face_aligner, face_classifier, face_detector, face_recognizer, inference_manager, state_manager, video_manager
+from facefusion import face_aligner, face_classifier, face_detector, face_recognizer, face_store, inference_manager, state_manager, video_manager
 from facefusion.common_helper import get_first, get_last
 from facefusion.download import conditional_download
 from facefusion.face_creator import get_many_faces, get_one_face
-from facefusion.face_store import clear_faces
 from facefusion.face_tracker import create_face_tracks, select_face_track, track_faces
 from facefusion.vision import read_static_video_frame, select_video_frames
 from .assert_helper import get_test_example_file, get_test_examples_directory
@@ -16,6 +15,8 @@ def before_all() -> None:
 	state_manager.init()
 	inference_manager.init()
 
+	face_store.init()
+
 	video_manager.init()
 
 	conditional_download(get_test_examples_directory(),
@@ -47,7 +48,8 @@ def before_each() -> None:
 	face_detector.clear_inference_pool()
 	face_aligner.clear_inference_pool()
 	face_recognizer.clear_inference_pool()
-	clear_faces()
+
+	face_store.clear()
 
 
 def test_track_faces() -> None: