Update app.py

app.py

@@ -1,49 +1,48 @@
 import gradio as gr
 import numpy as np
+import matplotlib.pyplot as plt
+import tempfile
 from PIL import Image
+import torch
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 
-# Initialize the model and processor
+# Load model and processor
 processor = TrOCRProcessor.from_pretrained("microsoft/trocr-large-handwritten")
 model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-large-handwritten")
 
 def display_sketch(sketch):
-    if isinstance(sketch, dict) and 'composite' in sketch:
-        image_data = sketch['composite']
-        if isinstance(image_data, np.ndarray):
-            img = Image.fromarray(image_data.astype('uint8'), 'RGBA')
-            temp_file = "/home/user/app/output.png"
-            img.save(temp_file)
-            return temp_file
-    return None
+    image_data = sketch['composite']
+    plt.imshow(image_data)
+    plt.axis('off')
+
+    temp_file_path = "/mnt/data/output.png"
+    plt.savefig(temp_file_path, bbox_inches='tight', pad_inches=0)
+    plt.close()
+
+    return temp_file_path
 
 def recognize_text(image_path):
-    # Open the image
-    image = Image.open(image_path)
-    # Convert image to RGB
-    image = image.convert("RGB")
+    # Open image and convert to grayscale
+    image = Image.open(image_path).convert("L")
+    # Resize image to 256x256
+    image = image.resize((256, 256))
+    # Binarize image (convert to black and white)
+    image = image.point(lambda p: p > 128 and 255)
     
-    # Process the image
+    # Preprocess the image
     pixel_values = processor(images=image, return_tensors="pt").pixel_values
+    # Generate prediction
     generated_ids = model.generate(pixel_values)
     generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    
     return generated_text
 
-# Define the Gradio interface
 with gr.Blocks() as demo:
-    sketchpad = gr.Sketchpad(label="Draw Something")
-    output_image = gr.Image(label="Your Sketch")
+    sketchpad = gr.Sketchpad(label="Draw Something", brush_radius=10)
+    sketchpad_output = gr.Image(label="Your Sketch")
     recognized_text = gr.Textbox(label="Recognized Text")
     
-    def process_and_recognize(sketch):
-        image_path = display_sketch(sketch)
-        if image_path:
-            text = recognize_text(image_path)
-            return image_path, text
-        return None, ""
-
-    sketchpad.change(process_and_recognize, inputs=sketchpad, outputs=[output_image, recognized_text])
+    sketchpad.submit(display_sketch, inputs=sketchpad, outputs=sketchpad_output).then(
+        recognize_text, inputs=sketchpad_output, outputs=recognized_text
+    )
 
-# Launch the demo
 demo.launch()