V1.1 -- added video filtering, updated readme

diff --git a/README.md b/README.md
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..ec35c04f251613229be4cfd462359d2dbc189a98 100644 (file)
--- a/README.md
+++ b/README.md
@@ -0,0 +1,23 @@
+## FHentaiBot
+---
+**FHentaiBot** -- это телеграм бот, для фильтрации NSFW контента. Для определения нежелательного контента используется нейросеть [Falconsai/nsfw_image_detection](https://huggingface.co/Falconsai/nsfw_image_detection)
+
+На данный момент поддерживается:
+- [x] Картинки
+- [x] Видео
+- [ ] GIF
+- [ ] Стикеры
+- [ ] Кружки
+
+Зависимости:
+- `python-telegram-bot`
+- `nest_asyncio`
+- `requests`
+- `pillow`
+- `transformers`
+- `torch`
+
+Для запуска впишите свой API токен в переменную `TOKEN`. Затем:
+```
+python fhentaibot
+```

diff --git a/fhentaibot.py b/fhentaibot.py
index 2bc553bda5457202c9505c7e1a8b697f8ad1093f..115d19ac57c3222ec172c74dcbc81a71c750254c 100755 (executable)
--- a/fhentaibot.py
+++ b/fhentaibot.py
@@ -51,11 +51,54 @@ async def analyze_image(update: Update, context) -> None:
         os.remove(file_name)
 
 
+async def analyze_video(update: Update, context) -> None:
+    file_id = update.message.video.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 video.')
+    else:
+        await update.message.reply_text('Failed to download the video.')
+
+    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)]:
+            img = Image.open(i)
+            classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
+            results = classifier(img)
+
+            if results[0]['label'] == '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
+
+    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()
     
     application.add_handler(CommandHandler('start', start))
     application.add_handler(MessageHandler(filters.PHOTO, analyze_image))
+    application.add_handler(MessageHandler(filters.VIDEO, analyze_video))
 
     await application.run_polling()