Create utils.py

utils.py

@@ -0,0 +1,135 @@
+import cv2
+import tempfile
+import requests
+import os
+from PIL import Image
+from transformers import pipeline
+import torch
+
+# 🔥 SPEED BOOST SETTINGS
+torch.set_grad_enabled(False)
+torch.set_num_threads(2)
+
+# 🔥 Faster NSFW model
+classifier = pipeline(
+    "image-classification",
+    model="AdamCodd/vit-base-nsfw-detector",
+    device=-1  # CPU
+)
+
+# -----------------------------
+# Download with retry + headers (FIX CATBOX)
+# -----------------------------
+def download_file(url):
+    headers = {
+        "User-Agent": "Mozilla/5.0",
+        "Accept": "*/*",
+        "Connection": "keep-alive",
+        "Range": "bytes=0-"
+    }
+
+    for _ in range(3):  # retry
+        try:
+            response = requests.get(
+                url,
+                headers=headers,
+                stream=True,
+                timeout=10
+            )
+
+            if response.status_code != 200:
+                continue
+
+            tmp = tempfile.NamedTemporaryFile(delete=False)
+
+            for chunk in response.iter_content(1024 * 1024):
+                if chunk:
+                    tmp.write(chunk)
+
+            tmp.close()
+            return tmp.name
+
+        except requests.exceptions.RequestException:
+            continue
+
+    raise Exception("Failed to fetch file")
+
+
+# -----------------------------
+# Video duration
+# -----------------------------
+def get_video_duration(video_path):
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
+    cap.release()
+    return frames / fps if fps > 0 else 0
+
+
+# -----------------------------
+# Extract frame
+# -----------------------------
+def extract_frame(video_path, second):
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS)
+
+    frame_no = int(fps * second)
+    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_no)
+
+    success, frame = cap.read()
+    cap.release()
+
+    if not success:
+        return None
+
+    tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
+    cv2.imwrite(tmp.name, frame)
+    return tmp.name
+
+
+# -----------------------------
+# FAST frame selection
+# -----------------------------
+def get_frame_times(duration):
+    if duration <= 3:
+        return [1]
+
+    elif duration <= 10:
+        return [2]
+
+    else:
+        return [3, 8]  # max 2 frames (FAST)
+
+
+# -----------------------------
+# Image NSFW check (OPTIMIZED)
+# -----------------------------
+def check_image_nsfw(image_path):
+    img = Image.open(image_path).convert("RGB")
+
+    # 🔥 Resize = BIG SPEED BOOST
+    img = img.resize((224, 224))
+
+    result = classifier(img)
+
+    for r in result:
+        if r["label"].lower() == "nsfw" and r["score"] > 0.5:
+            return True
+
+    return False
+
+
+# -----------------------------
+# Video NSFW check
+# -----------------------------
+def check_video_nsfw(video_path):
+    duration = get_video_duration(video_path)
+    times = get_frame_times(duration)
+
+    for t in times:
+        frame = extract_frame(video_path, t)
+        if frame:
+            if check_image_nsfw(frame):
+                return True  # 🚨 stop early
+
+    return False