lod/0x0
1
0
Fork 0
forked from mia/0x0
Code Pull requests Activity

PEP8 compliance

Browse source
This commit is contained in:
Mia Herkt 2024-09-27 17:39:18 +02:00
parent a2147cc964
commit de19212a71
No known key found for this signature in database
18 changed files with 376 additions and 299 deletions
Download patch file Download diff file Expand all files Collapse all files

14
nsfw_detect.py
View file

@ -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()