From 34e63104f9c77933dd3782e9f176a8f21105c1b8 Mon Sep 17 00:00:00 2001
From: nex <me@nexy7574.co.uk>
Date: Fri, 5 May 2023 10:35:17 +0100
Subject: [PATCH] Add OCR tech

---
 cogs/other.py    | 48 ++++++++++++++++++++++++++++++++++++++++++++++++
 requirements.txt |  2 ++
 2 files changed, 50 insertions(+)

diff --git a/cogs/other.py b/cogs/other.py
index 6dc7402..a4352d4 100644
--- a/cogs/other.py
+++ b/cogs/other.py
@@ -21,6 +21,8 @@ from typing import Literal
 from typing import Tuple, Optional, Dict
 from pathlib import Path
 from urllib.parse import urlparse
+from PIL import Image
+import pytesseract
 
 import aiohttp
 import discord
@@ -1521,6 +1523,52 @@ class OtherCog(commands.Cog):
         else:
             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):
     bot.add_cog(OtherCog(bot))
diff --git a/requirements.txt b/requirements.txt
index afe40a7..a191ea4 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -17,3 +17,5 @@ pyttsx3>=2.90
 yt-dlp
 beautifulsoup4>=4.12.0
 lxml>=4.9.2
+pytesseract>=0.3.10
+pillow>=9.5.0