Re-enabled the NSFW function (turn-off by default).

@refer to the PR #237 in the `experimental` branch.

Signed-off-by: Vic P <vic4key@gmail.com>
This commit is contained in:
Vic P 2024-08-21 02:02:00 +07:00
parent e4b494174d
commit 7313a332c8
5 changed files with 46 additions and 26 deletions

View File

@ -152,14 +152,15 @@ Additional command line arguments are given below. To learn out what they do, ch
```
options:
-h, --help show this help message and exit
-s SOURCE_PATH, --source SOURCE_PATH select an source image
-t TARGET_PATH, --target TARGET_PATH select an target image or video
-s SOURCE_PATH, --source SOURCE_PATH select a source image
-t TARGET_PATH, --target TARGET_PATH select a target image or video
-o OUTPUT_PATH, --output OUTPUT_PATH select output file or directory
--frame-processor FRAME_PROCESSOR [FRAME_PROCESSOR ...] frame processors (choices: face_swapper, face_enhancer, ...)
--keep-fps keep original fps
--keep-audio keep original audio
--keep-frames keep temporary frames
--many-faces process every face
--nsfw-filter filter the NSFW image or video.
--video-encoder {libx264,libx265,libvpx-vp9} adjust output video encoder
--video-quality [0-51] adjust output video quality
--max-memory MAX_MEMORY maximum amount of RAM in GB

View File

@ -39,6 +39,7 @@ def parse_args() -> None:
program.add_argument('--keep-audio', help='keep original audio', dest='keep_audio', action='store_true', default=True)
program.add_argument('--keep-frames', help='keep temporary frames', dest='keep_frames', action='store_true', default=False)
program.add_argument('--many-faces', help='process every face', dest='many_faces', action='store_true', default=False)
program.add_argument('--nsfw-filter', help='filter the NSFW image or video', dest='nsfw_filter', action='store_true', default=False)
program.add_argument('--video-encoder', help='adjust output video encoder', dest='video_encoder', default='libx264', choices=['libx264', 'libx265', 'libvpx-vp9'])
program.add_argument('--video-quality', help='adjust output video quality', dest='video_quality', type=int, default=18, choices=range(52), metavar='[0-51]')
program.add_argument('--max-memory', help='maximum amount of RAM in GB', dest='max_memory', type=int, default=suggest_max_memory())
@ -63,6 +64,7 @@ def parse_args() -> None:
modules.globals.keep_audio = args.keep_audio
modules.globals.keep_frames = args.keep_frames
modules.globals.many_faces = args.many_faces
modules.globals.nsfw_filter = args.nsfw_filter
modules.globals.video_encoder = args.video_encoder
modules.globals.video_quality = args.video_quality
modules.globals.max_memory = args.max_memory
@ -75,8 +77,6 @@ def parse_args() -> None:
else:
modules.globals.fp_ui['face_enhancer'] = False
modules.globals.nsfw = False
# translate deprecated args
if args.source_path_deprecated:
print('\033[33mArgument -f and --face are deprecated. Use -s and --source instead.\033[0m')
@ -169,13 +169,15 @@ def start() -> None:
for frame_processor in get_frame_processors_modules(modules.globals.frame_processors):
if not frame_processor.pre_start():
return
update_status('Processing...')
# process image to image
if has_image_extension(modules.globals.target_path):
if modules.globals.nsfw == False:
from modules.predicter import predict_image
if predict_image(modules.globals.target_path):
destroy()
shutil.copy2(modules.globals.target_path, modules.globals.output_path)
if modules.globals.nsfw_filter and ui.check_and_ignore_nsfw(modules.globals.target_path, destroy):
return
try:
shutil.copy2(modules.globals.target_path, modules.globals.output_path)
except Exception as e:
print("Error copying file:", str(e))
for frame_processor in get_frame_processors_modules(modules.globals.frame_processors):
update_status('Progressing...', frame_processor.NAME)
frame_processor.process_image(modules.globals.source_path, modules.globals.output_path, modules.globals.output_path)
@ -186,10 +188,8 @@ def start() -> None:
update_status('Processing to image failed!')
return
# process image to videos
if modules.globals.nsfw == False:
from modules.predicter import predict_video
if predict_video(modules.globals.target_path):
destroy()
if modules.globals.nsfw_filter and ui.check_and_ignore_nsfw(modules.globals.target_path, destroy):
return
update_status('Creating temp resources...')
create_temp(modules.globals.target_path)
update_status('Extracting frames...')
@ -225,10 +225,10 @@ def start() -> None:
update_status('Processing to video failed!')
def destroy() -> None:
def destroy(to_quit=True) -> None:
if modules.globals.target_path:
clean_temp(modules.globals.target_path)
quit()
if to_quit: quit()
def run() -> None:

View File

@ -17,6 +17,7 @@ keep_fps = None
keep_audio = None
keep_frames = None
many_faces = None
nsfw_filter = None
video_encoder = None
video_quality = None
max_memory = None
@ -25,6 +26,5 @@ execution_threads = None
headless = None
log_level = 'error'
fp_ui: Dict[str, bool] = {}
nsfw = None
camera_input_combobox = None
webcam_preview_running = False

View File

@ -6,11 +6,13 @@ from modules.typing import Frame
MAX_PROBABILITY = 0.85
model = None
def predict_frame(target_frame: Frame) -> bool:
image = Image.fromarray(target_frame)
image = opennsfw2.preprocess_image(image, opennsfw2.Preprocessing.YAHOO)
model = opennsfw2.make_open_nsfw_model()
global 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

View File

@ -10,7 +10,7 @@ import modules.metadata
from modules.face_analyser import get_one_face
from modules.capturer import get_video_frame, get_video_frame_total
from modules.processors.frame.core import get_frame_processors_modules
from modules.utilities import is_image, is_video, resolve_relative_path
from modules.utilities import is_image, is_video, resolve_relative_path, has_image_extension
ROOT = None
ROOT_HEIGHT = 700
@ -88,9 +88,9 @@ 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)
# nsfw_value = ctk.BooleanVar(value=modules.globals.nsfw)
# nsfw_switch = ctk.CTkSwitch(root, text='NSFW', variable=nsfw_value, cursor='hand2', command=lambda: setattr(modules.globals, 'nsfw', nsfw_value.get()))
# nsfw_switch.place(relx=0.6, rely=0.7)
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)
start_button = ctk.CTkButton(root, text='Start', cursor='hand2', command=lambda: select_output_path(start))
start_button.place(relx=0.15, rely=0.80, relwidth=0.2, relheight=0.05)
@ -192,6 +192,23 @@ def select_output_path(start: Callable[[], None]) -> None:
start()
def check_and_ignore_nsfw(target, destroy: Callable = None) -> bool:
''' Check the target is NSFW or not.
TODO: Consider to make blur the target.
'''
from numpy import ndarray
from modules.predicter import predict_image, predict_video, predict_frame
if type(target) is str: # image/video file path
check_nsfw = predict_image if has_image_extension(target) else predict_video
elif type(target) is ndarray: # frame object
check_nsfw = predict_frame
if check_nsfw and check_nsfw(target):
if destroy: destroy(to_quit=False) # Do not need to destroy the window frame if the target is NSFW
update_status('Processing ignored!')
return True
else: return False
def render_image_preview(image_path: str, size: Tuple[int, int]) -> ctk.CTkImage:
image = Image.open(image_path)
if size:
@ -219,7 +236,6 @@ def toggle_preview() -> None:
elif modules.globals.source_path and modules.globals.target_path:
init_preview()
update_preview()
PREVIEW.deiconify()
def init_preview() -> None:
@ -234,11 +250,10 @@ def init_preview() -> None:
def update_preview(frame_number: int = 0) -> None:
if modules.globals.source_path and modules.globals.target_path:
update_status('Processing...')
temp_frame = get_video_frame(modules.globals.target_path, frame_number)
if modules.globals.nsfw == False:
from modules.predicter import predict_frame
if predict_frame(temp_frame):
quit()
if modules.globals.nsfw_filter and check_and_ignore_nsfw(temp_frame):
return
for frame_processor in get_frame_processors_modules(modules.globals.frame_processors):
temp_frame = frame_processor.process_frame(
get_one_face(cv2.imread(modules.globals.source_path)),
@ -248,6 +263,8 @@ def update_preview(frame_number: int = 0) -> None:
image = ImageOps.contain(image, (PREVIEW_MAX_WIDTH, PREVIEW_MAX_HEIGHT), Image.LANCZOS)
image = ctk.CTkImage(image, size=image.size)
preview_label.configure(image=image)
update_status('Processing succeed!')
PREVIEW.deiconify()
def webcam_preview():
if modules.globals.source_path is None: