Files
wxocr/main.py
2025-03-27 14:05:30 +08:00

98 lines
2.8 KiB
Python

import wcocr
import os
import uuid
import base64
from flask import Flask, request, jsonify, render_template, send_from_directory
app = Flask(__name__)
wcocr.init("./wx/opt/wechat/wxocr", "./wx/opt/wechat")
@app.route("/ocr", methods=["POST"])
def ocr():
try:
# Get base64 image from request
image_data = request.json.get("image")
if not image_data:
return jsonify({"error": "No image data provided"}), 400
# Extract image type from base64 data
image_type, base64_data = extract_image_type(image_data)
if not image_type:
return jsonify({"error": "Invalid base64 image data"}), 400
# Create temp directory if not exists
temp_dir = "temp"
if not os.path.exists(temp_dir):
os.makedirs(temp_dir)
# Generate unique filename and save image
filename = os.path.join(temp_dir, f"{str(uuid.uuid4())}.{image_type}")
try:
image_bytes = base64.b64decode(base64_data)
with open(filename, "wb") as f:
f.write(image_bytes)
# Process image with OCR
result = wcocr.ocr(filename)
return jsonify({"result": result})
finally:
# Clean up temp file
if os.path.exists(filename):
os.remove(filename)
except Exception as e:
return jsonify({"error": str(e)}), 500
# 创建静态文件夹
static_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "static")
if not os.path.exists(static_dir):
os.makedirs(static_dir)
def extract_image_type(base64_data):
# Check if the base64 data has the expected prefix
if base64_data.startswith("data:image/"):
# Extract the image type from the prefix
prefix_end = base64_data.find(";base64,")
if prefix_end != -1:
return (
base64_data[len("data:image/") : prefix_end],
base64_data.split(";base64,")[-1],
)
return "png", base64_data
@app.route("/")
def index():
return render_template("index.html")
# Handle unsupported methods for /ocr route
@app.route("/ocr", methods=["GET", "PUT", "DELETE", "PATCH"])
def unsupported_method():
return jsonify({"error": "Method not allowed"}), 405
# Handle non-existent paths
@app.errorhandler(404)
def not_found(e):
return jsonify({"error": "Resource not found"}), 404
if __name__ == "__main__":
# 确保templates目录存在
templates_dir = os.path.join(
os.path.dirname(os.path.abspath(__file__)), "templates"
)
if not os.path.exists(templates_dir):
os.makedirs(templates_dir)
# 确保temp目录存在
temp_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "temp")
if not os.path.exists(temp_dir):
os.makedirs(temp_dir)
app.run(host="0.0.0.0", port=5000, threaded=True)