cookie.ebil.club
+7
-8
nsfw_detect.py
+7
-8
nsfw_detect.py
···
23
23
import sys
24
24
from io import BytesIO
25
25
from subprocess import run, PIPE, DEVNULL
26
-
27
-
import caffe
26
+
from pathlib import Path
28
27
29
28
os.environ["GLOG_minloglevel"] = "2" # seriously :|
30
-
29
+
import caffe
31
30
32
31
class NSFWDetector:
33
32
def __init__(self):
34
-
35
-
npath = os.path.join(os.path.dirname(__file__), "nsfw_model")
33
+
npath = Path(__file__).parent / "nsfw_model"
36
34
self.nsfw_net = caffe.Net(
37
-
os.path.join(npath, "deploy.prototxt"),
38
-
os.path.join(npath, "resnet_50_1by2_nsfw.caffemodel"),
39
-
caffe.TEST)
35
+
str(npath / "deploy.prototxt"),
36
+
caffe.TEST,
37
+
weights = str(npath / "resnet_50_1by2_nsfw.caffemodel")
38
+
)
40
39
self.caffe_transformer = caffe.io.Transformer({
41
40
'data': self.nsfw_net.blobs['data'].data.shape
42
41
})