Spaces:
Sleeping
Sleeping
File size: 2,439 Bytes
aa1a0f3 | 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 | from akaocr import TextEngine, BoxEngine
import cv2
import numpy as np
from PIL import Image
import re
def transform_image(image, box):
# Get perspective transform image
assert len(box) == 4, "Shape of points must be 4x2"
img_crop_width = int(
max(
np.linalg.norm(box[0] - box[1]),
np.linalg.norm(box[2] - box[3])))
img_crop_height = int(
max(
np.linalg.norm(box[0] - box[3]),
np.linalg.norm(box[1] - box[2])))
pts_std = np.float32([[0, 0],
[img_crop_width, 0],
[img_crop_width, img_crop_height],
[0, img_crop_height]])
box = np.array(box, dtype="float32")
M = cv2.getPerspectiveTransform(box, pts_std)
dst_img = cv2.warpPerspective(
image,
M, (img_crop_width, img_crop_height),
borderMode=cv2.BORDER_REPLICATE,
flags=cv2.INTER_CUBIC)
img_height, img_width = dst_img.shape[0:2]
if img_height/img_width >= 1.25:
dst_img = np.rot90(dst_img, k=3)
return dst_img
def two_pts(bounding_box):
# convert 4-points-bounding-box to 2-points-bounding-box
return (
(
round(min([x[0] for x in bounding_box])),
round(min([x[1] for x in bounding_box]))
),
(
round(max([x[0] for x in bounding_box])),
round(max([x[1] for x in bounding_box]))
)
)
class BIB_Extract:
def __init__(self):
# Initialize the OCR engines
self.box_engine = BoxEngine()
self.text_engine = TextEngine()
def __call__(self, image, bib_length):
boxes = self.box_engine(image)
images = []
# crop and transform images for recognition
for box in boxes[::-1]:
# org_image = cv2.polylines(org_image, [box.astype(
# np.int32)], isClosed=True, color=(0, 255, 0), thickness=2)
crop_img = transform_image(image, box)
images.append(crop_img)
# Get the texts from the boxes
texts = self.text_engine(images)
return self.BIB_filter(texts, bib_length)
def BIB_filter(self, texts, bib_length):
pattern = rf'^\d{{{bib_length}}}$'
return [s[0] for s in texts if re.match(pattern, s[0])]
if __name__ == '__main__':
image = cv2.imread("1.jpg")
engine = BIB_Extract()
print(engine(image, bib_length=4))
|