File size: 2,917 Bytes
c1cef3f b6654d0 c1cef3f b6654d0 19c5277 b6654d0 19c5277 b6654d0 19c5277 b6654d0 c1cef3f b6654d0 c1cef3f b6654d0 c1cef3f b6654d0 c1cef3f b6654d0 c1cef3f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
# Necessary imports
import re
import supervision as sv
from PIL import Image
# 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 punctuation marks.
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 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 punctuation if not followed by a space
cleaned_text = re.sub(
r"(?<=[a-z])(?=[A-Z])", r" ", cleaned_text
) # Add space between joined words where a lowercase letter is followed by an uppercase letter
cleaned_text = re.sub(
r"(\w)([A-Z][a-z])", r"\1 \2", cleaned_text
) # Add space between camel case words
# Return the cleaned text
return cleaned_text
# Draw OCR bounding boxes with enhanced visual elements
def draw_ocr_bboxes(image: Image, detections: sv.Detections) -> Image:
"""
Draws bounding boxes and labels on the input image based on the OCR detections.
Args:
image (PIL.Image): The input image on which to draw the bounding boxes and labels.
detections (sv.Detections): The OCR detections containing the bounding box coordinates and labels.
Returns:
PIL.Image: The annotated image with bounding boxes and labels.
"""
# Copy the input image to avoid modifying the original image
annotated_image = image.copy()
# Calculate the optimal line thickness and text scale based on the image resolution
thickness = sv.calculate_optimal_line_thickness(resolution_wh=image.size)
text_scale = sv.calculate_optimal_text_scale(resolution_wh=image.size)
# Initialize the bounding box and label annotators
bounding_box_annotator = sv.BoundingBoxAnnotator(
color_lookup=sv.ColorLookup.INDEX, thickness=thickness
)
label_annotator = sv.LabelAnnotator(
color_lookup=sv.ColorLookup.INDEX,
text_scale=text_scale,
text_thickness=thickness,
)
# Annotate the image with bounding boxes and labels
annotated_image = bounding_box_annotator.annotate(annotated_image, detections)
annotated_image = label_annotator.annotate(annotated_image, detections)
# Return the annotated image
return annotated_image
|