diff options
| author | Mia Herkt <mia@0x0.st> | 2022-11-29 21:42:46 +0100 |
|---|---|---|
| committer | Mia Herkt <mia@0x0.st> | 2022-11-29 21:44:07 +0100 |
| commit | 14cfe3da58a8194eb850a28e4170e684d1444831 (patch) | |
| tree | 6b0d5a1ea1bd6808cd6fcf00dbf56fcd57251b36 | |
| parent | aa443178e1fdd30504850fb1ebde138491cd3a48 (diff) | |
nsfw_detect: Use pathlib, fix deprecation warning
Also fix glog suppression
| -rwxr-xr-x | nsfw_detect.py | 15 |
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 }) |
