summaryrefslogtreecommitdiff
path: root/nsfw_detect.py
diff options
context:
space:
mode:
authorMia Herkt <mia@0x0.st>2022-11-29 21:42:46 +0100
committerMia Herkt <mia@0x0.st>2022-11-29 21:44:07 +0100
commit14cfe3da58a8194eb850a28e4170e684d1444831 (patch)
tree6b0d5a1ea1bd6808cd6fcf00dbf56fcd57251b36 /nsfw_detect.py
parentaa443178e1fdd30504850fb1ebde138491cd3a48 (diff)
nsfw_detect: Use pathlib, fix deprecation warning
Also fix glog suppression
Diffstat (limited to 'nsfw_detect.py')
-rwxr-xr-xnsfw_detect.py15
1 files changed, 7 insertions, 8 deletions
diff --git a/nsfw_detect.py b/nsfw_detect.py
index 225d8a7..eddf9cb 100755
--- a/nsfw_detect.py
+++ b/nsfw_detect.py
@@ -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
})