add toggle button for blueish cam fix (Force OpenCV2 BGR2RGB)

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)
 

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

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

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:

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)