Update app.py

app.py

@@ -4,9 +4,15 @@ import matplotlib.pyplot as plt
 import tempfile
 import os
 import logging
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+from PIL import Image
 
 logging.basicConfig(level=logging.DEBUG)
 
+# Initialize the TrOCR model and processor
+processor = TrOCRProcessor.from_pretrained('microsoft/trocr-large-handwritten')
+model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-large-handwritten')
+
 def display_sketch(sketch):
     logging.debug(f"Received sketch data: {sketch}")
     
@@ -29,10 +35,23 @@ def display_sketch(sketch):
         logging.error(error_message)
         return error_message
 
+def recognize_text(image_path):
+    # Load the image
+    image = Image.open(image_path).convert("RGB")
+    # Prepare the image for the model
+    pixel_values = processor(image, return_tensors="pt").pixel_values
+    # Generate the text
+    generated_ids = model.generate(pixel_values)
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    return generated_text
+
 with gr.Blocks() as demo:
     sketchpad = gr.Sketchpad(label="Draw Something")
     output_image = gr.Image(label="Your Sketch")
+    recognized_text = gr.Textbox(label="Recognized Text")
     submit_btn = gr.Button("Submit")
-    submit_btn.click(display_sketch, inputs=sketchpad, outputs=output_image)
+    
+    submit_btn.click(fn=display_sketch, inputs=sketchpad, outputs=output_image)
+    submit_btn.click(fn=recognize_text, inputs=output_image, outputs=recognized_text)
 
 demo.launch()