Update app.py

app.py

@@ -1,11 +1,7 @@
 import gradio as gr
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 import torch
-import subprocess
-import json
-from PIL import Image, ImageDraw
-import os
-import tempfile
+import spaces
 
 # Dictionary of model names and their corresponding HuggingFace model IDs
 MODEL_OPTIONS = {
@@ -36,78 +32,93 @@ def load_model(model_name):
         current_model = VisionEncoderDecoderModel.from_pretrained(model_id)
         current_model_name = model_name
         
-        # Move model to GPU if available
-        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        current_model = current_model.to(device)
+        # Move model to GPU
+        current_model = current_model.to('cuda')
     
     return current_processor, current_model
 
+@spaces.GPU
 def process_image(image, model_name):
-    # Save the uploaded image to a temporary file
-    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_img:
-        image.save(temp_img, format="JPEG")
-        temp_img_path = temp_img.name
-
-    # Run Kraken for line detection
-    lines_json_path = "lines.json"
-    kraken_command = f"kraken -i {temp_img_path} {lines_json_path} binarize segment -bl"
-    subprocess.run(kraken_command, shell=True, check=True)
-
-    # Load the lines from the JSON file
-    with open(lines_json_path, 'r') as f:
-        lines_data = json.load(f)
-
     processor, model = load_model(model_name)
-
-    # Process each line
-    transcriptions = []
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    for line in lines_data['lines']:
-        # Extract line coordinates
-        x1, y1 = line['baseline'][0]
-        x2, y2 = line['baseline'][-1]
-        
-        # Crop the line from the original image
-        line_image = image.crop((x1, y1, x2, y2))
-
-        # Prepare image for TrOCR
-        pixel_values = processor(line_image, return_tensors="pt").pixel_values
-        pixel_values = pixel_values.to(device)
-
-        # Generate (no beam search)
-        with torch.no_grad():
-            generated_ids = model.generate(pixel_values)
-
-        # Decode
-        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        transcriptions.append(generated_text)
-
-    # Clean up temporary files
-    os.unlink(temp_img_path)
-    os.unlink(lines_json_path)
-
-    # Create an image with bounding boxes
-    draw = ImageDraw.Draw(image)
-    for line in lines_data['lines']:
-        coords = line['baseline']
-        draw.line(coords, fill="red", width=2)
-
-    return image, "\n".join(transcriptions)
+    
+    # Prepare image
+    pixel_values = processor(image, return_tensors="pt").pixel_values
+    
+    # Move input to GPU
+    pixel_values = pixel_values.to('cuda')
+    
+    # Generate (no beam search)
+    with torch.no_grad():
+        generated_ids = model.generate(pixel_values)
+    
+    # Decode
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    return generated_text
+
+# Base URL for the images
+base_url = "https://huggingface.co/medieval-data/trocr-medieval-base/resolve/main/images/"
+
+# List of example images and their corresponding models
+examples = [
+    [f"{base_url}caroline-1.png", "Medieval Latin Caroline"],
+    [f"{base_url}caroline-2.png", "Medieval Latin Caroline"],
+    [f"{base_url}cursiva-1.png", "Medieval Cursiva"],
+    [f"{base_url}cursiva-2.png", "Medieval Cursiva"],
+    [f"{base_url}cursiva-3.png", "Medieval Cursiva"],
+    [f"{base_url}humanistica-1.png", "Medieval Humanistica"],
+    [f"{base_url}humanistica-2.png", "Medieval Humanistica"],
+    [f"{base_url}humanistica-3.png", "Medieval Humanistica"],
+    [f"{base_url}hybrida-1.png", "Medieval Castilian Hybrida"],
+    [f"{base_url}hybrida-2.png", "Medieval Castilian Hybrida"],
+    [f"{base_url}hybrida-3.png", "Medieval Castilian Hybrida"],
+    [f"{base_url}praegothica-1.png", "Medieval Praegothica"],
+    [f"{base_url}praegothica-2.png", "Medieval Praegothica"],
+    [f"{base_url}praegothica-3.png", "Medieval Praegothica"],
+    [f"{base_url}print-1.png", "Medieval Print"],
+    [f"{base_url}print-2.png", "Medieval Print"],
+    [f"{base_url}print-3.png", "Medieval Print"],
+    [f"{base_url}semihybrida-1.png", "Medieval Semihybrida"],
+    [f"{base_url}semihybrida-2.png", "Medieval Semihybrida"],
+    [f"{base_url}semihybrida-3.png", "Medieval Semihybrida"],
+    [f"{base_url}semitextualis-1.png", "Medieval Semitextualis"],
+    [f"{base_url}semitextualis-2.png", "Medieval Semitextualis"],
+    [f"{base_url}semitextualis-3.png", "Medieval Semitextualis"],
+    [f"{base_url}textualis-1.png", "Medieval Textualis"],
+    [f"{base_url}textualis-2.png", "Medieval Textualis"],
+    [f"{base_url}textualis-3.png", "Medieval Textualis"],
+]
+
+# Custom CSS to make the image wider
+custom_css = """
+#image_upload {
+    max-width: 100% !important;
+    width: 100% !important;
+    height: auto !important;
+}
+#image_upload > div:first-child {
+    width: 100% !important;
+}
+#image_upload img {
+    max-width: 100% !important;
+    width: 100% !important;
+    height: auto !important;
+}
+"""
 
 # Gradio interface
-with gr.Blocks() as iface:
-    gr.Markdown("# Medieval Document Transcription")
-    gr.Markdown("Upload an image of a medieval document and select a model to transcribe it. The tool will detect lines and transcribe each line separately.")
+with gr.Blocks(css=custom_css) as iface:
+    gr.Markdown("# Medieval TrOCR Model Switcher")
+    gr.Markdown("Upload an image of medieval text and select a model to transcribe it. Note: This tool is designed to work on a single line of text at a time for optimal results.")
     
     with gr.Row():
-        input_image = gr.Image(type="pil", label="Input Image")
+        input_image = gr.Image(type="pil", label="Input Image", elem_id="image_upload")
         model_dropdown = gr.Dropdown(choices=list(MODEL_OPTIONS.keys()), label="Select Model", value="Medieval Base")
     
-    with gr.Row():
-        output_image = gr.Image(type="pil", label="Detected Lines")
-        transcription_output = gr.Textbox(label="Transcription", lines=10)
+    transcription_output = gr.Textbox(label="Transcription")
     
     submit_button = gr.Button("Transcribe")
-    submit_button.click(fn=process_image, inputs=[input_image, model_dropdown], outputs=[output_image, transcription_output])
+    submit_button.click(fn=process_image, inputs=[input_image, model_dropdown], outputs=transcription_output)
+    
+    gr.Examples(examples, inputs=[input_image, model_dropdown], outputs=transcription_output)
 
 iface.launch()