christian.schwarz/Deep-Live-Cam
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

## Fix & Improve the NSFW function

Browse Source 
- Fixed incorrect state usage.
- Removed the redundant argument that caused exceptions.
- Prevented the app from closing when an image is flagged as NSFW.
This commit is contained in:
Vic P
2024-08-13 04:16:34 +07:00
parent 78c808ef36
commit 408b0f4cf7
2 changed files with 20 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/predicter.py
+5 -3
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)