issue #511

Images reading and writing refactored

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:

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:

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)

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)