Update waifu_name.py: use image hash matching for waifu name lookup

DragMusic/plugins/management/waifu_name.py

@@ -1,23 +1,32 @@
 from pyrogram import filters
 from DragMusic import app
 from DragMusic.core.userbot import Userbot
+from PIL import Image
+import imagehash
+import io
 
-WAIFU_CHANNEL_ID = "Grabber_Database"  # Use channel username instead of numeric ID
+WAIFU_CHANNEL_ID = "Grabber_Database"  # Use channel username
 
-async def get_waifu_name_by_file_id(file_id):
+async def get_waifu_name_by_image_hash(target_hash):
     userbot = Userbot().one
     async with userbot:
         async for msg in userbot.get_chat_history(WAIFU_CHANNEL_ID):
-            if msg.photo and msg.photo.file_id == file_id:
-                return msg.caption  # This is the waifu's name
+            if msg.photo:
+                file = await userbot.download_media(msg, in_memory=True)
+                img = Image.open(io.BytesIO(file.getvalue()))
+                waifu_hash = imagehash.average_hash(img)
+                if waifu_hash - target_hash < 3:  # 0 is identical, <3 is very similar
+                    return msg.caption
     return None
 
 @app.on_message(filters.command("waifus") & filters.reply)
 async def waifus_name(client, message):
     if not message.reply_to_message or not message.reply_to_message.photo:
         return await message.reply("Reply to a waifu image to get the name!")
-    file_id = message.reply_to_message.photo.file_id
-    waifu_name = await get_waifu_name_by_file_id(file_id)
+    file = await message.reply_to_message.download(in_memory=True)
+    img = Image.open(io.BytesIO(file.getvalue()))
+    target_hash = imagehash.average_hash(img)
+    waifu_name = await get_waifu_name_by_image_hash(target_hash)
     if waifu_name:
         await message.reply(f"The waifu's name is: {waifu_name}")
     else:

requirements.txt

@@ -44,5 +44,6 @@ fastapi
 torch
 transformers
 pillow
+imagehash
 # ffmpeg is installed at the system level, but if you want a Python wrapper, uncomment the next line:
 # ffmpeg-python