Files
ocr/main.py
2023-11-17 22:07:11 +08:00

89 lines
2.5 KiB
Python
Executable File

#!/usr/bin/env python3
import io
import requests
import oss2
import plyvel
from PIL import Image, ImageFile
# 创建或打开一个数据库
db = plyvel.DB('database', create_if_missing=True)
'''
# 写入一个键值对
db.put(b'key', b'value')
# 获取一个键的值
value = db.get(b'key')
# 删除一个键值对
db.delete(b'key')
# 批量写入
with db.write_batch() as wb:
for i in range(10000):
wb.put(b'key' + str(i).encode(), b'value' + str(i).encode())
# 迭代数据库中的所有键值对
for key, value in db:
print(key, value)
# 关闭数据库
db.close()
'''
# 读取 .env
from dotenv import dotenv_values
config = dotenv_values(".env")
ImageFile.LOAD_TRUNCATED_IMAGES = True
# 下载图片(使用OSS下载)
def download_image(url:str) -> Image.Image:
if url.startswith('http://image.gameuiux.cn/') or url.startswith('https://image.gameuiux.cn/'):
try:
url = url.replace('http://image.gameuiux.cn/', '').replace('https://image.gameuiux.cn/', '')
oss2.defaults.connection_pool_size = 100
oss_auth = oss2.Auth(config['OSS_ACCESS_KEY_ID'], config['OSS_ACCESS_KEY_SECRET'])
return Image.open(io.BytesIO(oss2.Bucket(oss_auth, f'http://{config["OSS_HOST"]}', config['OSS_BUCKET_NAME']).get_object(url).read()))
except Exception:
print('图片从OSS下载失败:', url)
return None
else:
try:
response = requests.get(url)
return Image.open(io.BytesIO(response.content))
except Exception:
print('图片从URL下载失败:', url)
return None
import pymysql
import pymysql.cursors
import cnocr
# 打开 mysql
ocr = cnocr.CnOcr(rec_model_name='ch_PP-OCRv3')
conn = pymysql.connect(host=config['MYSQL_HOST'], user=config['MYSQL_USER'], password=config['MYSQL_PASSWORD'], database=config['MYSQL_NAME'], cursorclass=pymysql.cursors.DictCursor)
cursor = conn.cursor()
cursor.execute("SELECT id, content FROM web_images LIMIT 5")
# 获取查询结果(跳过下载失败的)
for item in cursor.fetchall():
image = download_image(item['content'])
if image is None:
continue
# 将只包含那些非空非纯数字且长度大于1的'text'值
out = ocr.ocr(image)
texts = [x['text'] for x in out if x['text'] and not x['text'].isdigit() and len(x['text']) > 1]
print(item['id'], texts)
# 将结果存入 leveldb
# db.put(str(row['id']).encode(), ','.join(texts).encode())
# 关闭游标和连接
cursor.close()
conn.close()
# 关闭数据库
db.close()
print('Done')