# !/usr/bin/env python
# -*-coding:utf-8 -*-

"""
# File       : myself_train_model.py
# Time       ：2023/8/21 9:25
# Author     ：小吕同学
"""

from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
import gradio as gr
import os


class xiaolv_ocr_model():

    def __init__(self):
        model_small = r"./output_small"
        model_big = r"./output_big"
        self.ocr_recognition_small = pipeline(Tasks.ocr_recognition, model=model_small)
        self.ocr_recognition1_big = pipeline(Tasks.ocr_recognition, model=model_big)


    def run(self,pict_path,moshi = "small", context=[]):
        pict_path = pict_path.name
        context = [pict_path]

        if moshi == "small":
            result = self.ocr_recognition_small(pict_path)
        else:
            result = self.ocr_recognition1_big(pict_path)

        context += [str(result['text'][0])]
        responses = [(u, b) for u, b in zip(context[::2], context[1::2])]
        print(f"识别的结果为：{result}")
        os.remove(pict_path)
        return responses,context




if __name__ == "__main__":
    pict_path = r"C:\Users\admin\Desktop\图片识别测试\企业微信截图_16895911221007.png"
    ocr_model = xiaolv_ocr_model()
    # ocr_model.run(pict_path)



    with gr.Blocks() as demo:
        gr.HTML("""<h1 align="center">常见验证码识别——小吕</h1>""")
        with gr.Tab("验证码模型"):
            gr.Markdown("""
            #### 专门识别常见验证码的模型，训练模型有2个：
            * 1.**small**:训练数据大小为700MB，约8.4万张验证码图片，训练轮次27轮，最终的精度将近100%；
            * 2.**big**:训练数据大小为11G，约135万个验证码图片，训练轮次1轮，最终的精度将近93.95%；
            #### 训练验证码图片包括：
            * 类型：1. 纯数字型；2. 数字+字母型；3.纯字母型（大小写）
            * 长度：4位、5位、6位
            """)
            with gr.Row():
                with gr.Column():
                    select_types = gr.Radio(label="模型类型选择", choices=["small", "big"], value="small")
                    img_input = gr.File(label='输入图像,目前支持图片格式(png\jpg等)，建议使用png格式。',file_types=[".jpg",".jpeg",".png"])
                with gr.Column():
                    chatbot = gr.Chatbot([])
            with gr.Row():
                btn_submit = gr.Button(value="一键识别")


            state = gr.State([])

            btn_submit.click(fn=ocr_model.run, inputs=[img_input, select_types], outputs=[chatbot, state])


            gr.Markdown("""
            #### 后续计划：
            1. 推出专门识别手机号码和邮箱的模型
            """)



    demo.launch(show_error=True,server_name='192.168.8.173')