from flask import Flask, request, jsonify
import torch
import onnx
import onnxruntime as rt
from torchvision import transforms as T
from PIL import Image
import requests
from io import BytesIO
from tokenizer_base import Tokenizer

app = Flask(__name__)

model_file = "captcha.onnx"
img_size = (32, 128)
charset = r"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
tokenizer_base = Tokenizer(charset)

def get_transform(img_size):
    transforms = [
        T.Resize(img_size, T.InterpolationMode.BICUBIC),
        T.ToTensor(),
        T.Normalize(0.5, 0.5),
    ]
    return T.Compose(transforms)

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

def initialize_model(model_file):
    transform = get_transform(img_size)
    onnx_model = onnx.load(model_file)
    onnx.checker.check_model(onnx_model)
    ort_session = rt.InferenceSession(model_file)
    return transform, ort_session

transform, ort_session = initialize_model(model_file=model_file)

def get_text(img_url):
    try:
        response = requests.get(img_url)
        response.raise_for_status()
        img_org = Image.open(BytesIO(response.content))
    except Exception as e:
        return {"error": f"无法加载图片：{str(e)}"}

    try:
        x = transform(img_org.convert('RGB')).unsqueeze(0)
        ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
        logits = ort_session.run(None, ort_inputs)[0]
        probs = torch.tensor(logits).softmax(-1)
        preds, probs = tokenizer_base.decode(probs)
        return {"text": preds[0]}
    except Exception as e:
        return {"error": f"模型推理失败：{str(e)}"}

@app.route('/api/v1/captcha', methods=['POST'])
def captcha_recognition():
    data = request.json
    if not data or 'image_url' not in data:
        return jsonify({"error": "请提供有效的 'image_url'"}), 400
    
    result = get_text(data['image_url'])
    return jsonify(result)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)