Add NSFW plugin, requirements, and admin-only /nsfw command

DragMusic/plugins/plugins/nsfw.py

@@ -0,0 +1,147 @@
+import os
+import tempfile
+import subprocess
+from PIL import Image, UnidentifiedImageError
+import torch
+import torch.nn.functional as F
+from pyrogram import Client, filters
+from pyrogram.types import Message
+from transformers import AutoModelForImageClassification, ViTImageProcessor
+from DragMusic import app
+from motor.motor_asyncio import AsyncIOMotorClient
+from pyrogram.errors import PeerIdInvalid, UserNotParticipant
+from config import MONGO_DB_URI
+
+# Initialize MongoDB client and select database & collection
+db_client = AsyncIOMotorClient(MONGO_DB_URI)
+db = db_client['AnonXDatabase']
+nsfw_chats_collection = db['nsfw_chats']
+
+# Logger setup
+LOGGER_ID = -1002471976725
+LOGGER_TOPIC_ID = 6314
+
+# Load model and processor
+model = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection")
+processor = ViTImageProcessor.from_pretrained("Falconsai/nsfw_image_detection")
+
+async def is_nsfw_enabled(chat_id: int) -> bool:
+    chat = await nsfw_chats_collection.find_one({"chat_id": chat_id})
+    return bool(chat)
+
+async def enable_nsfw(chat_id: int):
+    await nsfw_chats_collection.update_one({"chat_id": chat_id}, {"$set": {"chat_id": chat_id}}, upsert=True)
+
+async def disable_nsfw(chat_id: int):
+    await nsfw_chats_collection.delete_one({"chat_id": chat_id})
+
+@app.on_message(filters.command("nsfw") & filters.group)
+async def toggle_nsfw_scan(_, message: Message):
+    if len(message.command) != 2:
+        return await message.reply("Usage: /nsfw enable or /nsfw disable")
+
+    try:
+        member = await app.get_chat_member(message.chat.id, message.from_user.id)
+        if not (member.status in ["administrator", "owner"]):
+            return await message.reply("Only group admins and the group owner can use this command.")
+    except (PeerIdInvalid, UserNotParticipant):
+        return await message.reply("You must be a group member to use this command.")
+
+    arg = message.command[1].lower()
+    if arg == "enable":
+        await enable_nsfw(message.chat.id)
+        await message.reply("Enabled NSFW scanning for this group.")
+    elif arg == "disable":
+        await disable_nsfw(message.chat.id)
+        await message.reply("Disabled NSFW scanning for this group.")
+    else:
+        await message.reply("Invalid argument. Use /nsfw enable or /nsfw disable.")
+        
+@app.on_message(filters.command("scan") & filters.reply)
+async def scan_command(_, message: Message):
+    if not message.reply_to_message:
+        return await message.reply("Reply to an image, video, or sticker to scan.")
+    await scan_media(message.reply_to_message, message)
+
+@app.on_message(filters.group & (filters.photo | filters.video | filters.sticker))
+async def auto_nsfw_scan(client: Client, message: Message):
+    if await is_nsfw_enabled(message.chat.id):
+        await scan_media(message, message)
+
+async def scan_media(media_msg: Message, reply_msg: Message):
+    media = media_msg.photo or media_msg.video or media_msg.sticker
+    path = None
+    image_path = None
+
+    try:
+        ext = ".jpg"
+        if media_msg.video:
+            ext = ".mp4"
+        elif media_msg.sticker:
+            ext = ".webp"
+
+        with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as temp:
+            path = temp.name
+
+        if hasattr(media, "download"):
+            await media.download(file_name=path)
+        elif media.file_id:
+            await media_msg.download(file_name=path)
+        else:
+            return await reply_msg.reply("Unable to download media.")
+
+        if path.endswith((".webm", ".mp4", ".webp")):
+            image_path = path + ".png"
+            subprocess.run(["ffmpeg", "-i", path, "-frames:v", "1", image_path, "-y"],
+                           stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+            if not os.path.exists(image_path) or os.path.getsize(image_path) == 0:
+                return await reply_msg.reply("Failed to extract image frame.")
+            img = Image.open(image_path).convert("RGB")
+        else:
+            if os.path.getsize(path) == 0:
+                return await reply_msg.reply("Downloaded file is empty, cannot scan.")
+            img = Image.open(path).convert("RGB")
+
+        inputs = processor(images=img, return_tensors="pt")
+        with torch.no_grad():
+            outputs = model(**inputs)
+            probs = F.softmax(outputs.logits, dim=1)
+            confidence = round(probs[0][1].item() * 100, 2)
+            label = "NSFW" if confidence > 50 else "SFW"
+
+        if label == "NSFW":
+            try:
+                await media_msg.delete()
+            except Exception as e:
+                print(f"[NSFW] Failed to delete message: {e}")
+
+            if confidence > 50:
+                await reply_msg.reply(f"NSFW Detected ({confidence}%) — Message deleted.")
+                log_text = (
+                    f"NSFW Content Detected\n"
+                    f"Chat: {reply_msg.chat.title} ({reply_msg.chat.id})\n"
+                    f"User: {media_msg.from_user.mention if media_msg.from_user else 'Unknown'}\n"
+                    f"Confidence: {confidence}%"
+                )
+                try:
+                    await app.send_message(
+                        chat_id=LOGGER_ID,
+                        text=log_text,
+                        message_thread_id=LOGGER_TOPIC_ID
+                    )
+                except Exception as log_err:
+                    print(f"[NSFW Logger Error] {log_err}")
+        elif confidence > 50:
+            await reply_msg.reply(f"Scan Result: {label} ({confidence}%)")
+
+    except UnidentifiedImageError:
+        print(f"[NSFW Detection Error] cannot identify image file '{path or image_path}'")
+        await reply_msg.reply("Failed to read image for scanning.")
+    except Exception as e:
+        print(f"[NSFW Detection Error] {e}")
+        await reply_msg.reply("Error occurred during NSFW scan.")
+    finally:
+        if path and os.path.exists(path):
+            os.remove(path)
+        if image_path and os.path.exists(image_path):
+            os.remove(image_path) 

requirements.txt

@@ -41,3 +41,8 @@ bing_image_downloader
 openai==0.28
 uvicorn
 fastapi
+torch
+transformers
+pillow
+# ffmpeg is installed at the system level, but if you want a Python wrapper, uncomment the next line:
+# ffmpeg-python