From c73d4bbf8b5d47334484a062e30c84f7a2bb40ef Mon Sep 17 00:00:00 2001 From: Pozitronik Date: Wed, 26 Jul 2023 00:52:10 +0400 Subject: [PATCH 1/6] refactoring code --- sinner/Core.py | 21 ++++++++++++++++++--- sinner/utilities.py | 9 --------- 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/sinner/Core.py b/sinner/Core.py index 31ca89fc..deeb0ad8 100644 --- a/sinner/Core.py +++ b/sinner/Core.py @@ -14,7 +14,7 @@ from sinner.processors.frame.BaseFrameProcessor import BaseFrameProcessor from sinner.State import State from sinner.typing import Frame -from sinner.utilities import is_image, is_video, delete_subdirectories, list_class_descendants, resolve_relative_path, get_app_dir, TEMP_DIRECTORY, suggest_output_path +from sinner.utilities import is_image, is_video, delete_subdirectories, list_class_descendants, resolve_relative_path, get_app_dir, TEMP_DIRECTORY from sinner.validators.AttributeLoader import AttributeLoader, Rules # single thread doubles cuda performance - needs to be set before torch import @@ -33,6 +33,7 @@ class Core(AttributeLoader, Status): target_path: str + output_path: str | None frame_processor: List[str] frame_handler: str temp_dir: str @@ -53,6 +54,12 @@ def rules(self) -> Rules: 'required': True, 'help': 'Select the target file or the directory' }, + { + 'parameter': {'output', 'output-path'}, + 'attribute': 'output_path', + 'default:': lambda: self.suggest_output_path(), + 'default': None, + }, { 'parameter': 'frame-processor', 'default': ['FaceSwapper'], @@ -114,7 +121,7 @@ def run(self, set_progress: Callable[[int], None] | None = None) -> None: temp_resources.append(state.in_dir) if temp_resources is not []: - output_filename = current_processor.output_path if current_processor is not None else suggest_output_path(self.target_path) + output_filename = current_processor.output_path if current_processor is not None else self.output_path final_handler = BaseFrameHandler.create(handler_name=self.frame_handler, parameters=self.parameters, target_path=self.target_path) if final_handler.result(from_dir=current_target_path, filename=output_filename, audio_target=self.target_path) is True: if self.keep_frames is False: @@ -157,10 +164,18 @@ def get_frame(self, frame_number: int = 0, extractor_handler: BaseFrameHandler | self.preview_processors[processor_name].load(self.parameters) frame = self.preview_processors[processor_name].process_frame(frame) result.append((frame, processor_name)) - except Exception as exception: # skip, if parameters is not enough for processors + except Exception as exception: # skip, if parameters is not enough for processor self.update_status(message=str(exception), mood=Mood.BAD) pass return result def stop(self) -> None: self._stop_flag = True + + def suggest_output_path(self) -> str: + target_name, target_extension = os.path.splitext(os.path.basename(self.target_path)) + if self.output_path is None: + return os.path.join(os.path.dirname(self.target_path), 'result-' + target_name + target_extension) + if os.path.isdir(self.output_path): + return os.path.join(self.output_path, 'result-' + target_name + target_extension) + return self.output_path diff --git a/sinner/utilities.py b/sinner/utilities.py index eadfcc75..d882114b 100644 --- a/sinner/utilities.py +++ b/sinner/utilities.py @@ -209,12 +209,3 @@ def declared_attr_type(obj: object, attribute: str) -> Any: if attribute in declared_typed_variables: return declared_typed_variables[attribute] return None - - -def suggest_output_path(target_path: str, output_path: str | None = None) -> str: - target_name, target_extension = os.path.splitext(os.path.basename(target_path)) - if output_path is None: - return os.path.join(os.path.dirname(target_path), 'result-' + target_name + target_extension) - if os.path.isdir(output_path): - return os.path.join(output_path, 'result-' + target_name + target_extension) - return output_path From 8460dc5d97d0dce1dc649dff6b41bd7aae4b2d33 Mon Sep 17 00:00:00 2001 From: Pozitronik Date: Wed, 26 Jul 2023 00:56:38 +0400 Subject: [PATCH 2/6] comment added --- sinner/handlers/frame/FFmpegVideoHandler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sinner/handlers/frame/FFmpegVideoHandler.py b/sinner/handlers/frame/FFmpegVideoHandler.py index ac2d7298..96bbc300 100644 --- a/sinner/handlers/frame/FFmpegVideoHandler.py +++ b/sinner/handlers/frame/FFmpegVideoHandler.py @@ -62,7 +62,7 @@ def detect_fps(self) -> float: def detect_fc(self) -> int: try: command = ['ffprobe', '-v', 'error', '-count_frames', '-select_streams', 'v:0', '-show_entries', 'stream=nb_frames', '-of', 'default=nokey=1:noprint_wrappers=1', self._target_path] - output = subprocess.check_output(command, stderr=subprocess.STDOUT).decode('utf-8').strip() + output = subprocess.check_output(command, stderr=subprocess.STDOUT).decode('utf-8').strip() # can be very slow! if 'N/A' == output: return 1 # non-frame files, still processable return int(output) From 3024dee704139b36041595a92f7814f2f9d52847 Mon Sep 17 00:00:00 2001 From: Pozitronik Date: Wed, 26 Jul 2023 01:24:02 +0400 Subject: [PATCH 3/6] fixed: CV2VideoHandler can't handle non-latin symbols in filenames fixed methods moved from utils to handler as static --- sinner/State.py | 4 +-- sinner/handlers/frame/CV2VideoHandler.py | 31 ++++++++++++++++--- sinner/handlers/frame/DirectoryHandler.py | 5 +-- sinner/handlers/frame/ImageHandler.py | 5 +-- sinner/processors/frame/BaseFrameProcessor.py | 5 +-- sinner/processors/frame/FaceSwapper.py | 5 +-- sinner/utilities.py | 21 ------------- 7 files changed, 40 insertions(+), 36 deletions(-) diff --git a/sinner/State.py b/sinner/State.py index c924a7ab..1b222aa3 100644 --- a/sinner/State.py +++ b/sinner/State.py @@ -4,8 +4,8 @@ from typing import Any, Dict, List from sinner.Status import Status, Mood +from sinner.handlers.frame.CV2VideoHandler import CV2VideoHandler from sinner.typing import Frame -from sinner.utilities import write_image from sinner.validators.AttributeLoader import AttributeLoader, Rules OUT_DIR = 'OUT' @@ -113,7 +113,7 @@ def state_path(self, dir_type: str) -> str: return os.path.join(self.temp_dir, *sub_path) def save_temp_frame(self, frame: Frame, index: int) -> None: - if not write_image(frame, self.get_frame_processed_name(index)): + if not CV2VideoHandler.write_image(frame, self.get_frame_processed_name(index)): raise Exception(f"Error saving frame: {self.get_frame_processed_name(index)}") # Checks if some frame already processed diff --git a/sinner/handlers/frame/CV2VideoHandler.py b/sinner/handlers/frame/CV2VideoHandler.py index c5a51b4d..30b4233c 100644 --- a/sinner/handlers/frame/CV2VideoHandler.py +++ b/sinner/handlers/frame/CV2VideoHandler.py @@ -1,16 +1,18 @@ import glob import os.path +import platform from pathlib import Path from typing import List import cv2 from cv2 import VideoCapture +from numpy import fromfile, uint8 from tqdm import tqdm from sinner.Status import Mood from sinner.handlers.frame.BaseFrameHandler import BaseFrameHandler -from sinner.typing import NumeratedFrame, NumeratedFramePath -from sinner.utilities import write_image, get_file_name +from sinner.typing import NumeratedFrame, NumeratedFramePath, Frame +from sinner.utilities import get_file_name from sinner.validators.AttributeLoader import Rules @@ -87,7 +89,7 @@ def get_frames_paths(self, path: str) -> List[NumeratedFramePath]: ret, frame = capture.read() if not ret: break - write_image(frame, os.path.join(path, str(i + 1).zfill(filename_length) + ".png")) + self.write_image(frame, os.path.join(path, str(i + 1).zfill(filename_length) + ".png")) progress.update() i += 1 capture.release() @@ -111,12 +113,12 @@ def result(self, from_dir: str, filename: str, audio_target: str | None = None) try: Path(os.path.dirname(filename)).mkdir(parents=True, exist_ok=True) frame_files = glob.glob(os.path.join(glob.escape(from_dir), '*.png')) - first_frame = cv2.imread(frame_files[0]) + first_frame = self.read_image(frame_files[0]) height, width, channels = first_frame.shape fourcc = self.suggest_codec() video_writer = cv2.VideoWriter(filename, fourcc, self.output_fps, (width, height)) for frame_path in frame_files: - frame = cv2.imread(frame_path) + frame = self.read_image(frame_path) video_writer.write(frame) video_writer.release() return True @@ -132,3 +134,22 @@ def suggest_codec(self) -> int: self.update_status(message=f"Suggested codec: {fourcc}", mood=Mood.NEUTRAL) return fourcc raise NotImplementedError('No supported codecs found') + + @staticmethod + def read_image(path: str) -> Frame: + if platform.system().lower() == 'windows': # issue #511 + image = cv2.imdecode(fromfile(path, dtype=uint8), cv2.IMREAD_UNCHANGED) + if image.shape[2] == 4: # fixes the alpha-channel issue + image = cv2.cvtColor(image, cv2.COLOR_BGRA2BGR) + return image + else: + return cv2.imread(path) + + @staticmethod + def write_image(image: Frame, 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, image) diff --git a/sinner/handlers/frame/DirectoryHandler.py b/sinner/handlers/frame/DirectoryHandler.py index b3070dd2..2982847a 100644 --- a/sinner/handlers/frame/DirectoryHandler.py +++ b/sinner/handlers/frame/DirectoryHandler.py @@ -5,8 +5,9 @@ from typing import List from sinner.handlers.frame.BaseFrameHandler import BaseFrameHandler +from sinner.handlers.frame.CV2VideoHandler import CV2VideoHandler from sinner.typing import NumeratedFrame, NumeratedFramePath -from sinner.utilities import read_image, is_image, get_file_name +from sinner.utilities import is_image, get_file_name class DirectoryHandler(BaseFrameHandler): @@ -27,7 +28,7 @@ def get_frames_paths(self, path: str) -> List[NumeratedFramePath]: return [(int(get_file_name(file_path)), file_path) for file_path in frames_path if is_image(file_path)] def extract_frame(self, frame_number: int) -> NumeratedFrame: - return frame_number, read_image(self.get_frames_paths(self._target_path)[frame_number - 1][1]) # zero-based sorted frames list + return frame_number, CV2VideoHandler.read_image(self.get_frames_paths(self._target_path)[frame_number - 1][1]) # zero-based sorted frames list def result(self, from_dir: str, filename: str, audio_target: str | None = None) -> bool: self.update_status(f"Copying results from {from_dir} to {filename}") diff --git a/sinner/handlers/frame/ImageHandler.py b/sinner/handlers/frame/ImageHandler.py index ddad9dd6..6ea67af6 100644 --- a/sinner/handlers/frame/ImageHandler.py +++ b/sinner/handlers/frame/ImageHandler.py @@ -6,8 +6,9 @@ from sinner.Status import Mood from sinner.handlers.frame.BaseFrameHandler import BaseFrameHandler +from sinner.handlers.frame.CV2VideoHandler import CV2VideoHandler from sinner.typing import NumeratedFrame, NumeratedFramePath -from sinner.utilities import read_image, is_image +from sinner.utilities import is_image class ImageHandler(BaseFrameHandler): @@ -27,7 +28,7 @@ def get_frames_paths(self, path: str) -> List[NumeratedFramePath]: return [(1, self._target_path)] def extract_frame(self, frame_number: int) -> NumeratedFrame: - return frame_number, read_image(self._target_path) + return frame_number, CV2VideoHandler.read_image(self._target_path) def result(self, from_dir: str, filename: str, audio_target: str | None = None) -> bool: try: diff --git a/sinner/processors/frame/BaseFrameProcessor.py b/sinner/processors/frame/BaseFrameProcessor.py index 0eb180d6..c72bf9d5 100644 --- a/sinner/processors/frame/BaseFrameProcessor.py +++ b/sinner/processors/frame/BaseFrameProcessor.py @@ -8,10 +8,11 @@ from argparse import Namespace from sinner.handlers.frame.BaseFrameHandler import BaseFrameHandler +from sinner.handlers.frame.CV2VideoHandler import CV2VideoHandler from sinner.validators.AttributeLoader import AttributeLoader, Rules from sinner.State import State from sinner.typing import Frame, FramesDataType, FrameDataType, NumeratedFrame -from sinner.utilities import load_class, get_mem_usage, read_image, suggest_execution_threads, suggest_execution_providers, decode_execution_providers, suggest_max_memory +from sinner.utilities import load_class, get_mem_usage, suggest_execution_threads, suggest_execution_providers, decode_execution_providers, suggest_max_memory class BaseFrameProcessor(ABC, AttributeLoader): @@ -96,7 +97,7 @@ def process_frames(self, frame_data: FrameDataType, state: State) -> None: # ty if isinstance(frame_data, int): frame_num, frame = self.extract_frame_method(frame_data) else: - frame = read_image(frame_data[1]) + frame = CV2VideoHandler.read_image(frame_data[1]) frame_num = frame_data[0] state.save_temp_frame(self.process_frame(frame), frame_num) except Exception as exception: diff --git a/sinner/processors/frame/FaceSwapper.py b/sinner/processors/frame/FaceSwapper.py index 453d4440..f2c10650 100644 --- a/sinner/processors/frame/FaceSwapper.py +++ b/sinner/processors/frame/FaceSwapper.py @@ -4,10 +4,11 @@ import insightface from sinner.FaceAnalyser import FaceAnalyser +from sinner.handlers.frame.CV2VideoHandler import CV2VideoHandler from sinner.validators.AttributeLoader import Rules from sinner.processors.frame.BaseFrameProcessor import BaseFrameProcessor from sinner.typing import Face, Frame, FaceSwapperType -from sinner.utilities import conditional_download, read_image, get_app_dir, is_image, is_video, get_file_name +from sinner.utilities import conditional_download, get_app_dir, is_image, is_video, get_file_name class FaceSwapper(BaseFrameProcessor): @@ -65,7 +66,7 @@ def suggest_output_path(self) -> str: @property def source_face(self) -> Face | None: if self._source_face is None: - self._source_face = self.face_analyser.get_one_face(read_image(self.source_path)) + self._source_face = self.face_analyser.get_one_face(CV2VideoHandler.read_image(self.source_path)) return self._source_face @property diff --git a/sinner/utilities.py b/sinner/utilities.py index d882114b..28ed2690 100644 --- a/sinner/utilities.py +++ b/sinner/utilities.py @@ -9,11 +9,9 @@ import urllib from typing import List, Literal, Any, get_type_hints -import cv2 import onnxruntime import psutil import tensorflow -from numpy import uint8, fromfile from tqdm import tqdm from sinner.typing import Frame @@ -77,25 +75,6 @@ def resolve_relative_path(path: str, from_file: str | None = None) -> str: return os.path.abspath(os.path.join(os.path.dirname(from_file), path)) # type: ignore[arg-type] -def read_image(path: str) -> Frame: - if platform.system().lower() == 'windows': # issue #511 - image = cv2.imdecode(fromfile(path, dtype=uint8), cv2.IMREAD_UNCHANGED) - if image.shape[2] == 4: # fixes the alpha-channel issue - image = cv2.cvtColor(image, cv2.COLOR_BGRA2BGR) - return image - else: - return cv2.imread(path) - - -def write_image(image: Frame, 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, image) - - def get_mem_usage(param: Literal['rss', 'vms', 'shared', 'text', 'lib', 'data', 'dirty'] = 'rss', size: Literal['b', 'k', 'm', 'g'] = 'm') -> int: """ The `memory_info()` method of the `psutil.Process` class provides information about the memory usage of a process. It returns a named tuple containing the following attributes: From e77f038f5da0b318dc9a50ad88edf396f0c869e8 Mon Sep 17 00:00:00 2001 From: Pozitronik Date: Wed, 26 Jul 2023 01:33:01 +0400 Subject: [PATCH 4/6] type fixed --- sinner/Core.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/sinner/Core.py b/sinner/Core.py index deeb0ad8..135ba784 100644 --- a/sinner/Core.py +++ b/sinner/Core.py @@ -33,7 +33,7 @@ class Core(AttributeLoader, Status): target_path: str - output_path: str | None + output_path: str frame_processor: List[str] frame_handler: str temp_dir: str @@ -58,7 +58,6 @@ def rules(self) -> Rules: 'parameter': {'output', 'output-path'}, 'attribute': 'output_path', 'default:': lambda: self.suggest_output_path(), - 'default': None, }, { 'parameter': 'frame-processor', From 0b963b03e3c55ccfb17cdaf893f79de60da6a698 Mon Sep 17 00:00:00 2001 From: Pozitronik Date: Wed, 26 Jul 2023 01:33:22 +0400 Subject: [PATCH 5/6] unused import --- sinner/utilities.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/sinner/utilities.py b/sinner/utilities.py index 28ed2690..ca607f98 100644 --- a/sinner/utilities.py +++ b/sinner/utilities.py @@ -14,8 +14,6 @@ import tensorflow from tqdm import tqdm -from sinner.typing import Frame - TEMP_DIRECTORY = 'temp' From 88fdabc09860c86d13416474d6a77f7e8b764e03 Mon Sep 17 00:00:00 2001 From: Pozitronik Date: Wed, 26 Jul 2023 01:43:01 +0400 Subject: [PATCH 6/6] tests fixed --- tests/processors/test_base_processor.py | 4 ++-- tests/processors/test_face_enhancer.py | 4 ++-- tests/processors/test_face_swapper.py | 4 ++-- tests/test_face_analyser.py | 8 ++++---- 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/processors/test_base_processor.py b/tests/processors/test_base_processor.py index 3f408685..3c1e37a1 100644 --- a/tests/processors/test_base_processor.py +++ b/tests/processors/test_base_processor.py @@ -8,12 +8,12 @@ from sinner.Parameters import Parameters from sinner.handlers.frame.BaseFrameHandler import BaseFrameHandler +from sinner.handlers.frame.CV2VideoHandler import CV2VideoHandler from sinner.handlers.frame.VideoHandler import VideoHandler from sinner.processors.frame.BaseFrameProcessor import BaseFrameProcessor from sinner.processors.frame.DummyProcessor import DummyProcessor from sinner.State import State from sinner.typing import Frame -from sinner.utilities import read_image from tests.constants import source_jpg, target_png, IMAGE_SHAPE, target_mp4, tmp_dir, TARGET_FC parameters: Namespace = Parameters(f'--frame-processor=DummyProcessor --execution-provider=cpu --execution-threads={multiprocessing.cpu_count()} --source-path="{source_jpg}" --target-path="{target_mp4}" --output-path="{tmp_dir}"').parameters @@ -61,7 +61,7 @@ def test_init(): def test_process_frame(): - processed_frame = get_test_object().process_frame(read_image(target_png)) + processed_frame = get_test_object().process_frame(CV2VideoHandler.read_image(target_png)) assert (processed_frame, Frame) assert processed_frame.shape == IMAGE_SHAPE diff --git a/tests/processors/test_face_enhancer.py b/tests/processors/test_face_enhancer.py index 52a00628..1d14f6c5 100644 --- a/tests/processors/test_face_enhancer.py +++ b/tests/processors/test_face_enhancer.py @@ -5,10 +5,10 @@ from sinner.Parameters import Parameters from sinner.FaceAnalyser import FaceAnalyser +from sinner.handlers.frame.CV2VideoHandler import CV2VideoHandler from sinner.processors.frame.FaceEnhancer import FaceEnhancer from sinner.State import State from sinner.typing import Frame -from sinner.utilities import read_image from tests.constants import target_png, IMAGE_SHAPE, tmp_dir parameters: Namespace = Parameters(f'--execution-provider=cpu --execution-threads={multiprocessing.cpu_count()} --max-memory=12 --target-path="{target_png}" --output-path="{tmp_dir}"').parameters @@ -36,6 +36,6 @@ def test_init(): def test_process_frame(): - processed_frame = get_test_object().process_frame(read_image(target_png)) + processed_frame = get_test_object().process_frame(CV2VideoHandler.read_image(target_png)) assert (processed_frame, Frame) assert processed_frame.shape == IMAGE_SHAPE diff --git a/tests/processors/test_face_swapper.py b/tests/processors/test_face_swapper.py index a1ef40ce..794e1c29 100644 --- a/tests/processors/test_face_swapper.py +++ b/tests/processors/test_face_swapper.py @@ -3,9 +3,9 @@ from sinner.Parameters import Parameters from sinner.FaceAnalyser import FaceAnalyser +from sinner.handlers.frame.CV2VideoHandler import CV2VideoHandler from sinner.processors.frame.FaceSwapper import FaceSwapper from sinner.typing import Frame, FaceSwapperType, Face -from sinner.utilities import read_image from tests.constants import source_jpg, target_png, IMAGE_SHAPE, tmp_dir parameters: Namespace = Parameters(f'--execution-provider=cpu --execution-threads={multiprocessing.cpu_count()} --max-memory=12 --source-path="{source_jpg}" --target-path="{target_png}" --output-path="{tmp_dir}"').parameters @@ -30,6 +30,6 @@ def test_face_analysis(): def test_process_frame(): - processed_frame = get_test_object().process_frame(read_image(target_png)) + processed_frame = get_test_object().process_frame(CV2VideoHandler.read_image(target_png)) assert (processed_frame, Frame) assert processed_frame.shape == IMAGE_SHAPE diff --git a/tests/test_face_analyser.py b/tests/test_face_analyser.py index bdb20464..84b06081 100644 --- a/tests/test_face_analyser.py +++ b/tests/test_face_analyser.py @@ -1,8 +1,8 @@ from typing import List from sinner.FaceAnalyser import FaceAnalyser +from sinner.handlers.frame.CV2VideoHandler import CV2VideoHandler from sinner.typing import Face -from sinner.utilities import read_image from tests.constants import source_jpg, target_faces @@ -12,7 +12,7 @@ def get_test_object() -> FaceAnalyser: def test_one_face(): analyser = get_test_object() - face = analyser.get_one_face(read_image(source_jpg)) + face = analyser.get_one_face(CV2VideoHandler.read_image(source_jpg)) assert (face, Face) assert face.age == 31 assert face.sex == 'F' @@ -20,7 +20,7 @@ def test_one_face(): def test_one_face_from_many(): analyser = get_test_object() - face = analyser.get_one_face(read_image(target_faces)) + face = analyser.get_one_face(CV2VideoHandler.read_image(target_faces)) assert (face, Face) assert face.age == 47 assert face.sex == 'M' @@ -28,7 +28,7 @@ def test_one_face_from_many(): def test_many_faces(): analyser = get_test_object() - faces = analyser.get_many_faces(read_image(target_faces)) + faces = analyser.get_many_faces(CV2VideoHandler.read_image(target_faces)) assert (faces, List) assert len(faces) == 2 assert faces[0].age == 28