Add OCR tech

This commit is contained in:
Nexus 2023-05-05 10:35:17 +01:00
parent 7d9ed0f6ad
commit 34e63104f9
Signed by: nex
GPG key ID: 0FA334385D0B689F
2 changed files with 50 additions and 0 deletions

View file

@ -21,6 +21,8 @@ from typing import Literal
from typing import Tuple, Optional, Dict from typing import Tuple, Optional, Dict
from pathlib import Path from pathlib import Path
from urllib.parse import urlparse from urllib.parse import urlparse
from PIL import Image
import pytesseract
import aiohttp import aiohttp
import discord import discord
@ -1521,6 +1523,52 @@ class OtherCog(commands.Cog):
else: else:
return await ctx.respond(result, view=GenerateNewView()) return await ctx.respond(result, view=GenerateNewView())
@commands.slash_command()
@commands.cooldown(1, 30, commands.BucketType.user)
@commands.max_concurrency(1, commands.BucketType.user)
async def ocr(
self,
ctx: discord.ApplicationContext,
attachment: discord.Option(
discord.SlashCommandOptionType.attachment,
description="Image to perform OCR on",
)
):
"""OCRs an image"""
await ctx.defer()
attachment: discord.Attachment
data = await attachment.read()
file = io.BytesIO(data)
file.seek(0)
img = await self.bot.loop.run_in_executor(None, Image.open, file)
try:
text = await self.bot.loop.run_in_executor(None, pytesseract.image_to_string, img)
except pytesseract.TesseractError as e:
return await ctx.respond(f"Failed to perform OCR: `{e}`")
if len(text) > ctx.guild.filesize_limit - 100:
try:
response = await self.http.put(
"https://api.mystb.in/paste",
json={
"files": [
{
"filename": "ocr.txt",
"content": text
}
],
}
)
response.raise_for_status()
except httpx.HTTPError:
return await ctx.respond("OCR content too large to post.")
else:
data = response.json()
return await ctx.respond("https://mystb.in/%s" % data["id"])
out_file = io.BytesIO(text.encode("utf-8", "replace"))
return await ctx.respond(file=discord.File(out_file, filename="ocr.txt"))
def setup(bot): def setup(bot):
bot.add_cog(OtherCog(bot)) bot.add_cog(OtherCog(bot))

View file

@ -17,3 +17,5 @@ pyttsx3>=2.90
yt-dlp yt-dlp
beautifulsoup4>=4.12.0 beautifulsoup4>=4.12.0
lxml>=4.9.2 lxml>=4.9.2
pytesseract>=0.3.10
pillow>=9.5.0