Aadhithya
commited on
Commit
·
3f47bc4
1
Parent(s):
e0cff75
Delete roop/predictor.py
Browse files- roop/predictor.py +0 -43
roop/predictor.py
DELETED
@@ -1,43 +0,0 @@
|
|
1 |
-
import threading
|
2 |
-
import numpy
|
3 |
-
import opennsfw2
|
4 |
-
from PIL import Image
|
5 |
-
from keras import Model
|
6 |
-
|
7 |
-
from roop.typing import Frame
|
8 |
-
|
9 |
-
PREDICTOR = None
|
10 |
-
THREAD_LOCK = threading.Lock()
|
11 |
-
MAX_PROBABILITY = 0.85
|
12 |
-
|
13 |
-
|
14 |
-
def get_predictor() -> Model:
|
15 |
-
global PREDICTOR
|
16 |
-
|
17 |
-
with THREAD_LOCK:
|
18 |
-
if PREDICTOR is None:
|
19 |
-
PREDICTOR = opennsfw2.make_open_nsfw_model()
|
20 |
-
return PREDICTOR
|
21 |
-
|
22 |
-
|
23 |
-
def clear_predictor() -> None:
|
24 |
-
global PREDICTOR
|
25 |
-
|
26 |
-
PREDICTOR = None
|
27 |
-
|
28 |
-
|
29 |
-
def predict_frame(target_frame: Frame) -> bool:
|
30 |
-
image = Image.fromarray(target_frame)
|
31 |
-
image = opennsfw2.preprocess_image(image, opennsfw2.Preprocessing.YAHOO)
|
32 |
-
views = numpy.expand_dims(image, axis=0)
|
33 |
-
_, probability = get_predictor().predict(views)[0]
|
34 |
-
return probability > MAX_PROBABILITY
|
35 |
-
|
36 |
-
|
37 |
-
def predict_image(target_path: str) -> bool:
|
38 |
-
return opennsfw2.predict_image(target_path) > MAX_PROBABILITY
|
39 |
-
|
40 |
-
|
41 |
-
def predict_video(target_path: str) -> bool:
|
42 |
-
_, probabilities = opennsfw2.predict_video_frames(video_path=target_path, frame_interval=100)
|
43 |
-
return any(probability > MAX_PROBABILITY for probability in probabilities)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|