add face tools to assets

facefusion/apis/asset_helper.py

@@ -7,8 +7,9 @@ from starlette.datastructures import UploadFile
 
 import facefusion.choices
 from facefusion import ffmpeg, process_manager, state_manager
+from facefusion.face_creator import get_many_faces
 from facefusion.filesystem import create_directory, get_file_extension, get_file_format, is_audio, is_image, is_video
-from facefusion.types import ImageAsset, ImageMetadata, MediaType, VideoAsset, VisionFrame
+from facefusion.types import Face, ImageAsset, ImageMetadata, MediaType, VideoAsset, VisionFrame
 from facefusion.vision import detect_image_resolution, fit_contain_frame, is_vision_frame, read_static_image, read_static_video_frame, unpack_resolution
 
 
@@ -120,3 +121,30 @@ def capture_asset_frames(asset : ImageAsset | VideoAsset, frame_numbers : List[s
 		capture_vision_frames.append(capture_vision_frame)
 
 	return capture_vision_frames
+
+
+def crop_face_frame(vision_frame : VisionFrame, face : Face) -> VisionFrame:
+	bounding_box = face.bounding_box.astype(int)
+	start_x = max(0, bounding_box[0])
+	start_y = max(0, bounding_box[1])
+	end_x = min(vision_frame.shape[1], bounding_box[2])
+	end_y = min(vision_frame.shape[0], bounding_box[3])
+	crop_vision_frame = vision_frame[start_y:end_y, start_x:end_x]
+	return crop_vision_frame
+
+
+def capture_asset_faces(asset : ImageAsset | VideoAsset, frame_numbers : List[str], resolution : str) -> List[VisionFrame]:
+	capture_vision_frames = []
+	temp_vision_frames = read_asset_frames(asset, frame_numbers)
+
+	for temp_vision_frame in temp_vision_frames:
+		faces = get_many_faces([ temp_vision_frame ])
+
+		for face in faces:
+			crop_vision_frame = crop_face_frame(temp_vision_frame, face)
+
+			if crop_vision_frame.shape[0] and crop_vision_frame.shape[1]:
+				capture_vision_frame = fit_contain_frame(crop_vision_frame, unpack_resolution(resolution))
+				capture_vision_frames.append(capture_vision_frame)
+
+	return capture_vision_frames

facefusion/apis/endpoints/assets.py

@@ -7,7 +7,7 @@ from starlette.status import HTTP_200_OK, HTTP_201_CREATED, HTTP_400_BAD_REQUEST
 
 from facefusion import session_context, session_manager
 from facefusion.apis import asset_store
-from facefusion.apis.asset_helper import capture_asset_frames, save_asset_files, validate_asset_files
+from facefusion.apis.asset_helper import capture_asset_faces, capture_asset_frames, save_asset_files, validate_asset_files
 from facefusion.apis.session_helper import extract_access_token
 from facefusion.filesystem import remove_file
 from facefusion.vision import is_vision_frames, to_strip_buffer
@@ -90,12 +90,18 @@ async def get_asset(request : Request) -> Response:
 		asset = asset_store.get_asset(session_id, asset_id)
 
 		if asset:
-			if asset.get('media') in [ 'image', 'video' ] and request.query_params.get('action') == 'capture' and request.query_params.get('subject') == 'frame':
+			if asset.get('media') in [ 'image', 'video' ] and request.query_params.get('action') == 'capture':
 				resolution = request.query_params.get('resolution')
 				frame_numbers = request.query_params.getlist('frame_number')
-				vision_frames = capture_asset_frames(asset, frame_numbers, resolution) #type:ignore[arg-type]
+				vision_frames = []
 
-				if is_vision_frames(vision_frames):
+				if request.query_params.get('subject') == 'frame':
+					vision_frames = capture_asset_frames(asset, frame_numbers, resolution) #type:ignore[arg-type]
+
+				if request.query_params.get('subject') == 'face':
+					vision_frames = capture_asset_faces(asset, frame_numbers, resolution) #type:ignore[arg-type]
+
+				if vision_frames and is_vision_frames(vision_frames):
 					return Response(content = to_strip_buffer(vision_frames), media_type = 'image/jpeg')
 
 				return Response(status_code = HTTP_400_BAD_REQUEST)

facefusion/types.py

@@ -315,7 +315,7 @@ AssetSet : TypeAlias = Dict[AssetId, AudioAsset | ImageAsset | VideoAsset]
 AssetStore : TypeAlias = Dict[SessionId, AssetSet]
 
 AssetAction = Literal['capture']
-AssetSubject = Literal['frame']
+AssetSubject = Literal['frame', 'face']
 
 BenchmarkMode = Literal['warm', 'cold']
 BenchmarkResolution = Literal['240p', '360p', '540p', '720p', '1080p', '1440p', '2160p']