Aadhithya commited on
Commit
3f47bc4
·
1 Parent(s): e0cff75

Delete roop/predictor.py

Browse files
Files changed (1) hide show
  1. 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)