Create utils.py

utils.py

@@ -0,0 +1,84 @@
+import cv2
+import tempfile
+import requests
+import os
+from PIL import Image
+from transformers import pipeline
+
+classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
+
+
+def download_file(url):
+    response = requests.get(url, stream=True)
+    tmp = tempfile.NamedTemporaryFile(delete=False)
+    for chunk in response.iter_content(1024):
+        tmp.write(chunk)
+    tmp.close()
+    return tmp.name
+
+
+def get_video_duration(video_path):
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    frame_count = cap.get(cv2.CAP_PROP_FRAME_COUNT)
+    cap.release()
+    return frame_count / fps if fps > 0 else 0
+
+
+def extract_frame(video_path, second):
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    frame_number = int(fps * second)
+
+    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
+    success, frame = cap.read()
+    cap.release()
+
+    if not success:
+        return None
+
+    tmp_file = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
+    cv2.imwrite(tmp_file.name, frame)
+    return tmp_file.name
+
+
+def get_frame_times(duration, file_size_mb):
+    if duration <= 2:
+        return [1, 2]
+
+    elif duration <= 10:
+        return [2]
+
+    elif duration <= 15:
+        return [4, 9, 13]
+
+    if file_size_mb > 14:
+        return [4, 9, 13]
+
+    return [2]
+
+
+def check_image_nsfw(image_path):
+    image = Image.open(image_path).convert("RGB")
+    result = classifier(image)
+
+    for r in result:
+        if r["label"] == "nsfw" and r["score"] > 0.5:
+            return True
+
+    return False
+
+
+def check_video_nsfw(video_path):
+    size_mb = os.path.getsize(video_path) / (1024 * 1024)
+    duration = get_video_duration(video_path)
+
+    times = get_frame_times(duration, size_mb)
+
+    for t in times:
+        frame = extract_frame(video_path, t)
+        if frame:
+            if check_image_nsfw(frame):
+                return True  # 🚨 return immediately if ANY frame is NSFW
+
+    return False