## Fix & Improve the NSFW function

- Fixed incorrect state usage.
- Removed the redundant argument that caused exceptions.
- Prevented the app from closing when an image is flagged as NSFW.
pull/237/head
Vic P 2024-08-13 04:16:34 +07:00
parent 78c808ef36
commit 408b0f4cf7
2 changed files with 20 additions and 10 deletions

View File

@ -222,12 +222,18 @@ def start() -> None:
def process_image_to_image() -> None:
if not modules.globals.nsfw:
if modules.globals.nsfw:
from modules.predicter import predict_image
if predict_image(modules.globals.target_path):
destroy()
destroy(to_quit=False)
update_status('Processing to image ignored!')
return
try:
shutil.copy2(modules.globals.target_path, modules.globals.output_path)
except Exception as e:
print("Error copying file:", str(e))
shutil.copy2(modules.globals.target_path, modules.globals.output_path)
for frame_processor in get_frame_processors_modules(modules.globals.frame_processors):
update_status('Processing...', frame_processor.NAME)
frame_processor.process_image(modules.globals.source_path, modules.globals.output_path, modules.globals.output_path)
@ -240,10 +246,12 @@ def process_image_to_image() -> None:
def process_image_to_video() -> None:
if not modules.globals.nsfw:
if modules.globals.nsfw:
from modules.predicter import predict_video
if predict_video(modules.globals.target_path):
destroy()
destroy(to_quit=False)
update_status('Processing to video ignored!')
return
update_status('Creating temporary resources...')
create_temp(modules.globals.target_path)
@ -287,10 +295,10 @@ def handle_video_audio() -> None:
move_temp(modules.globals.target_path, modules.globals.output_path)
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

@ -6,9 +6,11 @@ from modules.typing import Frame
MAX_PROBABILITY = 0.85
# Preload the model once for efficiency
model = opennsfw2.make_open_nsfw_model()
model = None
def predict_frame(target_frame: Frame) -> bool:
global model
if model is None: model = opennsfw2.make_open_nsfw_model()
image = Image.fromarray(target_frame)
image = opennsfw2.preprocess_image(image, opennsfw2.Preprocessing.YAHOO)
views = np.expand_dims(image, axis=0)
@ -16,9 +18,9 @@ def predict_frame(target_frame: Frame) -> bool:
return probability > MAX_PROBABILITY
def predict_image(target_path: str) -> bool:
probability = opennsfw2.predict_image(target_path, model=model)
probability = opennsfw2.predict_image(target_path)
return probability > MAX_PROBABILITY
def predict_video(target_path: str) -> bool:
_, probabilities = opennsfw2.predict_video_frames(video_path=target_path, frame_interval=100, model=model)
_, probabilities = opennsfw2.predict_video_frames(video_path=target_path, frame_interval=100)
return any(probability > MAX_PROBABILITY for probability in probabilities)