diff --git a/facefusion/apis/endpoints/process.py b/facefusion/apis/endpoints/process.py index c974615d..821d32b9 100644 --- a/facefusion/apis/endpoints/process.py +++ b/facefusion/apis/endpoints/process.py @@ -1,48 +1,11 @@ -import os -import uuid -from time import time -from typing import Optional - +import cv2 +import numpy from starlette.websockets import WebSocket, WebSocketDisconnect -from facefusion import ffmpeg, process_manager, session_context, session_manager, state_manager +from facefusion import session_context, session_manager, state_manager from facefusion.apis.api_helper import get_sec_websocket_protocol from facefusion.apis.endpoints.session import extract_access_token -from facefusion.filesystem import create_directory, remove_file -from facefusion.workflows import image_to_image - - -def save_image_bytes(image_bytes : bytes, temp_path : str, file_extension : str) -> str: - raw_path = os.path.join(temp_path, str(uuid.uuid4()) + file_extension) - - with open(raw_path, 'wb') as raw_file: - raw_file.write(image_bytes) - - return raw_path - - -def sanitize_upload(raw_path : str, temp_path : str, file_extension : str) -> Optional[str]: - target_path = os.path.join(temp_path, str(uuid.uuid4()) + file_extension) - - process_manager.start() - sanitized = ffmpeg.sanitize_image(raw_path, target_path) - process_manager.end() - remove_file(raw_path) - - if sanitized: - return target_path - return None - - -def process_target(target_path : str, output_path : str) -> bool: - state_manager.set_item('target_path', target_path) - state_manager.set_item('output_path', output_path) - state_manager.set_item('workflow', 'image-to-image') - - error_code = image_to_image.process(time()) - remove_file(target_path) - - return error_code == 0 +from facefusion.streamer import process_stream_frame async def websocket_process_image(websocket : WebSocket) -> None: @@ -56,23 +19,17 @@ async def websocket_process_image(websocket : WebSocket) -> None: await websocket.accept(subprotocol = subprotocol) if source_paths: - temp_path = state_manager.get_temp_path() - create_directory(temp_path) - try: while True: image_bytes = await websocket.receive_bytes() - raw_path = save_image_bytes(image_bytes, temp_path, '.jpg') - target_path = sanitize_upload(raw_path, temp_path, '.jpg') + target_vision_frame = cv2.imdecode(numpy.frombuffer(image_bytes, numpy.uint8), cv2.IMREAD_COLOR) - if target_path: - output_path = os.path.join(temp_path, str(uuid.uuid4()) + '.jpg') + if numpy.any(target_vision_frame): + temp_vision_frame = process_stream_frame(target_vision_frame) + success, result_bytes = cv2.imencode('.jpg', temp_vision_frame) - if process_target(target_path, output_path): - with open(output_path, 'rb') as output_file: - await websocket.send_bytes(output_file.read()) - - remove_file(output_path) + if success: + await websocket.send_bytes(result_bytes.tobytes()) except (WebSocketDisconnect, OSError): pass