1
0
Fork 0
forked from mia/0x0

nsfw_detect: Use pathlib, fix deprecation warning

Also fix glog suppression
This commit is contained in:
Mia Herkt 2022-11-29 21:42:46 +01:00
parent aa443178e1
commit 14cfe3da58
No known key found for this signature in database
GPG key ID: 72E154B8622EC191

View file

@ -23,20 +23,19 @@ import os
import sys import sys
from io import BytesIO from io import BytesIO
from subprocess import run, PIPE, DEVNULL from subprocess import run, PIPE, DEVNULL
from pathlib import Path
import caffe
os.environ["GLOG_minloglevel"] = "2" # seriously :| os.environ["GLOG_minloglevel"] = "2" # seriously :|
import caffe
class NSFWDetector: class NSFWDetector:
def __init__(self): def __init__(self):
npath = Path(__file__).parent / "nsfw_model"
npath = os.path.join(os.path.dirname(__file__), "nsfw_model")
self.nsfw_net = caffe.Net( self.nsfw_net = caffe.Net(
os.path.join(npath, "deploy.prototxt"), str(npath / "deploy.prototxt"),
os.path.join(npath, "resnet_50_1by2_nsfw.caffemodel"), caffe.TEST,
caffe.TEST) weights = str(npath / "resnet_50_1by2_nsfw.caffemodel")
)
self.caffe_transformer = caffe.io.Transformer({ self.caffe_transformer = caffe.io.Transformer({
'data': self.nsfw_net.blobs['data'].data.shape 'data': self.nsfw_net.blobs['data'].data.shape
}) })