Update app.py

app.py

@@ -1,23 +1,27 @@
 import gradio as gr
-from datasets import load_dataset
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 from PIL import Image
 
-# Load dataset
-dataset = load_dataset("Pisethan/sangapac_ocr")
+# Load model and processor
+processor = TrOCRProcessor.from_pretrained("Pisethan/sangapac_ocr")
+model = VisionEncoderDecoderModel.from_pretrained("Pisethan/sangapac_ocr")
 
-# Function to display image and text
-def display_data(index):
-    row = dataset["train"][index]
-    image = Image.open(row["image_path"])
-    text = row["text"]
-    return image, text
+# Define the prediction function
+def recognize_text(image):
+    image = Image.open(image).convert("RGB")
+    inputs = processor(images=image, return_tensors="pt").pixel_values
+    outputs = model.generate(inputs)
+    predicted_text = processor.tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return predicted_text
 
-# Create Gradio Interface
-with gr.Blocks() as demo:
-    gr.Markdown("### OCR Dataset Viewer")
-    index_input = gr.Slider(0, len(dataset["train"]) - 1, step=1, value=0, label="Select Index")
-    output_image = gr.Image(label="Image")
-    output_text = gr.Textbox(label="OCR Text")
-    index_input.change(display_data, inputs=index_input, outputs=[output_image, output_text])
+# Gradio Interface
+interface = gr.Interface(
+    fn=recognize_text,
+    inputs=gr.inputs.Image(type="file", label="Upload an Image"),
+    outputs=gr.outputs.Textbox(label="Recognized Text"),
+    title="OCR Model Demo",
+    description="Upload an image containing text to test the OCR model.",
+)
 
-demo.launch()
+# Launch the app
+interface.launch()