浏览代码

Added stats and ocr

TheHamkerCat 5 年之前
父节点
当前提交
6e846d5ffa
共有 5 个文件被更改,包括 122 次插入 和 3 次删除
  1. 30 0
      wbb/modules/arq.py
  2. 2 1
      wbb/modules/misc.py
  3. 79 0
      wbb/modules/ocr.py
  4. 1 1
      wbb/modules/sudoers.py
  5. 10 1
      wbb/utils/functions.py

+ 30 - 0
wbb/modules/arq.py

@@ -0,0 +1,30 @@
+from wbb import app, arq, ARQ_API_BASE_URL
+from pyrogram import filters
+
+__MODULE__ = "ARQ"
+__HELP__ = "/arq - Stats Of ARQ API."
+
+
+@app.on_message(filters.command("arq"))
+async def arq_stats(_, message):
+    data = await arq.stats()
+    uptime = data.uptime
+    requests = data.requests
+    cpu = data.cpu
+    server_mem = data.memory.server
+    api_mem = data.memory.api
+    disk = data.disk
+    platform = data.platform
+    python_version = data.python
+    statistics = f"""
+**Uptime:** `{uptime}`
+**Requests:** `{requests}`
+**CPU:** `{cpu}`
+**Memory:**
+    **Total Used:** `{server_mem}`
+    **Used By API:** `{api_mem}`
+**Disk:** `{disk}`
+**Platform:** `{platform}`
+**Python:** `{python_version}`
+"""
+    await message.reply_text(statistics)

+ 2 - 1
wbb/modules/misc.py

@@ -44,7 +44,8 @@ __HELP__ = '''/commit - Generate Funny Commit Messages
 /cheat [Language] [Query] - Get Programming Related Help
 /weather [City] - To Get Weather Info
 /tr [en] - Translate A Message
-/json [URL] - Get Response From An API or Something.
+/json [URL] - Get JSON Response From An API or Something.
+/ocr - Perform OCR on an image.
 #RTFM - Check it lol'''
 
 

+ 79 - 0
wbb/modules/ocr.py

@@ -0,0 +1,79 @@
+"""
+MIT License
+
+Copyright (c) 2021 TheHamkerCat
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+from wbb import app, arq
+from wbb.core.decorators.errors import capture_err
+from wbb.utils.functions import transfer_sh
+from pyrogram import filters
+from random import randint
+import os
+
+@app.on_message(filters.command("ocr"))
+@capture_err
+async def image_ocr(_, message):
+    if not message.reply_to_message:
+        await message.reply_text("Reply to an image or document to perform OCR on it.")
+        return
+    reply = message.reply_to_message
+    if not reply.document and not reply.photo and not reply.sticker and not reply.animation and not reply.video:
+        await message.reply_text("Reply to an image/document/sticker/animation to perform OCR on it.")
+        return
+    m = await message.reply_text("Scanning")
+    if reply.document:
+        if int(reply.document.file_size) > 3145728:
+            await m.edit("File too large")
+            return
+        mime_type = reply.document.mime_type
+        if mime_type != "image/png" and mime_type != "image/jpeg":
+            await m.edit("Document Mimetype Invalid")
+            return
+        file_id = reply.document.file_id
+    if reply.sticker:
+        if reply.sticker.is_animated:
+            if not reply.sticker.thumbs:
+                await m.edit("Sticker Has No Thumb")
+                return
+            file_id = reply.sticker.thumbs[0].file_id
+        else:
+            file_id = reply.sticker.file_id
+
+    if reply.photo:
+        file_id = reply.photo.file_id
+
+    if reply.animation:
+        if not reply.animation.thumbs:
+            await m.edit("Gif Has No Thumbnail, so it cannot be scanned for OCR")
+            return
+        file_id = reply.animation.thumbs[0].file_id
+
+    if reply.video:
+        if not reply.video.thumbs:
+            return
+        file_id = reply.video.thumbs[0].file_id
+
+    filename = f"{randint(1000, 10000)}.jpg"
+    file = await app.download_media(file_id)
+    url = await transfer_sh(file)
+    os.remove(file)
+    data = await arq.ocr(url)
+    await m.edit(data.ocr)

+ 1 - 1
wbb/modules/sudoers.py

@@ -54,7 +54,7 @@ __HELP__ = '''
 
 async def bot_sys_stats():
     bot_uptime = int(time.time() - bot_start_time)
-    cpu = psutil.cpu_percent(interval=0.5)
+    cpu = psutil.cpu_percent()
     mem = psutil.virtual_memory().percent
     disk = psutil.disk_usage("/").percent
     process = psutil.Process(os.getpid())

+ 10 - 1
wbb/utils/functions.py

@@ -27,7 +27,7 @@ from wbb.utils import aiodownloader
 from carbonnow import Carbon
 import speedtest
 import aiohttp
-
+import aiofiles
 
 """
 Just import 'downloader' anywhere and do downloader.download() to
@@ -119,3 +119,12 @@ async def make_carbon(code):
     carbon = Carbon(code=code)
     image = await carbon.save(str(randint(1000, 10000)))
     return image
+
+
+async def transfer_sh(file):
+    async with aiofiles.open(file, 'rb') as f:
+        params = {file: await f.read()}
+    async with aiohttp.ClientSession() as session:
+        async with session.post('https://transfer.sh/', data=params) as resp:
+            download_link = str(await resp.text()).strip()
+    return download_link