skip image duplicates

This commit is contained in:
bakatrouble 2022-08-08 04:35:45 +03:00
parent 76dfae89ad
commit f296ea4121

View File

@ -5,6 +5,7 @@ import tempfile
from io import BytesIO from io import BytesIO
from uuid import uuid4 from uuid import uuid4
import imagehash
import requests import requests
from PIL import Image from PIL import Image
from django.db import models from django.db import models
@ -43,6 +44,10 @@ class ChannelHelperBotModuleConfig(TelegramBotModuleConfig):
except: except:
raise RuntimeError('Could not load image') raise RuntimeError('Could not load image')
im = Image.open(f) # type: Image.Image im = Image.open(f) # type: Image.Image
if self.queued_items.filter(image_hash=imagehash.phash(im), type='photo').count() > 0:
return False
width, height = im.size width, height = im.size
if width > 2000 or height > 2000: if width > 2000 or height > 2000:
im.thumbnail((2000, 2000)) im.thumbnail((2000, 2000))
@ -59,10 +64,11 @@ class ChannelHelperBotModuleConfig(TelegramBotModuleConfig):
return True return True
def periodic_task(self, bot: Bot): def periodic_task(self, bot: Bot):
i = self.queued_items.order_by('?').first() # type: QueuedItem i = self.queued_items.filter(processed=False).order_by('?').first() # type: QueuedItem
if i: if i:
i.send(bot) i.send(bot)
i.delete() i.processed = True
i.save()
def handle_message(self, update: Update, ctx: CallbackContext): def handle_message(self, update: Update, ctx: CallbackContext):
if self.users.count() and not self.users.filter(user_id=update.effective_user.id).count(): if self.users.count() and not self.users.filter(user_id=update.effective_user.id).count():
@ -84,6 +90,12 @@ class ChannelHelperBotModuleConfig(TelegramBotModuleConfig):
p = m.photo p = m.photo
i.type = 'photo' i.type = 'photo'
i.args = json.dumps([p[-1].file_id]) i.args = json.dumps([p[-1].file_id])
io = BytesIO()
io.name = 'file.jpg'
bot.get_file(p[-1]).download(out=io)
im = Image.open(io)
if self.queued_items.filter(image_hash=imagehash.phash(im), type='photo').count() > 0:
return
elif hasattr(m, 'sticker') and m.sticker: elif hasattr(m, 'sticker') and m.sticker:
s = m.sticker s = m.sticker
i.type = 'sticker' i.type = 'sticker'
@ -152,6 +164,7 @@ class QueuedItem(models.Model):
args = models.TextField() args = models.TextField()
message_id = models.PositiveBigIntegerField(default=None, db_index=True, null=True, blank=True) message_id = models.PositiveBigIntegerField(default=None, db_index=True, null=True, blank=True)
image_hash = models.CharField(max_length=64, null=True, blank=True) image_hash = models.CharField(max_length=64, null=True, blank=True)
processed = models.BooleanField(default=False)
def send(self, bot: Bot): def send(self, bot: Bot):
getattr(bot, 'send_' + self.type)(self.config.chat_id, *json.loads(self.args)) getattr(bot, 'send_' + self.type)(self.config.chat_id, *json.loads(self.args))