summaryrefslogtreecommitdiff
path: root/nsfw_detect.py
diff options
context:
space:
mode:
authorMia Herkt <mia@0x0.st>2024-09-27 17:39:18 +0200
committerMia Herkt <mia@0x0.st>2024-09-27 18:30:33 +0200
commitde19212a719b0bc6ea5455b3e32072a76d021df6 (patch)
tree6f95f343bbb3587f50739c97019f3159b4b8a095 /nsfw_detect.py
parenta2147cc964d385a59f077add9d7e5ba009bf9b2e (diff)
PEP8 compliance
Diffstat (limited to 'nsfw_detect.py')
-rwxr-xr-xnsfw_detect.py14
1 files changed, 8 insertions, 6 deletions
diff --git a/nsfw_detect.py b/nsfw_detect.py
index 9c8b117..96fa9a2 100755
--- a/nsfw_detect.py
+++ b/nsfw_detect.py
@@ -18,32 +18,34 @@
and limitations under the License.
"""
-import os
import sys
-from pathlib import Path
-
import av
from transformers import pipeline
+
class NSFWDetector:
def __init__(self):
- self.classifier = pipeline("image-classification", model="giacomoarienti/nsfw-classifier")
+ self.classifier = pipeline("image-classification",
+ model="giacomoarienti/nsfw-classifier")
def detect(self, fpath):
try:
with av.open(fpath) as container:
- try: container.seek(int(container.duration / 2))
+ try:
+ container.seek(int(container.duration / 2))
except: container.seek(0)
frame = next(container.decode(video=0))
img = frame.to_image()
res = self.classifier(img)
- return max([x["score"] for x in res if x["label"] not in ["neutral", "drawings"]])
+ return max([x["score"] for x in res
+ if x["label"] not in ["neutral", "drawings"]])
except: pass
return -1.0
+
if __name__ == "__main__":
n = NSFWDetector()