[mirror of https://git.0x0.st/mia/0x0] No-bullshit file hosting and URL shortening service https://0x0.st
cookie.ebil.club
1#!/usr/bin/env python3
2
3"""
4 Copyright © 2024 Mia Herkt
5 Licensed under the EUPL, Version 1.2 or - as soon as approved
6 by the European Commission - subsequent versions of the EUPL
7 (the "License");
8 You may not use this work except in compliance with the License.
9 You may obtain a copy of the license at:
10
11 https://joinup.ec.europa.eu/software/page/eupl
12
13 Unless required by applicable law or agreed to in writing,
14 software distributed under the License is distributed on an
15 "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
16 either express or implied.
17 See the License for the specific language governing permissions
18 and limitations under the License.
19"""
20
21import sys
22import av
23from transformers import pipeline
24
25
26class NSFWDetector:
27 def __init__(self):
28 self.classifier = pipeline("image-classification",
29 model="giacomoarienti/nsfw-classifier")
30
31 def detect(self, fpath):
32 try:
33 with av.open(fpath) as container:
34 try:
35 container.seek(int(container.duration / 2))
36 except: container.seek(0)
37
38 frame = next(container.decode(video=0))
39 img = frame.to_image()
40 res = self.classifier(img)
41
42 return max([x["score"] for x in res
43 if x["label"] not in ["neutral", "drawings"]])
44 except: pass
45
46 return -1.0
47
48
49if __name__ == "__main__":
50 n = NSFWDetector()
51
52 for inf in sys.argv[1:]:
53 score = n.detect(inf)
54 print(inf, score)