File size: 2,597 Bytes
a7b8c18 |
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 86 87 88 89 90 91 92 93 |
from typing import List, Tuple, Union
import cv2
import numpy as np
from PIL import Image
class TableVisualizer:
"""
Utility class for visualizing detected tables and OCR results.
"""
@staticmethod
def draw_boxes(
image: Union[np.ndarray, Image.Image],
boxes: List[List[int]],
color: Tuple[int, int, int] = (0, 255, 0),
thickness: int = 2
) -> Image.Image:
"""
Draw bounding boxes on an image.
Args:
image: Input image
boxes: List of bounding box coordinates [x1, y1, x2, y2]
color: RGB color for the boxes
thickness: Line thickness
Returns:
Image with drawn bounding boxes
"""
if isinstance(image, Image.Image):
image = np.array(image)
if len(image.shape) == 2:
image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
elif image.shape[2] == 4:
image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
image_copy = image.copy()
for box in boxes:
cv2.rectangle(
image_copy,
(box[0], box[1]),
(box[2], box[3]),
color,
thickness
)
return Image.fromarray(image_copy)
@staticmethod
def draw_text_boxes(
image: Union[np.ndarray, Image.Image],
text_data: List[Tuple[str, List[int]]],
color: Tuple[int, int, int] = (255, 0, 0),
thickness: int = 1
) -> Image.Image:
"""
Draw text boxes with labels on an image.
Args:
image: Input image
text_data: List of (text, bbox) tuples
color: RGB color for the boxes
thickness: Line thickness
Returns:
Image with drawn text boxes
"""
if isinstance(image, Image.Image):
image = np.array(image)
image_copy = image.copy()
for text, bbox in text_data:
cv2.rectangle(
image_copy,
(bbox[0], bbox[1]),
(bbox[2], bbox[3]),
color,
thickness
)
cv2.putText(
image_copy,
text[:20],
(bbox[0], bbox[1] - 5),
cv2.FONT_HERSHEY_SIMPLEX,
0.5,
color,
thickness
)
return Image.fromarray(image_copy) |