diff options
| author | Mia Herkt <mia@0x0.st> | 2024-09-27 17:39:18 +0200 |
|---|---|---|
| committer | Mia Herkt <mia@0x0.st> | 2024-09-27 18:30:33 +0200 |
| commit | de19212a719b0bc6ea5455b3e32072a76d021df6 (patch) | |
| tree | 6f95f343bbb3587f50739c97019f3159b4b8a095 /nsfw_detect.py | |
| parent | a2147cc964d385a59f077add9d7e5ba009bf9b2e (diff) | |
PEP8 compliance
Diffstat (limited to 'nsfw_detect.py')
| -rwxr-xr-x | nsfw_detect.py | 14 |
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() |
