From de19212a719b0bc6ea5455b3e32072a76d021df6 Mon Sep 17 00:00:00 2001 From: Mia Herkt Date: Fri, 27 Sep 2024 17:39:18 +0200 Subject: PEP8 compliance --- nsfw_detect.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) (limited to 'nsfw_detect.py') 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() -- cgit v1.2.3