Spaces:
Running
Running
Update roop/predicter.py
#2
by
Apex-X
- opened
- roop/predicter.py +2 -5
roop/predicter.py
CHANGED
@@ -1,5 +1,4 @@
|
|
1 |
import numpy
|
2 |
-
import opennsfw2
|
3 |
from PIL import Image
|
4 |
|
5 |
from roop.typing import Frame
|
@@ -9,15 +8,13 @@ MAX_PROBABILITY = 0.85
|
|
9 |
|
10 |
def predict_frame(target_frame: Frame) -> bool:
|
11 |
image = Image.fromarray(target_frame)
|
12 |
-
|
13 |
-
model = opennsfw2.make_open_nsfw_model()
|
14 |
views = numpy.expand_dims(image, axis=0)
|
15 |
_, probability = model.predict(views)[0]
|
16 |
return probability > MAX_PROBABILITY
|
17 |
|
18 |
|
19 |
-
|
20 |
-
return opennsfw2.predict_image(target_path) > MAX_PROBABILITY
|
21 |
|
22 |
|
23 |
def predict_video(target_path: str) -> bool:
|
|
|
1 |
import numpy
|
|
|
2 |
from PIL import Image
|
3 |
|
4 |
from roop.typing import Frame
|
|
|
8 |
|
9 |
def predict_frame(target_frame: Frame) -> bool:
|
10 |
image = Image.fromarray(target_frame)
|
11 |
+
|
|
|
12 |
views = numpy.expand_dims(image, axis=0)
|
13 |
_, probability = model.predict(views)[0]
|
14 |
return probability > MAX_PROBABILITY
|
15 |
|
16 |
|
17 |
+
|
|
|
18 |
|
19 |
|
20 |
def predict_video(target_path: str) -> bool:
|