From 17d771060eb83e39e8d6646f1174bca82084004f Mon Sep 17 00:00:00 2001 From: Myned Date: Tue, 17 Sep 2019 02:59:35 -0400 Subject: [PATCH] Add support for SauceNAO as fallback reverse image search engine --- src/cogs/booru.py | 76 ++++++++++++++++++++++++++------------------ src/utils/scraper.py | 48 ++++++++++++++++++++-------- src/utils/utils.py | 2 +- 3 files changed, 80 insertions(+), 46 deletions(-) diff --git a/src/cogs/booru.py b/src/cogs/booru.py index 139258a..31d4b3a 100644 --- a/src/cogs/booru.py +++ b/src/cogs/booru.py @@ -106,8 +106,10 @@ class MsG(cmds.Cog): return args - def _get_score(self, score): - if score < 0: + def _get_icon(self, score): + if score is 'SauceNAO': + return 'https://d2.alternativeto.net/dist/icons/saucenao_23437.png?width=64&height=64&mode=crop&upscale=false' + elif score < 0: return 'https://emojipedia-us.s3.amazonaws.com/thumbs/320/twitter/103/pouting-face_1f621.png' elif score == 0: return 'https://emojipedia-us.s3.amazonaws.com/thumbs/320/mozilla/36/pile-of-poo_1f4a9.png' @@ -301,7 +303,7 @@ class MsG(cmds.Cog): embed.set_author(name=f'{post["width"]} x {post["height"]}', url=f'https://e621.net/post?tags=ratio:{post["width"]/post["height"]:.2f}', icon_url=ctx.author.avatar_url) embed.set_footer(text=post['score'], - icon_url=self._get_score(post['score'])) + icon_url=self._get_icon(post['score'])) except exc.MissingArgument: await ctx.send('\N{CROSS MARK} **Invalid url**') @@ -399,17 +401,21 @@ class MsG(cmds.Cog): try: await ctx.trigger_typing() - post = await scraper.get_post(url) + post, source, similarity = await scraper.get_post(url) embed = d.Embed( - title=', '.join(post['artist']), url=f'https://e621.net/post/show/{post["id"]}', color=ctx.me.color if isinstance(ctx.channel, d.TextChannel) else u.color) + title=', '.join(post['artist']), + url=source, + color=ctx.me.color if isinstance(ctx.channel, d.TextChannel) else u.color) embed.set_image(url=post['file_url']) - embed.set_author(name=f'{post["width"]} x {post["height"]}', - url=f'https://e621.net/post?tags=ratio:{post["width"]/post["height"]:.2f}', icon_url=ctx.author.avatar_url) - embed.set_footer(text=post['score'], - icon_url=self._get_score(post['score'])) + embed.set_author( + name=similarity, + icon_url=ctx.author.avatar_url) + embed.set_footer( + text=post['score'], + icon_url=self._get_icon(post['score'])) - await ctx.send('**Probable match**', embed=embed) + await ctx.send(embed=embed) c += 1 @@ -470,17 +476,21 @@ class MsG(cmds.Cog): try: await ctx.trigger_typing() - post = await scraper.get_post(url) + post, source, similarity = await scraper.get_post(url) embed = d.Embed( - title=', '.join(post['artist']), url=f'https://e621.net/post/show/{post["id"]}', color=ctx.me.color if isinstance(ctx.channel, d.TextChannel) else u.color) + title=', '.join(post['artist']), + url=source, + color=ctx.me.color if isinstance(ctx.channel, d.TextChannel) else u.color) embed.set_image(url=post['file_url']) - embed.set_author(name=f'{post["width"]} x {post["height"]}', - url=f'https://e621.net/post?tags=ratio:{post["width"]/post["height"]:.2f}', icon_url=ctx.author.avatar_url) + embed.set_author( + name=similarity, + icon_url=message.author.avatar_url) embed.set_footer( - text=post['score'], icon_url=self._get_score(post['score'])) + text=post['score'], + icon_url=self._get_icon(post['score'])) - await dest.send(f'**Probable match from** {message.author.display_name}', embed=embed) + await dest.send(embed=embed) await message.add_reaction('\N{WHITE HEAVY CHECK MARK}') if remove: @@ -529,17 +539,21 @@ class MsG(cmds.Cog): try: await message.channel.trigger_typing() - post = await scraper.get_post(url) + post, source, similarity = await scraper.get_post(url) embed = d.Embed( - title=', '.join(post['artist']), url=f'https://e621.net/post/show/{post["id"]}', color=message.channel.guild.me.color if isinstance(message.channel, d.TextChannel) else u.color) + title=', '.join(post['artist']), + url=source, + color=message.channel.guild.me.color if isinstance(message.channel, d.TextChannel) else u.color) embed.set_image(url=post['file_url']) - embed.set_author(name=f'{post["width"]} x {post["height"]}', - url=f'https://e621.net/post?tags=ratio:{post["width"]/post["height"]:.2f}', icon_url=message.author.avatar_url) - embed.set_footer(text=post['score'], - icon_url=self._get_score(post['score'])) + embed.set_author( + name=similarity, + icon_url=message.author.avatar_url) + embed.set_footer( + text=post['score'], + icon_url=self._get_icon(post['score'])) - await message.channel.send('**Probable match from** {}'.format(message.author.display_name), embed=embed) + await message.channel.send(embed=embed) await message.add_reaction('\N{WHITE HEAVY CHECK MARK}') @@ -768,7 +782,7 @@ class MsG(cmds.Cog): embed.set_author(name=pool['name'], url='https://e621.net/pool/show?id={}'.format(pool['id']), icon_url=ctx.author.avatar_url) embed.set_footer(text='{} / {}'.format(c, len(posts)), - icon_url=self._get_score(values[c - 1]['score'])) + icon_url=self._get_icon(values[c - 1]['score'])) paginator = await ctx.send(embed=embed) @@ -799,7 +813,7 @@ class MsG(cmds.Cog): embed.url = 'https://e621.net/post/show/{}'.format( keys[c - 1]) embed.set_footer(text='{} / {}'.format(c, len(posts)), - icon_url=self._get_score(values[c - 1]['score'])) + icon_url=self._get_icon(values[c - 1]['score'])) embed.set_image(url=values[c - 1]['file_url']) await paginator.edit(content='\N{HEAVY BLACK HEART}' if keys[c - 1] in hearted.keys() else None, embed=embed) @@ -817,7 +831,7 @@ class MsG(cmds.Cog): embed.url = 'https://e621.net/post/show/{}'.format( keys[c - 1]) embed.set_footer(text='{} / {}'.format(c, len(posts)), - icon_url=self._get_score(values[c - 1]['score'])) + icon_url=self._get_icon(values[c - 1]['score'])) embed.set_image(url=values[c - 1]['file_url']) if ctx.channel is d.TextChannel: @@ -833,7 +847,7 @@ class MsG(cmds.Cog): embed.url = 'https://e621.net/post/show/{}'.format( keys[c - 1]) embed.set_footer(text='{} / {}'.format(c, len(posts)), - icon_url=self._get_score(values[c - 1]['score'])) + icon_url=self._get_icon(values[c - 1]['score'])) embed.set_image(url=values[c - 1]['file_url']) await paginator.edit(content='\N{HEAVY BLACK HEART}' if keys[c - 1] in hearted.keys() else None, embed=embed) @@ -904,7 +918,7 @@ class MsG(cmds.Cog): embed.set_author(name=' '.join(tags) if tags else order, url='https://{}.net/post?tags={}'.format(booru, ','.join(tags)), icon_url=ctx.author.avatar_url) embed.set_footer(text=values[c - 1]['score'], - icon_url=self._get_score(values[c - 1]['score'])) + icon_url=self._get_icon(values[c - 1]['score'])) paginator = await ctx.send(embed=embed) @@ -936,7 +950,7 @@ class MsG(cmds.Cog): booru, keys[c - 1]) embed.set_footer(text=values[c - 1]['score'], - icon_url=self._get_score(values[c - 1]['score'])) + icon_url=self._get_icon(values[c - 1]['score'])) embed.set_image(url=values[c - 1]['file_url']) await paginator.edit(content='\N{HEAVY BLACK HEART}' if keys[c - 1] in hearted.keys() else None, embed=embed) @@ -960,7 +974,7 @@ class MsG(cmds.Cog): booru, keys[c - 1]) embed.set_footer(text=values[c - 1]['score'], - icon_url=self._get_score(values[c - 1]['score'])) + icon_url=self._get_icon(values[c - 1]['score'])) embed.set_image(url=values[c - 1]['file_url']) await paginator.edit(content='\N{HEAVY BLACK HEART}' if keys[c - 1] in hearted.keys() else None, embed=embed) @@ -1035,7 +1049,7 @@ class MsG(cmds.Cog): embed.set_author(name=' '.join(tags) if tags else order, url='https://{}.net/post?tags={}'.format(booru, ','.join(tags)), icon_url=ctx.author.avatar_url) embed.set_footer( - text=post['score'], icon_url=self._get_score(post['score'])) + text=post['score'], icon_url=self._get_icon(post['score'])) message = await ctx.send(embed=embed) diff --git a/src/utils/scraper.py b/src/utils/scraper.py index 6fd0266..702bb50 100644 --- a/src/utils/scraper.py +++ b/src/utils/scraper.py @@ -16,31 +16,51 @@ async def get_post(url): if filesize > 8192 * 1024: raise exc.SizeError(size(filesize, system=alternative)) - except (ValueError, KeyError): - raise exc.MissingArgument + content = await u.fetch('http://iqdb.harry.lu', params={'url': url}) + soup = BeautifulSoup(content, 'html.parser') + source = soup.find_all('a', limit=2)[1].get('href') - content = await u.fetch('http://iqdb.harry.lu', params={'url': url}) - - try: - value = BeautifulSoup(content, 'html.parser').find_all('a')[1].get('href') - if value != '#': - ident = re.search('show/([0-9]+)', value).group(1) + if source != '#': + ident = re.search('show/([0-9]+)', source).group(1) post = await u.fetch('http://e621.net/post/show.json', params={'id': ident}, json=True) - if (post['status'] == 'deleted'): ident = re.search('#(\\d+)', post['delreason']).group(1) post = await u.fetch('http://e621.net/post/show.json', params={'id': ident}, json=True) + source = f'https://e621.net/post/show/{post["id"]}' + similarity = re.search('\\d+', soup.find(string=re.compile('similarity'))).group(0) + '% Match' - return post + return post, source, similarity else: raise IndexError except IndexError: - try: - raise exc.MatchError(re.search('\\/([^\\/]+)$', url).group(1)) + content = await u.fetch( + 'https://saucenao.com/search.php', + params={ + 'url': url, + 'api_key': u.config['saucenao_api'], + 'output_type': 2}, + json=True) + result = content['results'][0] + if 'author_name' in result['data']: + artist = 'author_name' + elif 'member_name' in result['data']: + artist = 'member_name' + else: + artist = 'creator' + post = { + 'file_url': result['header']['thumbnail'], + 'artist': [result['data'][artist]], + 'score': 'SauceNAO'} + source = result['data']['ext_urls'][0] + similarity = re.search('(\\d+)\\.', result['header']['similarity']).group(1) + '% Match' - except AttributeError: - raise exc.MissingArgument + return post, source, similarity + + raise exc.MatchError(re.search('\\/([^\\/]+)$', url).group(1)) + + except (AttributeError, ValueError, KeyError): + raise exc.MissingArgument async def get_image(url): diff --git a/src/utils/utils.py b/src/utils/utils.py index c01ce63..e907744 100644 --- a/src/utils/utils.py +++ b/src/utils/utils.py @@ -26,7 +26,7 @@ try: except FileNotFoundError: with open('config.json', 'w') as outfile: jsn.dump({'client_id': 0, 'owner_id': 0, 'permissions': 126016, - 'playing': 'a game', 'prefix': [',', 'm,'], 'selfbot': False, 'token': 'str'}, outfile, indent=4, sort_keys=True) + 'playing': 'a game', 'prefix': [',', 'm,'], 'selfbot': False, 'token': 'str', 'saucenao_api': 'str'}, outfile, indent=4, sort_keys=True) print('FILE NOT FOUND : config.json created with abstract values. Restart run.py with correct values')