siina/ boop.icu
1
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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
1 changed files with 7 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
nsfw_detect.py
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
}) })