You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
151 lines
5.0 KiB
151 lines
5.0 KiB
import pathlib
|
|
|
|
import discord
|
|
from nudenet import NudeClassifier
|
|
from redbot.core import Config, commands
|
|
from redbot.core.bot import Red
|
|
from redbot.core.data_manager import cog_data_path
|
|
|
|
|
|
class Nudity(commands.Cog):
|
|
"""Monitor images for NSFW content and moves them to a nsfw channel if possible"""
|
|
|
|
def __init__(self, bot: Red):
|
|
super().__init__()
|
|
self.bot = bot
|
|
self.config = Config.get_conf(self, identifier=9811198108111121, force_registration=True)
|
|
|
|
default_guild = {"enabled": False, "channel_id": None}
|
|
|
|
self.config.register_guild(**default_guild)
|
|
|
|
# self.detector = NudeDetector()
|
|
self.classifier = NudeClassifier()
|
|
|
|
self.data_path: pathlib.Path = cog_data_path(self)
|
|
|
|
self.current_processes = 0
|
|
|
|
async def red_delete_data_for_user(self, **kwargs):
|
|
"""Nothing to delete"""
|
|
return
|
|
|
|
@commands.command(aliases=["togglenudity"], name="nudity")
|
|
async def nudity(self, ctx: commands.Context):
|
|
"""Toggle nude-checking on or off"""
|
|
is_on = await self.config.guild(ctx.guild).enabled()
|
|
await self.config.guild(ctx.guild).enabled.set(not is_on)
|
|
await ctx.send("Nude checking is now set to {}".format(not is_on))
|
|
|
|
@commands.command()
|
|
async def nsfwchannel(self, ctx: commands.Context, channel: discord.TextChannel = None):
|
|
if channel is None:
|
|
await self.config.guild(ctx.guild).channel_id.set(None)
|
|
await ctx.send("NSFW Channel cleared")
|
|
else:
|
|
if not channel.is_nsfw():
|
|
await ctx.send("This channel isn't NSFW!")
|
|
return
|
|
else:
|
|
await self.config.guild(ctx.guild).channel_id.set(channel.id)
|
|
await ctx.send("NSFW channel has been set to {}".format(channel.mention))
|
|
|
|
async def get_nsfw_channel(self, guild: discord.Guild):
|
|
channel_id = await self.config.guild(guild).channel_id()
|
|
|
|
if channel_id is None:
|
|
return None
|
|
else:
|
|
return guild.get_channel(channel_id=channel_id)
|
|
|
|
async def nsfw(self, message: discord.Message, images: dict):
|
|
content = message.content
|
|
guild: discord.Guild = message.guild
|
|
if not content:
|
|
content = "*`None`*"
|
|
try:
|
|
await message.delete()
|
|
except discord.Forbidden:
|
|
await message.channel.send("NSFW Image detected!")
|
|
return
|
|
|
|
embed = discord.Embed(title="NSFW Image Detected")
|
|
embed.add_field(name="Original Message", value=content)
|
|
embed.set_author(name=message.author.name, icon_url=message.author.avatar_url)
|
|
await message.channel.send(embed=embed)
|
|
|
|
nsfw_channel = await self.get_nsfw_channel(guild)
|
|
|
|
if nsfw_channel is None:
|
|
return
|
|
else:
|
|
for image, r in images.items():
|
|
if r["unsafe"] > 0.7:
|
|
await nsfw_channel.send(
|
|
"NSFW Image from {}".format(message.channel.mention),
|
|
file=discord.File(
|
|
image,
|
|
),
|
|
)
|
|
|
|
@commands.Cog.listener()
|
|
async def on_message(self, message: discord.Message):
|
|
is_private = isinstance(message.channel, discord.abc.PrivateChannel)
|
|
|
|
if not message.attachments or is_private or message.author.bot:
|
|
# print("did not qualify")
|
|
return
|
|
|
|
if await self.bot.cog_disabled_in_guild(self, message.guild):
|
|
return
|
|
|
|
try:
|
|
is_on = await self.config.guild(message.guild).enabled()
|
|
except AttributeError:
|
|
return
|
|
|
|
if not is_on:
|
|
print("Not on")
|
|
return
|
|
|
|
channel: discord.TextChannel = message.channel
|
|
|
|
if channel.is_nsfw():
|
|
print("nsfw channel is okay")
|
|
return
|
|
|
|
check_list = []
|
|
for attachment in message.attachments:
|
|
# async with aiohttp.ClientSession() as session:
|
|
# img = await fetch_img(session, attachment.url)
|
|
|
|
ext = attachment.filename
|
|
|
|
temp_name = self.data_path / f"nudecheck{self.current_processes}_{ext}"
|
|
|
|
self.current_processes += 1
|
|
|
|
print("Pre attachment save")
|
|
await attachment.save(temp_name)
|
|
check_list.append(temp_name)
|
|
|
|
print("Pre nude check")
|
|
# nude_results = self.detector.detect(temp_name)
|
|
nude_results = self.classifier.classify([str(n) for n in check_list])
|
|
# print(nude_results)
|
|
|
|
if True in [r["unsafe"] > 0.7 for r in nude_results.values()]:
|
|
# print("Is nude")
|
|
await message.add_reaction("❌")
|
|
await self.nsfw(message, nude_results)
|
|
else:
|
|
# print("Is not nude")
|
|
await message.add_reaction("✅")
|
|
|
|
|
|
# async def fetch_img(session, url):
|
|
# with aiohttp.Timeout(10):
|
|
# async with session.get(url) as response:
|
|
# assert response.status == 200
|
|
# return await response.read()
|