Upload initial file structure

app.py

@@ -0,0 +1,38 @@
+# Importing the requirements
+import gradio as gr
+from src.task import ocr_task
+
+
+# Image input for the interface
+image = gr.Image(type="pil", label="Image")
+
+# Output for the interface (image and text)
+ocr_text_output = gr.Textbox(label="OCR Text")
+ocr_image_output = gr.Image(type="pil", label="Output Image")
+
+# Examples for the interface (image paths)
+examples = [
+    ["images/ocr_image_1jpg"],
+    ["images/ocr_image_2.jpg"],
+    ["images/ocr_image_3.jpg"],
+]
+
+# Title, description, and article for the interface
+title = "OCR Text Extraction and Visualization"
+description = "Gradio Demo for the Florence-2-large Vision Language Model. This application performs Optical Character Recognition (OCR) on images and provides both extracted text and visualized bounding boxes around detected text regions. To use it, simply upload your image and click 'Submit'. The application will return the detected text and an image with bounding boxes drawn around the detected text regions. This is useful for various OCR-related tasks including document digitization, text extraction, and visual verification of detected text regions."
+article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2311.06242' target='_blank'>Florence-2: Advancing a Unified Representation for a Variety of Vision Tasks</a> | <a href='https://huggingface.co/microsoft/Florence-2-large-ft' target='_blank'>Model Page</a></p>"
+
+
+# Launch the interface
+interface = gr.Interface(
+    fn=ocr_task,
+    inputs=[image],
+    outputs=[ocr_image_output, ocr_text_output],
+    examples=examples,
+    title=title,
+    description=description,
+    article=article,
+    theme="soft",
+    allow_flagging="never",
+)
+interface.launch(debug=False)

requirements.txt

@@ -0,0 +1,5 @@
+numpy==2.0.0
+Pillow==10.4.0
+gradio==4.38.1
+transformers==4.42.4
+timm==1.0.7

src/model.py

@@ -0,0 +1,52 @@
+# Importing necessary libraries
+import spaces
+from transformers import AutoProcessor, AutoModelForCausalLM
+
+
+# Load model and processor from Hugging Face
+model_id = "microsoft/Florence-2-large-ft"
+model = (
+    AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).eval().cuda()
+)
+processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+
+
+@spaces.GPU(duration=120)
+def run_example(task_prompt, image, text_input=None):
+    """
+    Runs an example using the given task prompt and image.
+
+    Args:
+        task_prompt (str): The task prompt for the example.
+        image (PIL.Image.Image): The image to be processed.
+        text_input (str, optional): Additional text input to be appended to the task prompt. Defaults to None.
+
+    Returns:
+        str: The parsed answer generated by the model.
+    """
+
+    # If there is no text input, use the task prompt as the prompt
+    if text_input is None:
+        prompt = task_prompt
+    else:
+        prompt = task_prompt + text_input
+
+    # Process the image and text input
+    inputs = processor(text=prompt, images=image, return_tensors="pt")
+
+    # Generate the answer using the model
+    generated_ids = model.generate(
+        input_ids=inputs["input_ids"].cuda(),
+        pixel_values=inputs["pixel_values"].cuda(),
+        max_new_tokens=1024,
+        early_stopping=False,
+        do_sample=False,
+        num_beams=3,
+    )
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+    parsed_answer = processor.post_process_generation(
+        generated_text, task=task_prompt, image_size=(image.width, image.height)
+    )
+
+    # Return the parsed answer
+    return parsed_answer

src/task.py

@@ -0,0 +1,33 @@
+import copy
+from src.utils import run_example, clean_text, draw_ocr_bboxes
+
+
+def ocr_task(image):
+    """
+    Perform OCR (Optical Character Recognition) on the given image.
+
+    Args:
+        image (PIL.Image.Image): The input image to perform OCR on.
+
+    Returns:
+        tuple: A tuple containing the output image with OCR bounding boxes drawn and the cleaned OCR text.
+
+    """
+
+    # Task prompts
+    ocr_prompt = "<OCR>"
+    ocr_with_region_prompt = "<OCR_WITH_REGION>"
+
+    # Get OCR text
+    ocr_results = run_example(ocr_prompt, image)
+    cleaned_text = clean_text(ocr_results["<OCR>"])
+
+    # Get OCR with region
+    ocr_with_region_results = run_example(ocr_with_region_prompt, image)
+    output_image = copy.deepcopy(image)
+    output_image = draw_ocr_bboxes(
+        output_image, ocr_with_region_results["<OCR_WITH_REGION>"]
+    )
+
+    # Return the output image and cleaned OCR text
+    return output_image, cleaned_text

src/utils.py

@@ -0,0 +1,172 @@
+from PIL import ImageDraw
+import numpy as np
+import re
+
+
+# Use a color map for bounding boxes
+colormap = [
+    "#0000FF",
+    "#FFA500",
+    "#008000",
+    "#800080",
+    "#A52A2A",
+    "#FFC0CB",
+    "#808080",
+    "#808000",
+    "#00FFFF",
+    "#FF0000",
+    "#00FF00",
+    "#4B0082",
+    "#4B0082",
+    "#EE82EE",
+    "#00FFFF",
+    "#FF00FF",
+    "#FF7F50",
+    "#FFD700",
+    "#87CEEB",
+]
+
+
+# Text cleaning function
+def clean_text(text):
+    """
+    Cleans the given text by removing unwanted tokens, extra spaces,
+    and ensures proper spacing between words and after periods.
+
+    Args:
+        text (str): The input text to be cleaned.
+
+    Returns:
+        str: The cleaned and properly formatted text.
+    """
+
+    # Remove unwanted tokens
+    text = text.replace("<pad>", "").replace("</s>", "").strip()
+
+    # Split the text into lines and clean each line
+    lines = text.split("\n")
+    cleaned_lines = [line.strip() for line in lines if line.strip()]
+
+    # Join the cleaned lines into a single string with a space between each line
+    cleaned_text = " ".join(cleaned_lines)
+
+    # Ensure proper spacing between words and after periods using regex
+    cleaned_text = re.sub(
+        r"\s+", " ", cleaned_text
+    )  # Replace multiple spaces with a single space
+    cleaned_text = re.sub(
+        r"(?<=[.])(?=[^\s])", r" ", cleaned_text
+    )  # Add space after a period if not followed by a space
+
+    # Return the cleaned text
+    return cleaned_text
+
+
+# Convert hex color to RGBA with the given alpha
+def hex_to_rgba(hex_color, alpha):
+    """
+    Convert a hexadecimal color code to RGBA format.
+
+    Args:
+        hex_color (str): The hexadecimal color code (e.g., "#FF0000").
+        alpha (int): The alpha value for the RGBA color (0-255).
+
+    Returns:
+        tuple: A tuple representing the RGBA color values (red, green, blue, alpha).
+    """
+    hex_color = hex_color.lstrip("#")
+    r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
+    return (r, g, b, alpha)
+
+
+# Draw OCR bounding boxes with enhanced visual elements
+def draw_ocr_bboxes(image, prediction):
+    """
+    Draw bounding boxes with enhanced visual elements on the given image based on the OCR prediction.
+
+    Args:
+        image (PIL.Image.Image): The input image on which the bounding boxes will be drawn.
+        prediction (dict): The OCR prediction containing 'quad_boxes' and 'labels'.
+
+    Returns:
+        PIL.Image.Image: The image with the bounding boxes drawn.
+    """
+
+    # Create a drawing object for the image with RGBA mode
+    draw = ImageDraw.Draw(image, "RGBA")
+
+    # Extract bounding boxes and labels from the prediction
+    bboxes, labels = prediction["quad_boxes"], prediction["labels"]
+
+    for i, (box, label) in enumerate(zip(bboxes, labels)):
+        # Select color for the bounding box and label
+        color = colormap[i % len(colormap)]
+        new_box = (np.array(box)).tolist()
+
+        # Define the outline width and corner radius for the bounding box
+        box_outline_width = 3
+        corner_radius = 10
+
+        # Draw rounded corners for the bounding box
+        for j in range(4):
+            start_x, start_y = new_box[j * 2], new_box[j * 2 + 1]
+            end_x, end_y = new_box[(j * 2 + 2) % 8], new_box[(j * 2 + 3) % 8]
+
+            # Draw the arcs for the rounded corners
+            draw.arc(
+                [
+                    (start_x - corner_radius, start_y - corner_radius),
+                    (start_x + corner_radius, start_y + corner_radius),
+                ],
+                90 + j * 90,
+                180 + j * 90,
+                fill=color,
+                width=box_outline_width,
+            )
+            draw.arc(
+                [
+                    (end_x - corner_radius, end_y - corner_radius),
+                    (end_x + corner_radius, end_y + corner_radius),
+                ],
+                j * 90,
+                90 + j * 90,
+                fill=color,
+                width=box_outline_width,
+            )
+
+            # Draw the lines connecting the arcs
+            if j in [0, 1, 2]:
+                draw.line(
+                    [
+                        (start_x + corner_radius if j != 1 else start_x, start_y),
+                        (end_x - corner_radius if j != 1 else end_x, end_y),
+                    ],
+                    fill=color,
+                    width=box_outline_width,
+                )
+            else:
+                draw.line(
+                    [
+                        (start_x, start_y + corner_radius),
+                        (end_x, end_y - corner_radius),
+                    ],
+                    fill=color,
+                    width=box_outline_width,
+                )
+
+        # Calculate the position for the text label
+        text_x, text_y = min(new_box[0::2]), min(new_box[1::2]) - 20
+        text_w, text_h = draw.textsize(label)
+        rgba_color = hex_to_rgba(color, 200)  # Semi-transparent background for text
+
+        # Draw the background rectangle for the text
+        draw.rectangle(
+            [text_x, text_y, text_x + text_w + 10, text_y + text_h + 10],
+            fill=rgba_color,
+        )
+
+        # Draw the text label
+        draw.text((text_x + 5, text_y + 5), label, fill=(0, 0, 0, 255))
+
+    # Return the image with the OCR boxes drawn
+    return image