Update app.py

app.py

@@ -1,61 +1,46 @@
 import gradio as gr
 import numpy as np
+from PIL import Image
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 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')
+# Initialize the 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}")
-
     if isinstance(sketch, dict) and 'composite' in sketch:
-        image_data = np.array(sketch['composite'], dtype=np.uint8)
-        logging.debug(f"Image data type: {type(image_data)}")
-        logging.debug(f"Image data shape: {image_data.shape}")
-
-        # Ensure the image is in the correct format
-        image = Image.fromarray(image_data, 'RGBA').convert('RGB')
-
-        temp_file_path = os.path.join(os.getcwd(), "output.png")
-        image.save(temp_file_path)
-        logging.debug(f"Image saved to: {temp_file_path}")
-
-        return temp_file_path
-    else:
-        error_message = f"Unexpected sketch data format: {type(sketch)}"
-        logging.error(error_message)
-        return error_message
+        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
 
 def recognize_text(image_path):
-    try:
-        # 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]
-        logging.debug(f"Recognized text: {generated_text}")
-        return generated_text
-    except Exception as e:
-        logging.error(f"Error in recognizing text: {e}")
-        return "Error in recognizing text"
-
+    # Open the image
+    image = Image.open(image_path)
+    # Convert image to RGB
+    image = image.convert("RGB")
+    
+    # Process the image
+    pixel_values = processor(images=image, return_tensors="pt").pixel_values
+    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")
     recognized_text = gr.Textbox(label="Recognized Text")
-    submit_btn = gr.Button("Submit")
-
-    submit_btn.click(fn=display_sketch, inputs=sketchpad, outputs=output_image)
-    submit_btn.click(fn=recognize_text, inputs=output_image, outputs=recognized_text)
+    
+    sketchpad.submit(display_sketch, inputs=sketchpad, outputs=output_image).then(
+        recognize_text, inputs=output_image, outputs=recognized_text
+    )
 
+# Launch the demo
 demo.launch()