File size: 4,232 Bytes
e32cd03 |
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 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
import matplotlib.pyplot as plt
import numpy as np
import cv2
import easyocr
class OCR:
def __init__(self, image) -> None:
self.image = image
self.reader = easyocr.Reader(
lang_list=['en'],
gpu=False,
model_storage_directory='./EasyOCR/model/',
download_enabled=False,
user_network_directory='./EasyOCR/user_network/'
)
def detection(self):
img_arr = np.array(self.image, dtype=np.uint8)
response = self.reader.readtext(image=img_arr)
plot_inputs = []
for box, text, conf in response:
plot_inputs.append((text, np.array(box, dtype=np.float32)))
fig, ax = plt.subplots(nrows=1, ncols=1)
plot = self.drawAnnotations(
image=img_arr, predictions=plot_inputs, ax=ax)
return fig
def drawAnnotations(self, image, predictions, ax=None):
"""Draw text annotations onto image.
Args:
image: The image on which to draw
predictions: The predictions as provided by `pipeline.recognize`.
ax: A matplotlib axis on which to draw.
"""
if ax is None:
_, ax = plt.subplots()
ax.imshow(self.drawBoxes(image=image, boxes=predictions,
boxes_format="predictions"))
predictions = sorted(predictions, key=lambda p: p[1][:, 1].min())
left = []
right = []
for word, box in predictions:
if box[:, 0].min() < image.shape[1] / 2:
left.append((word, box))
else:
right.append((word, box))
ax.set_yticks([])
ax.set_xticks([])
for side, group in zip(["left", "right"], [left, right]):
for index, (text, box) in enumerate(group):
y = 1 - (index / len(group))
xy = box[0] / np.array([image.shape[1], image.shape[0]])
xy[1] = 1 - xy[1]
ax.annotate(
text=text,
xy=xy,
xytext=(-0.05 if side == "left" else 1.05, y),
xycoords="axes fraction",
arrowprops={"arrowstyle": "->", "color": "r"},
color="r",
fontsize=14,
horizontalalignment="right" if side == "left" else "left",
)
return ax
def drawBoxes(self, image, boxes, color=(255, 0, 0), thickness=1, boxes_format="boxes"):
"""Draw boxes onto an image.
Args:
image: The image on which to draw the boxes.
boxes: The boxes to draw.
color: The color for each box.
thickness: The thickness for each box.
boxes_format: The format used for providing the boxes. Options are
"boxes" which indicates an array with shape(N, 4, 2) where N is the
number of boxes and each box is a list of four points) as provided
by `keras_ocr.detection.Detector.detect`, "lines" (a list of
lines where each line itself is a list of (box, character) tuples) as
provided by `keras_ocr.data_generation.get_image_generator`,
or "predictions" where boxes is by itself a list of (word, box) tuples
as provided by `keras_ocr.pipeline.Pipeline.recognize` or
`keras_ocr.recognition.Recognizer.recognize_from_boxes`.
"""
if len(boxes) == 0:
return image
canvas = image.copy()
if boxes_format == "lines":
revised_boxes = []
for line in boxes:
for box, _ in line:
revised_boxes.append(box)
boxes = revised_boxes
if boxes_format == "predictions":
revised_boxes = []
for _, box in boxes:
revised_boxes.append(box)
boxes = revised_boxes
for box in boxes:
cv2.polylines(
img=canvas,
pts=box[np.newaxis].astype("int32"),
color=color,
thickness=thickness,
isClosed=True,
)
return canvas
|