forked from mia/0x0
nsfw_detect: Use pathlib, fix deprecation warning
Also fix glog suppression
This commit is contained in:
parent
aa443178e1
commit
14cfe3da58
1 changed files with 7 additions and 8 deletions
|
@ -23,20 +23,19 @@ import os
|
|||
import sys
|
||||
from io import BytesIO
|
||||
from subprocess import run, PIPE, DEVNULL
|
||||
|
||||
import caffe
|
||||
from pathlib import Path
|
||||
|
||||
os.environ["GLOG_minloglevel"] = "2" # seriously :|
|
||||
|
||||
import caffe
|
||||
|
||||
class NSFWDetector:
|
||||
def __init__(self):
|
||||
|
||||
npath = os.path.join(os.path.dirname(__file__), "nsfw_model")
|
||||
npath = Path(__file__).parent / "nsfw_model"
|
||||
self.nsfw_net = caffe.Net(
|
||||
os.path.join(npath, "deploy.prototxt"),
|
||||
os.path.join(npath, "resnet_50_1by2_nsfw.caffemodel"),
|
||||
caffe.TEST)
|
||||
str(npath / "deploy.prototxt"),
|
||||
caffe.TEST,
|
||||
weights = str(npath / "resnet_50_1by2_nsfw.caffemodel")
|
||||
)
|
||||
self.caffe_transformer = caffe.io.Transformer({
|
||||
'data': self.nsfw_net.blobs['data'].data.shape
|
||||
})
|
||||
|
|
Loading…
Reference in a new issue