1
0
Fork 0
forked from mia/0x0

Fix some flake8 errors in cleanup and nsfw_detect

Just some minor code cleanup
This commit is contained in:
Nikolas Nyby 2021-08-06 15:18:02 -04:00
parent 47ff3a1152
commit b8def71a94
2 changed files with 30 additions and 14 deletions

View file

@ -18,7 +18,10 @@
and limitations under the License. and limitations under the License.
""" """
import os, sys, time, datetime import os
import sys
import time
import datetime
from fhost import app from fhost import app
os.chdir(os.path.dirname(sys.argv[0])) os.chdir(os.path.dirname(sys.argv[0]))
@ -33,7 +36,7 @@ maxd = 365
for f in files: for f in files:
stat = os.stat(f) stat = os.stat(f)
systime = time.time() systime = time.time()
age = datetime.timedelta(seconds = systime - stat.st_mtime).days age = datetime.timedelta(seconds=(systime - stat.st_mtime)).days
maxage = mind + (-maxd + mind) * (stat.st_size / maxs - 1) ** 3 maxage = mind + (-maxd + mind) * (stat.st_size / maxs - 1) ** 3

View file

@ -24,21 +24,30 @@ import sys
from io import BytesIO from io import BytesIO
from subprocess import run, PIPE, DEVNULL from subprocess import run, PIPE, DEVNULL
os.environ["GLOG_minloglevel"] = "2" # seriously :|
import caffe import caffe
os.environ["GLOG_minloglevel"] = "2" # seriously :|
class NSFWDetector: class NSFWDetector:
def __init__(self): def __init__(self):
npath = os.path.join(os.path.dirname(__file__), "nsfw_model") npath = os.path.join(os.path.dirname(__file__), "nsfw_model")
self.nsfw_net = caffe.Net(os.path.join(npath, "deploy.prototxt"), self.nsfw_net = caffe.Net(
os.path.join(npath, "resnet_50_1by2_nsfw.caffemodel"), os.path.join(npath, "deploy.prototxt"),
caffe.TEST) os.path.join(npath, "resnet_50_1by2_nsfw.caffemodel"),
self.caffe_transformer = caffe.io.Transformer({'data': self.nsfw_net.blobs['data'].data.shape}) caffe.TEST)
self.caffe_transformer.set_transpose('data', (2, 0, 1)) # move image channels to outermost self.caffe_transformer = caffe.io.Transformer({
self.caffe_transformer.set_mean('data', np.array([104, 117, 123])) # subtract the dataset-mean value in each channel 'data': self.nsfw_net.blobs['data'].data.shape
self.caffe_transformer.set_raw_scale('data', 255) # rescale from [0, 1] to [0, 255] })
self.caffe_transformer.set_channel_swap('data', (2, 1, 0)) # swap channels from RGB to BGR # move image channels to outermost
self.caffe_transformer.set_transpose('data', (2, 0, 1))
# subtract the dataset-mean value in each channel
self.caffe_transformer.set_mean('data', np.array([104, 117, 123]))
# rescale from [0, 1] to [0, 255]
self.caffe_transformer.set_raw_scale('data', 255)
# swap channels from RGB to BGR
self.caffe_transformer.set_channel_swap('data', (2, 1, 0))
def _compute(self, img): def _compute(self, img):
image = caffe.io.load_image(BytesIO(img)) image = caffe.io.load_image(BytesIO(img))
@ -54,8 +63,8 @@ class NSFWDetector:
input_name = self.nsfw_net.inputs[0] input_name = self.nsfw_net.inputs[0]
output_layers = ["prob"] output_layers = ["prob"]
all_outputs = self.nsfw_net.forward_all(blobs=output_layers, all_outputs = self.nsfw_net.forward_all(
**{input_name: transformed_image}) blobs=output_layers, **{input_name: transformed_image})
outputs = all_outputs[output_layers[0]][0].astype(float) outputs = all_outputs[output_layers[0]][0].astype(float)
@ -63,7 +72,10 @@ class NSFWDetector:
def detect(self, fpath): def detect(self, fpath):
try: try:
ff = run(["ffmpegthumbnailer", "-m", "-o-", "-s256", "-t50%", "-a", "-cpng", "-i", fpath], stdout=PIPE, stderr=DEVNULL, check=True) ff = run([
"ffmpegthumbnailer", "-m", "-o-", "-s256", "-t50%", "-a",
"-cpng", "-i", fpath
], stdout=PIPE, stderr=DEVNULL, check=True)
image_data = ff.stdout image_data = ff.stdout
except: except:
return -1.0 return -1.0
@ -72,6 +84,7 @@ class NSFWDetector:
return scores[1] return scores[1]
if __name__ == "__main__": if __name__ == "__main__":
n = NSFWDetector() n = NSFWDetector()