Skip to content

Commit

Permalink
[reverseimagesearch] implement tracemoe
Browse files Browse the repository at this point in the history
  • Loading branch information
@fixator10
fixator10 committed May 19, 2019
1 parent ad895df commit c9ee006
Show file tree
Hide file tree
Showing 3 changed files with 107 additions and 58 deletions.
3 changes: 2 additions & 1 deletion reverseimagesearch/info.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
"saucenao"
],
"requirements": [
"python-dateutil"
"python-dateutil",
"pillow>=6"
]
}
145 changes: 90 additions & 55 deletions reverseimagesearch/reverseimagesearch.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,9 @@
from redbot.core.i18n import Translator, cog_i18n
from redbot.core.utils.menus import menu, DEFAULT_CONTROLS

# from .tracemoe import TraceMoe
from .converters import ImageFinder
from .saucenao import SauceNAO
from .tracemoe import TraceMoe

_ = Translator("ReverseImageSearch", __file__)

Expand Down Expand Up @@ -116,9 +116,9 @@ async def maxres(self, ctx, results: int = 6):
await self.config.numres.set(results)
await ctx.tick()

@saucenao.command()
@saucenao.command(name="stats")
@commands.is_owner()
async def stats(self, ctx):
async def saucenao_stats(self, ctx):
"""See how many requests are left"""
if any(limit is not None for limit in self.saucenao_limits.values()):
await ctx.send(
Expand All @@ -138,55 +138,90 @@ async def stats(self, ctx):
_("Command `{}` has not been used yet").format(self.saucenao)
)

# @commands.group(invoke_without_command=True, aliases=["WAIT"])
# async def tracemoe(self, ctx, image: ImageFinder = None):
# """Reverse search image via WAIT"""
# if image is None:
# try:
# image = await ImageFinder().search_for_images(ctx)
# except ValueError as e:
# await ctx.send(e)
# return
# image = image[0]
# try:
# search = await TraceMoe.from_image(ctx, image)
# except ValueError as e:
# await ctx.send(e)
# return
# embeds = []
# page = 0
# for doc in search.docs:
# page += 1
# e = Embed(
# title=doc.title,
# description="\n".join(
# [
# doc.title_native
# and "🇯🇵" + _("Native title: {}").format(doc.title_native)
# or "",
# doc.title_romaji
# and "🇯🇵"
# + _("Romaji transcription: {}").format(doc.title_romaji)
# or "",
# doc.title_chinese
# and "🇨🇳" + _("Chinese title: {}").format(doc.title_chinese)
# or "",
# doc.title_english
# and "🇺🇸" + _("English title: {}").format(doc.title_english)
# or "",
# ]
# ),
# url=doc.mal_id
# and f"https://myanimelist.net/anime/{doc.mal_id}"
# or f"https://anilist.co/anime/{doc.anilist_id}",
# color=await ctx.embed_color(),
# )
# e.set_thumbnail(url=doc.thumbnail)
# e.set_footer(
# text=_("Via WAIT (trace.moe) • Page {}/{}").format(
# page, len(search.docs)
# ),
# icon_url="https://trace.moe/favicon128.png",
# )
# embeds.append(e)
# await menu(ctx, embeds, DEFAULT_CONTROLS)
@commands.group(invoke_without_command=True, aliases=["WAIT"])
async def tracemoe(self, ctx, image: ImageFinder = None):
"""Reverse search image via WAIT
If search performed not in NSFW channel, NSFW results will be not shown"""
if image is None:
try:
image = await ImageFinder().search_for_images(ctx)
except ValueError as e:
await ctx.send(e)
return
image = image[0]
try:
search = await TraceMoe.from_image(ctx, image)
except ValueError as e:
await ctx.send(e)
return
embeds = []
page = 0
for doc in search.docs:
page += 1
if ctx.channel.nsfw and doc.is_adult:
continue
# this design is kinda shit too, ideas, plssss
e = Embed(
title=doc.title,
description="\n".join(
[
s
for s in [
_("Similarity: {:.2f}%").format(doc.similarity * 100),
doc.title_native
and "🇯🇵 " + _("Native title: {}").format(doc.title_native),
doc.title_romaji
and "🇯🇵 "
+ _("Romaji transcription: {}").format(doc.title_romaji),
doc.title_chinese
and "🇨🇳 "
+ _("Chinese title: {}").format(doc.title_chinese),
doc.title_english
and "🇺🇸 "
+ _("English title: {}").format(doc.title_english),
_("Est. Time: {}").format(doc.time_str),
_("Episode: {}").format(doc.episode),
doc.synonyms
and _("Also known as: {}").format(", ".join(doc.synonyms)),
]
if s
]
),
url=doc.mal_id
and f"https://myanimelist.net/anime/{doc.mal_id}"
or f"https://anilist.co/anime/{doc.anilist_id}",
color=await ctx.embed_color(),
)
e.set_thumbnail(url=doc.thumbnail)
e.set_footer(
text=_("Via WAIT (trace.moe) • Page {}/{}").format(
page, len(search.docs)
),
icon_url="https://trace.moe/favicon128.png",
)
embeds.append(e)
await menu(ctx, embeds, DEFAULT_CONTROLS)

@tracemoe.command(name="stats")
@commands.is_owner()
async def tracemoe_stats(self, ctx):
"""See how many requests are left and time until reset"""
stats = await TraceMoe.me(ctx)
await ctx.send(
_(
"Remaining requests (ratelimit): {}/{}\n"
"Remaining requests (quota): {}/{}\n"
"Ratelimit reset in {}/{}\n"
"Quota reset in {}/{}\n"
).format(
stats.limit,
stats.user_limit,
stats.quota,
stats.user_quota,
stats.limit_ttl,
stats.user_limit_ttl,
stats.quota_ttl,
stats.user_quota_ttl,
)
)
17 changes: 15 additions & 2 deletions reverseimagesearch/tracemoe.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
from base64 import b64encode
from collections import namedtuple
from io import BytesIO
from urllib.parse import quote

from PIL import Image
from aiohttp import ClientResponseError
from redbot.core.i18n import Translator

Expand Down Expand Up @@ -35,6 +37,12 @@ def __init__(self, data: dict):
f"&file={quote(self.filename)}&t={self.time}&token={tokenthumb}"
)

@property
def time_str(self):
hours, minutes = divmod(self.time, 3600)
minutes, seconds = divmod(minutes, 60)
return "{:02}:{:02}:{:02}".format(int(hours), int(minutes), int(seconds))


class TraceMoe:
def __init__(self, data: dict):
Expand All @@ -59,14 +67,19 @@ async def from_image(cls, ctx, image_url):
async with ctx.cog.session.get(
image_url, raise_for_status=True
) as resp:
image = b64encode(await resp.read())
image = BytesIO(await resp.read())
image = Image.open(image)
image = image.convert("RGB")
image.thumbnail((2048, 2048))
image_file = BytesIO()
image.save(image_file, "JPEG")
except ClientResponseError as e:
raise ValueError(_("Unable to get image: {}").format(e.message))
try:
async with ctx.cog.session.post(
f"{BASE_API_URL}/search",
params={"token": apikeys["tracemoe"]},
json={"image": str(image)},
json={"image": b64encode(image_file.getvalue()).decode()},
raise_for_status=True,
) as data:
return cls(await data.json())
Expand Down

0 comments on commit c9ee006

Please sign in to comment.