Apex-X commited on
Commit
25b194f
1 Parent(s): 12e09f7

Update roop/predictor.py

Browse files
Files changed (1) hide show
  1. roop/predictor.py +10 -31
roop/predictor.py CHANGED
@@ -1,43 +1,22 @@
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)
 
 
1
  import threading
2
  import numpy
 
3
  from PIL import Image
 
4
 
5
  from roop.typing import Frame
6
 
7
+ # Define any other necessary variables or constants here
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8
 
9
  def predict_frame(target_frame: Frame) -> bool:
10
+ # Modify this function as needed for your specific use case, without NSFW prediction
11
+ # For example, you can implement custom image analysis or processing here
12
+ return False
 
 
 
13
 
14
  def predict_image(target_path: str) -> bool:
15
+ # Modify this function as needed for your specific use case, without NSFW prediction
16
+ # For example, you can check the image based on your application's requirements
17
+ return False
18
 
19
  def predict_video(target_path: str) -> bool:
20
+ # Modify this function as needed for your specific use case, without NSFW prediction
21
+ # For example, you can analyze video frames for other purposes
22
+ return False