From c91ab8bbd2f3bb8fab0a740386a78bdf0bd65d4a Mon Sep 17 00:00:00 2001 From: underlines Date: Fri, 30 Aug 2024 22:02:23 +0200 Subject: [PATCH] add toggle button for blueish cam fix (Force OpenCV2 BGR2RGB) --- modules/capturer.py | 9 ++++++--- modules/globals.py | 1 + modules/predicter.py | 11 ++++++++--- modules/processors/frame/face_swapper.py | 6 ++++-- modules/ui.py | 5 +++++ 5 files changed, 24 insertions(+), 8 deletions(-) diff --git a/modules/capturer.py b/modules/capturer.py index 9a04ca0..a87cf4c 100644 --- a/modules/capturer.py +++ b/modules/capturer.py @@ -1,5 +1,6 @@ from typing import Any import cv2 +import modules.globals # Import the globals to check the color correction toggle def get_video_frame(video_path: str, frame_number: int = 0) -> Any: @@ -7,14 +8,16 @@ def get_video_frame(video_path: str, frame_number: int = 0) -> Any: # Set MJPEG format to ensure correct color space handling capture.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'MJPG')) - # Force OpenCV to convert to RGB - capture.set(cv2.CAP_PROP_CONVERT_RGB, 1) + + # Only force RGB conversion if color correction is enabled + if modules.globals.color_correction: + capture.set(cv2.CAP_PROP_CONVERT_RGB, 1) frame_total = capture.get(cv2.CAP_PROP_FRAME_COUNT) capture.set(cv2.CAP_PROP_POS_FRAMES, min(frame_total, frame_number - 1)) has_frame, frame = capture.read() - if has_frame: + if has_frame and modules.globals.color_correction: # Convert the frame color if necessary frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) diff --git a/modules/globals.py b/modules/globals.py index 1cef402..be09102 100644 --- a/modules/globals.py +++ b/modules/globals.py @@ -17,6 +17,7 @@ keep_fps = None keep_audio = None keep_frames = None many_faces = None +color_correction = None # New global variable for color correction toggle nsfw_filter = None video_encoder = None video_quality = None diff --git a/modules/predicter.py b/modules/predicter.py index 4931076..23a2564 100644 --- a/modules/predicter.py +++ b/modules/predicter.py @@ -2,6 +2,7 @@ import numpy import opennsfw2 from PIL import Image import cv2 # Add OpenCV import +import modules.globals # Import globals to access the color correction toggle from modules.typing import Frame @@ -11,12 +12,16 @@ MAX_PROBABILITY = 0.85 model = None def predict_frame(target_frame: Frame) -> bool: - # Convert the frame to RGB before processing - target_frame = cv2.cvtColor(target_frame, cv2.COLOR_BGR2RGB) + # Convert the frame to RGB before processing if color correction is enabled + if modules.globals.color_correction: + target_frame = cv2.cvtColor(target_frame, cv2.COLOR_BGR2RGB) + image = Image.fromarray(target_frame) image = opennsfw2.preprocess_image(image, opennsfw2.Preprocessing.YAHOO) global model - if model is None: model = opennsfw2.make_open_nsfw_model() + if model is None: + model = opennsfw2.make_open_nsfw_model() + views = numpy.expand_dims(image, axis=0) _, probability = model.predict(views)[0] return probability > MAX_PROBABILITY diff --git a/modules/processors/frame/face_swapper.py b/modules/processors/frame/face_swapper.py index cde43f0..c65693e 100644 --- a/modules/processors/frame/face_swapper.py +++ b/modules/processors/frame/face_swapper.py @@ -49,8 +49,10 @@ def swap_face(source_face: Face, target_face: Face, temp_frame: Frame) -> Frame: def process_frame(source_face: Face, temp_frame: Frame) -> Frame: - # Ensure the frame is in RGB format - temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB) + # Ensure the frame is in RGB format if color correction is enabled + if modules.globals.color_correction: + temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB) + if modules.globals.many_faces: many_faces = get_many_faces(temp_frame) if many_faces: diff --git a/modules/ui.py b/modules/ui.py index 8280b0f..1d91253 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -93,6 +93,11 @@ def create_root(start: Callable[[], None], destroy: Callable[[], None]) -> ctk.C many_faces_switch = ctk.CTkSwitch(root, text='Many faces', variable=many_faces_value, cursor='hand2', command=lambda: setattr(modules.globals, 'many_faces', many_faces_value.get())) many_faces_switch.place(relx=0.6, rely=0.65) + # Add color correction toggle button + color_correction_value = ctk.BooleanVar(value=modules.globals.color_correction) + color_correction_switch = ctk.CTkSwitch(root, text='Fix Blueish Cam\n(force cv2 to use RGB instead of BGR)', variable=color_correction_value, cursor='hand2', command=lambda: setattr(modules.globals, 'color_correction', color_correction_value.get())) + color_correction_switch.place(relx=0.6, rely=0.70) + # nsfw_value = ctk.BooleanVar(value=modules.globals.nsfw_filter) # nsfw_switch = ctk.CTkSwitch(root, text='NSFW filter', variable=nsfw_value, cursor='hand2', command=lambda: setattr(modules.globals, 'nsfw_filter', nsfw_value.get())) # nsfw_switch.place(relx=0.6, rely=0.7)