mia/0x0
1
0
Fork 1
mirror of https://git.0x0.st/mia/0x0.git synced 2024-11-22 08:57:11 +01:00
Code Issues Projects Releases Packages Wiki Activity

nsfw_detect: Use pathlib, fix deprecation warning

Browse source 
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
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
})