import gradio as gr
from transformers import pipeline

# 创建用于图像到文本的 pipeline
ocr_pipe = pipeline("image-to-text", model="microsoft/trocr-base-handwritten")

# 创建用于语法纠正的 pipeline
grammar_correction_pipe = pipeline("text2text-generation", model="flexudy/t5-base-multi-sentence-doctor")

def process_image_and_correct_grammar(input_image):
    # 使用 OCR 模型识别图像中的文本
    ocr_results = ocr_pipe(input_image)
    recognized_text = ocr_results[0]['generated_text']
    
    # 使用语法纠正模型纠正文本
    corrected_text_results = grammar_correction_pipe(recognized_text)
    corrected_text = corrected_text_results[0]['generated_text']
    
    # 返回原始文本和纠正后的文本
    return {"Original Text": recognized_text, "Corrected Text": corrected_text}

# 创建 Gradio 接口
iface = gr.Interface(
    fn=process_image_and_correct_grammar,  # 指定处理函数
    inputs=gr.inputs.Image(type='pil', label="Upload your image here or use your camera"),  # 设置输入为图片
    outputs=[gr.outputs.Textbox(label="Original Text"), gr.outputs.Textbox(label="Corrected Text")],  # 设置输出为两个文本框
    title="OCR and Grammar Correction",  # 接口标题
    description="Upload an image and the text in the image will be recognized and grammatically corrected."  # 接口描述
)

# 启动 Gradio 应用
iface.launch()