V1.2 -- added animation and circle filtering

diff --git a/README.md b/README.md
index ebe2757da53c06f58faaa3ae83ad16add3a37d6e..499f8d719490b52d14863ae8f7c151da4f3d30fc 100644 (file)
--- a/README.md
+++ b/README.md
@@ -5,9 +5,9 @@
 На данный момент поддерживается:
 - [x] Картинки
 - [x] Видео
-- [ ] GIF
+- [x] GIF
+- [x] Кружки
 - [ ] Стикеры
-- [ ] Кружки
 
 Зависимости:
 - `python-telegram-bot`

diff --git a/fhentaibot.py b/fhentaibot.py
index 69b30b7cea10dabe323d3d3df0fe3a851f272f93..5d075c0f8a28558d1978cac08e11022647b4d9e4 100755 (executable)
--- a/fhentaibot.py
+++ b/fhentaibot.py
@@ -15,6 +15,13 @@ async def start(update: Update, context) -> None:
     await update.message.reply_text('Hi! I\'m your NSFW Image Detector Bot. Send me an image to check if it contains NSFW content.')
 
 
+async def analyze(file_name):
+        img = Image.open(file_name)
+        classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
+        result = classifier(img)
+        return result[0]['label']
+
+
 async def analyze_image(update: Update, context) -> None:
     file_id = update.message.photo[-1].file_id
     file_info = await context.bot.get_file(file_id)
@@ -35,11 +42,9 @@ async def analyze_image(update: Update, context) -> None:
         await update.message.reply_text('Failed to download the photo.')
 
     try:
-        img = Image.open(file_name)
-        classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
-        results = classifier(img)
+        result = await analyze(file_name)
 
-        if results[0]['label'] == 'nsfw':
+        if result == 'nsfw':
             await update.message.reply_text(f"Warning: This image contains NSFW content.")
             await context.bot.delete_message(chat_id=update.message.chat_id,message_id=update.message.message_id)
 
@@ -74,11 +79,9 @@ async def analyze_video(update: Update, context) -> None:
         ffmpeg.input(file_name).filter('fps', fps='1').output(f'{file_name}-%d.jpg', start_number=0).overwrite_output().run(quiet=True)
 
         for i in [i for i in os.listdir('.') if (file_name in i) and ('.jpg' in i)]:
-            img = Image.open(i)
-            classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
-            results = classifier(img)
+            result = await analyze(i)
 
-            if results[0]['label'] == 'nsfw':
+            if result == 'nsfw':
                 await update.message.reply_text(f"Warning: This video contains NSFW content.")
                 await context.bot.delete_message(chat_id=update.message.chat_id,message_id=update.message.message_id)
                 break
@@ -92,6 +95,83 @@ async def analyze_video(update: Update, context) -> None:
             os.remove(i)
 
 
+async def analyze_gif(update: Update, context) -> None:
+    file_id = update.message.animation.file_id
+    file_info = await context.bot.get_file(file_id)
+
+    if file_info:
+         # Get the URL of the file
+        file_path = file_info.file_path
+        response = requests.get(file_path)
+
+        if response.status_code == 200:
+            # Save the video to disk
+            file_name = f"{file_id}.mp4"
+            with open(file_name, 'wb') as f:
+                f.write(response.content)
+        else:
+            await update.message.reply_text('Failed to download the animation.')
+    else:
+        await update.message.reply_text('Failed to download the animation.')
+
+    try:
+        ffmpeg.input(file_name).filter('fps', fps='1').output(f'{file_name}-%d.jpg', start_number=0).overwrite_output().run(quiet=True)
+
+        for i in [i for i in os.listdir('.') if (file_name in i) and ('.jpg' in i)]:
+            result = await analyze(i)
+
+            if result == 'nsfw':
+                await update.message.reply_text(f"Warning: This animation contains NSFW content.")
+                await context.bot.delete_message(chat_id=update.message.chat_id,message_id=update.message.message_id)
+                break
+
+    except Exception as e:
+        await update.message.reply_text(f'Some kind of mistake')
+        print(f"An error occurred: {e}")
+
+    finally:
+        for i in [i for i in os.listdir('.') if file_name in i]:
+            os.remove(i)
+
+
+async def analyze_video_note(update: Update, context) -> None:
+    file_id = update.message.video_note.file_id
+    file_info = await context.bot.get_file(file_id)
+
+    if file_info:
+         # Get the URL of the file
+        file_path = file_info.file_path
+        response = requests.get(file_path)
+
+        if response.status_code == 200:
+            # Save the video to disk
+            file_name = f"{file_id}.mp4"
+            with open(file_name, 'wb') as f:
+                f.write(response.content)
+        else:
+            await update.message.reply_text('Failed to download the circle.')
+    else:
+        await update.message.reply_text('Failed to download the circle.')
+
+    try:
+        ffmpeg.input(file_name).filter('fps', fps='1').output(f'{file_name}-%d.jpg', start_number=0).overwrite_output().run(quiet=True)
+
+        for i in [i for i in os.listdir('.') if (file_name in i) and ('.jpg' in i)]:
+            result = await analyze(i)
+
+            if result == 'nsfw':
+                await update.message.reply_text(f"Warning: This circle contains NSFW content.")
+                await context.bot.delete_message(chat_id=update.message.chat_id,message_id=update.message.message_id)
+                break
+
+    except Exception as e:
+        await update.message.reply_text(f'Some kind of mistake')
+        print(f"An error occurred: {e}")
+
+    finally:
+        for i in [i for i in os.listdir('.') if file_name in i]:
+            os.remove(i)
+
 
 async def main() -> None:
     application = ApplicationBuilder().token(TOKEN).build()
@@ -99,6 +179,8 @@ async def main() -> None:
     application.add_handler(CommandHandler('start', start))
     application.add_handler(MessageHandler(filters.PHOTO, analyze_image))
     application.add_handler(MessageHandler(filters.VIDEO, analyze_video))
+    application.add_handler(MessageHandler(filters.ANIMATION, analyze_gif))
+    application.add_handler(MessageHandler(filters.VIDEO_NOTE, analyze_video_note))
 
     await application.run_polling()