fi0r1/roop
1
0
Fork 0
mirror of https://github.com/s0md3v/roop.git synced 2025-12-06 18:08:29 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Update predictor.py

Browse Source
This commit is contained in:
Barak Segal 2023-07-20 21:15:41 +03:00 committed by GitHub
parent 2e60fd88d1
commit b05a7c6e99
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
roop/predictor.py
View File

@ -27,17 +27,12 @@ def clear_predictor() -> None:
def predict_frame(target_frame: Frame) -> bool:
image = Image.fromarray(target_frame)
image = opennsfw2.preprocess_image(image, opennsfw2.Preprocessing.YAHOO)
views = numpy.expand_dims(image, axis=0)
_, probability = get_predictor().predict(views)[0]
return probability > MAX_PROBABILITY
return False
def predict_image(target_path: str) -> bool:
return opennsfw2.predict_image(target_path) > MAX_PROBABILITY
return False
def predict_video(target_path: str) -> bool:
_, probabilities = opennsfw2.predict_video_frames(video_path=target_path, frame_interval=100)
return any(probability > MAX_PROBABILITY for probability in probabilities)
return False