diff --git a/roop/processors/frame/face_enhancer.py b/roop/processors/frame/face_enhancer.py index c07d049..0eb3930 100644 --- a/roop/processors/frame/face_enhancer.py +++ b/roop/processors/frame/face_enhancer.py @@ -1,14 +1,12 @@ from typing import Any, List -import cv2 import threading import gfpgan -import numpy import roop.globals import roop.processors.frame.core from roop.core import update_status from roop.face_analyser import get_one_face -from roop.utilities import conditional_download, resolve_relative_path, is_image, is_video +from roop.utilities import conditional_download, resolve_relative_path, is_image, is_video, read_image, write_image FACE_ENHANCER = None THREAD_SEMAPHORE = threading.Semaphore() @@ -58,19 +56,17 @@ def process_frame(source_face: Any, temp_frame: Any) -> Any: def process_frames(source_path: str, temp_frame_paths: List[str], progress=None) -> None: for temp_frame_path in temp_frame_paths: - temp_frame = cv2.imdecode(numpy.fromfile(temp_frame_path, dtype=numpy.uint8), cv2.IMREAD_UNCHANGED) + temp_frame = read_image(temp_frame_path) result = process_frame(None, temp_frame) - is_success, im_buf_arr = cv2.imencode(".png", result) - im_buf_arr.tofile(temp_frame_path) + write_image(result, temp_frame_path) if progress: progress.update(1) def process_image(source_path: str, target_path: str, output_path: str) -> None: - target_frame = cv2.imdecode(numpy.fromfile(target_path, dtype=numpy.uint8), cv2.IMREAD_UNCHANGED) + target_frame = read_image(target_path) result = process_frame(None, target_frame) - is_success, im_buf_arr = cv2.imencode(".png", result) - im_buf_arr.tofile(output_path) + write_image(result, output_path) def process_video(source_path: str, temp_frame_paths: List[str]) -> None: diff --git a/roop/processors/frame/face_swapper.py b/roop/processors/frame/face_swapper.py index d0f0325..18af0dc 100644 --- a/roop/processors/frame/face_swapper.py +++ b/roop/processors/frame/face_swapper.py @@ -1,15 +1,12 @@ from typing import Any, List -import cv2 import insightface import threading -import numpy - import roop.globals import roop.processors.frame.core from roop.core import update_status from roop.face_analyser import get_one_face, get_many_faces -from roop.utilities import conditional_download, resolve_relative_path, is_image, is_video +from roop.utilities import conditional_download, resolve_relative_path, is_image, is_video, read_image, write_image FACE_SWAPPER = None THREAD_LOCK = threading.Lock() @@ -26,7 +23,7 @@ def pre_start() -> bool: if not is_image(roop.globals.source_path): update_status('Select an image for source path.', NAME) return False - elif not get_one_face(cv2.imdecode(numpy.fromfile(roop.globals.source_path, dtype=numpy.uint8), cv2.IMREAD_UNCHANGED)): + elif not get_one_face(read_image(roop.globals.source_path)): update_status('No face in source path detected.', NAME) return False if not is_image(roop.globals.target_path) and not is_video(roop.globals.target_path): @@ -63,13 +60,12 @@ def process_frame(source_face: Any, temp_frame: Any) -> Any: def process_frames(source_path: str, temp_frame_paths: List[str], progress=None) -> None: - source_face = get_one_face(cv2.imdecode(numpy.fromfile(source_path, dtype=numpy.uint8), cv2.IMREAD_UNCHANGED)) + source_face = get_one_face(read_image(source_path)) for temp_frame_path in temp_frame_paths: - temp_frame = cv2.imdecode(numpy.fromfile(temp_frame_path, dtype=numpy.uint8), cv2.IMREAD_UNCHANGED) + temp_frame = read_image(temp_frame_path) try: result = process_frame(source_face, temp_frame) - is_success, im_buf_arr = cv2.imencode(".png", result) - im_buf_arr.tofile(temp_frame_path) + write_image(result, temp_frame_path) except Exception as exception: print(exception) pass @@ -78,11 +74,10 @@ def process_frames(source_path: str, temp_frame_paths: List[str], progress=None) def process_image(source_path: str, target_path: str, output_path: str) -> None: - source_face = get_one_face(cv2.imdecode(numpy.fromfile(source_path, dtype=numpy.uint8), cv2.IMREAD_UNCHANGED)) - target_frame = cv2.imdecode(numpy.fromfile(target_path, dtype=numpy.uint8), cv2.IMREAD_UNCHANGED) + source_face = get_one_face(read_image(source_path)) + target_frame = read_image(target_path) result = process_frame(source_face, target_frame) - is_success, im_buf_arr = cv2.imencode(".png", result) - im_buf_arr.tofile(output_path) + write_image(result, output_path) def process_video(source_path: str, temp_frame_paths: List[str]) -> None: diff --git a/roop/ui.py b/roop/ui.py index 2eec032..3220582 100644 --- a/roop/ui.py +++ b/roop/ui.py @@ -11,7 +11,7 @@ from roop.face_analyser import get_one_face from roop.capturer import get_video_frame, get_video_frame_total from roop.predicter import predict_frame from roop.processors.frame.core import get_frame_processors_modules -from roop.utilities import is_image, is_video, resolve_relative_path +from roop.utilities import is_image, is_video, resolve_relative_path, read_image WINDOW_HEIGHT = 700 WINDOW_WIDTH = 600 @@ -205,10 +205,7 @@ def update_preview(frame_number: int = 0) -> None: if predict_frame(temp_frame): quit() for frame_processor in get_frame_processors_modules(roop.globals.frame_processors): - temp_frame = frame_processor.process_frame( - get_one_face(cv2.imdecode(numpy.fromfile(roop.globals.source_path, dtype=numpy.uint8), cv2.IMREAD_UNCHANGED)), - temp_frame - ) + temp_frame = frame_processor.process_frame(get_one_face(read_image(roop.globals.source_path)), temp_frame) image = Image.fromarray(cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB)) image = ImageOps.contain(image, (PREVIEW_MAX_WIDTH, PREVIEW_MAX_HEIGHT), Image.LANCZOS) image = ctk.CTkImage(image, size=image.size) diff --git a/roop/utilities.py b/roop/utilities.py index f97e175..4d9f8d6 100644 --- a/roop/utilities.py +++ b/roop/utilities.py @@ -8,6 +8,9 @@ import subprocess import urllib from pathlib import Path from typing import List + +import cv2 +from numpy import array, uint8, fromfile from tqdm import tqdm import roop.globals @@ -139,3 +142,19 @@ def conditional_download(download_directory_path: str, urls: List[str]) -> None: def resolve_relative_path(path: str) -> str: return os.path.abspath(os.path.join(os.path.dirname(__file__), path)) + + +def read_image(path: str) -> array: + if platform.system().lower() == 'windows': # issue #511 + return cv2.imdecode(fromfile(path, dtype=uint8), cv2.IMREAD_UNCHANGED) + else: + return cv2.imread(path) + + +def write_image(image: array, path: str) -> bool: + if platform.system().lower() == 'windows': # issue #511 + is_success, im_buf_arr = cv2.imencode(".png", image) + im_buf_arr.tofile(path) + return is_success + else: + return cv2.imwrite(path, array) \ No newline at end of file