370 lines
14 KiB
Python
370 lines
14 KiB
Python
|
from PIL import ImageFont, ImageDraw, Image, ImageOps
|
||
|
|
# import matplotlib.pyplot as plt
|
||
|
|
import numpy as np
|
||
|
|
import cv2
|
||
|
|
import os
|
||
|
|
import time
|
||
|
|
from typing import Generator, Union, List, overload, Tuple, Callable
|
||
|
|
import glob
|
||
|
|
import math
|
||
|
|
from pathlib import Path
|
||
|
|
from pdf2image import convert_from_path
|
||
|
|
# from deskew import determine_skew
|
||
|
|
# from jdeskew.estimator import get_angle
|
||
|
|
# from jdeskew.utility import rotate as jrotate
|
||
|
|
|
||
|
|
|
||
|
|
def post_process_recog(text: str) -> str:
|
||
|
|
text = text.replace("✪", " ")
|
||
|
|
return text
|
||
|
|
|
||
|
|
|
||
|
|
def find_maximum_without_outliers(lst: list[int], threshold: float = 1.):
|
||
|
|
'''
|
||
|
|
To find the maximum number in a list while excluding its outlier values, you can follow these steps:
|
||
|
|
Determine the range within which you consider values as outliers. This can be based on a specific threshold or a statistical measure such as the interquartile range (IQR).
|
||
|
|
Iterate through the list and filter out the outlier values based on the defined range. Keep track of the non-outlier values.
|
||
|
|
Find the maximum value among the non-outlier values.
|
||
|
|
'''
|
||
|
|
# Calculate the lower and upper boundaries for outliers
|
||
|
|
q1 = np.percentile(lst, 25)
|
||
|
|
q3 = np.percentile(lst, 75)
|
||
|
|
iqr = q3 - q1
|
||
|
|
lower_bound = q1 - threshold * iqr
|
||
|
|
upper_bound = q3 + threshold * iqr
|
||
|
|
|
||
|
|
# Filter out outlier values
|
||
|
|
non_outliers = [x for x in lst if lower_bound <= x <= upper_bound]
|
||
|
|
|
||
|
|
# Find the maximum value among non-outliers
|
||
|
|
max_value = max(non_outliers)
|
||
|
|
|
||
|
|
return max_value
|
||
|
|
|
||
|
|
|
||
|
|
class Timer:
|
||
|
|
def __init__(self, name: str) -> None:
|
||
|
|
self.name = name
|
||
|
|
|
||
|
|
def __enter__(self):
|
||
|
|
self.start_time = time.perf_counter()
|
||
|
|
return self
|
||
|
|
|
||
|
|
def __exit__(self, func: Callable, *args):
|
||
|
|
self.end_time = time.perf_counter()
|
||
|
|
self.elapsed_time = self.end_time - self.start_time
|
||
|
|
print(f"[INFO]: {self.name} took : {self.elapsed_time:.6f} seconds")
|
||
|
|
|
||
|
|
|
||
|
|
# def rotate(
|
||
|
|
# image: np.ndarray, angle: float, background: Union[int, Tuple[int, int, int]]
|
||
|
|
# ) -> np.ndarray:
|
||
|
|
# old_width, old_height = image.shape[:2]
|
||
|
|
# angle_radian = math.radians(angle)
|
||
|
|
# width = abs(np.sin(angle_radian) * old_height) + abs(np.cos(angle_radian) * old_width)
|
||
|
|
# height = abs(np.sin(angle_radian) * old_width) + abs(np.cos(angle_radian) * old_height)
|
||
|
|
# image_center = tuple(np.array(image.shape[1::-1]) / 2)
|
||
|
|
# rot_mat = cv2.getRotationMatrix2D(image_center, angle, 1.0)
|
||
|
|
# rot_mat[1, 2] += (width - old_width) / 2
|
||
|
|
# rot_mat[0, 2] += (height - old_height) / 2
|
||
|
|
# return cv2.warpAffine(image, rot_mat, (int(round(height)), int(round(width))), borderValue=background)
|
||
|
|
|
||
|
|
|
||
|
|
# def rotate_bbox(bbox: list, angle: float) -> list:
|
||
|
|
# # Compute the center point of the bounding box
|
||
|
|
# cx = bbox[0] + bbox[2] / 2
|
||
|
|
# cy = bbox[1] + bbox[3] / 2
|
||
|
|
|
||
|
|
# # Define the scale factor for the rotated bounding box
|
||
|
|
# scale = 1.0 # following the deskew and jdeskew function
|
||
|
|
# angle_radian = math.radians(angle)
|
||
|
|
|
||
|
|
# # Obtain the rotation matrix using cv2.getRotationMatrix2D()
|
||
|
|
# M = cv2.getRotationMatrix2D((cx, cy), angle_radian, scale)
|
||
|
|
|
||
|
|
# # Apply the rotation matrix to the four corners of the bounding box
|
||
|
|
# corners = np.array([[bbox[0], bbox[1]],
|
||
|
|
# [bbox[0] + bbox[2], bbox[1]],
|
||
|
|
# [bbox[0] + bbox[2], bbox[1] + bbox[3]],
|
||
|
|
# [bbox[0], bbox[1] + bbox[3]]], dtype=np.float32)
|
||
|
|
# rotated_corners = cv2.transform(np.array([corners]), M)[0]
|
||
|
|
|
||
|
|
# # Compute the bounding box of the rotated corners
|
||
|
|
# x = int(np.min(rotated_corners[:, 0]))
|
||
|
|
# y = int(np.min(rotated_corners[:, 1]))
|
||
|
|
# w = int(np.max(rotated_corners[:, 0]) - np.min(rotated_corners[:, 0]))
|
||
|
|
# h = int(np.max(rotated_corners[:, 1]) - np.min(rotated_corners[:, 1]))
|
||
|
|
# rotated_bbox = [x, y, w, h]
|
||
|
|
|
||
|
|
# return rotated_bbox
|
||
|
|
|
||
|
|
# def rotate_bbox(bbox: List[int], angle: float, old_shape: Tuple[int, int]) -> List[int]:
|
||
|
|
# # https://medium.com/@pokomaru/image-and-bounding-box-rotation-using-opencv-python-2def6c39453
|
||
|
|
# bbox_ = [bbox[0], bbox[1], bbox[2], bbox[1], bbox[2], bbox[3], bbox[0], bbox[3]]
|
||
|
|
# h, w = old_shape
|
||
|
|
# cx, cy = (int(w / 2), int(h / 2))
|
||
|
|
|
||
|
|
# bbox_tuple = [
|
||
|
|
# (bbox_[0], bbox_[1]),
|
||
|
|
# (bbox_[2], bbox_[3]),
|
||
|
|
# (bbox_[4], bbox_[5]),
|
||
|
|
# (bbox_[6], bbox_[7]),
|
||
|
|
# ] # put x and y coordinates in tuples, we will iterate through the tuples and perform rotation
|
||
|
|
|
||
|
|
# rotated_bbox = []
|
||
|
|
|
||
|
|
# for i, coord in enumerate(bbox_tuple):
|
||
|
|
# M = cv2.getRotationMatrix2D((cx, cy), angle, 1.0)
|
||
|
|
# cos, sin = abs(M[0, 0]), abs(M[0, 1])
|
||
|
|
# newW = int((h * sin) + (w * cos))
|
||
|
|
# newH = int((h * cos) + (w * sin))
|
||
|
|
# M[0, 2] += (newW / 2) - cx
|
||
|
|
# M[1, 2] += (newH / 2) - cy
|
||
|
|
# v = [coord[0], coord[1], 1]
|
||
|
|
# adjusted_coord = np.dot(M, v)
|
||
|
|
# rotated_bbox.insert(i, (adjusted_coord[0], adjusted_coord[1]))
|
||
|
|
# result = [int(x) for t in rotated_bbox for x in t]
|
||
|
|
# return [result[i] for i in [0, 1, 2, -1]] # reformat to xyxy
|
||
|
|
|
||
|
|
|
||
|
|
# def deskew(image: np.ndarray) -> Tuple[np.ndarray, float]:
|
||
|
|
# grayscale = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||
|
|
# angle = 0.
|
||
|
|
# try:
|
||
|
|
# angle = determine_skew(grayscale)
|
||
|
|
# except Exception:
|
||
|
|
# pass
|
||
|
|
# rotated = rotate(image, angle, (0, 0, 0)) if angle else image
|
||
|
|
# return rotated, angle
|
||
|
|
|
||
|
|
|
||
|
|
# def jdeskew(image: np.ndarray) -> Tuple[np.ndarray, float]:
|
||
|
|
# angle = 0.
|
||
|
|
# try:
|
||
|
|
# angle = get_angle(image)
|
||
|
|
# except Exception:
|
||
|
|
# pass
|
||
|
|
# # TODO: change resize = True and scale the bounding box
|
||
|
|
# rotated = jrotate(image, angle, resize=False) if angle else image
|
||
|
|
# return rotated, angle
|
||
|
|
# def deskew()
|
||
|
|
|
||
|
|
class ImageReader:
|
||
|
|
"""
|
||
|
|
accept anything, return numpy array image
|
||
|
|
"""
|
||
|
|
supported_ext = [".png", ".jpg", ".jpeg", ".pdf", ".gif"]
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def validate_img_path(img_path: str) -> None:
|
||
|
|
if not os.path.exists(img_path):
|
||
|
|
raise FileNotFoundError(img_path)
|
||
|
|
if os.path.isdir(img_path):
|
||
|
|
raise IsADirectoryError(img_path)
|
||
|
|
if not Path(img_path).suffix.lower() in ImageReader.supported_ext:
|
||
|
|
raise NotImplementedError("Not supported extension at {}".format(img_path))
|
||
|
|
|
||
|
|
@overload
|
||
|
|
@staticmethod
|
||
|
|
def read(img: Union[str, np.ndarray, Image.Image]) -> np.ndarray: ...
|
||
|
|
|
||
|
|
@overload
|
||
|
|
@staticmethod
|
||
|
|
def read(img: List[Union[str, np.ndarray, Image.Image]]) -> List[np.ndarray]: ...
|
||
|
|
|
||
|
|
@overload
|
||
|
|
@staticmethod
|
||
|
|
def read(img: str) -> List[np.ndarray]: ... # for pdf or directory
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def read(img):
|
||
|
|
if isinstance(img, list):
|
||
|
|
return ImageReader.from_list(img)
|
||
|
|
elif isinstance(img, str) and os.path.isdir(img):
|
||
|
|
return ImageReader.from_dir(img)
|
||
|
|
elif isinstance(img, str) and img.endswith(".pdf"):
|
||
|
|
return ImageReader.from_pdf(img)
|
||
|
|
else:
|
||
|
|
return ImageReader._read(img)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def from_dir(dir_path: str) -> List[np.ndarray]:
|
||
|
|
if os.path.isdir(dir_path):
|
||
|
|
image_files = glob.glob(os.path.join(dir_path, "*"))
|
||
|
|
return ImageReader.from_list(image_files)
|
||
|
|
else:
|
||
|
|
raise NotADirectoryError(dir_path)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def from_str(img_path: str) -> np.ndarray:
|
||
|
|
ImageReader.validate_img_path(img_path)
|
||
|
|
return ImageReader.from_PIL(Image.open(img_path))
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def from_np(img_array: np.ndarray) -> np.ndarray:
|
||
|
|
return img_array
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def from_PIL(img_pil: Image.Image, transpose=True) -> np.ndarray:
|
||
|
|
# if img_pil.is_animated:
|
||
|
|
# raise NotImplementedError("Only static images are supported, animated image found")
|
||
|
|
if transpose:
|
||
|
|
img_pil = ImageOps.exif_transpose(img_pil)
|
||
|
|
if img_pil.mode != "RGB":
|
||
|
|
img_pil = img_pil.convert("RGB")
|
||
|
|
|
||
|
|
return np.array(img_pil)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def from_list(img_list: List[Union[str, np.ndarray, Image.Image]]) -> List[np.ndarray]:
|
||
|
|
limgs = list()
|
||
|
|
for img_path in img_list:
|
||
|
|
try:
|
||
|
|
if isinstance(img_path, str):
|
||
|
|
ImageReader.validate_img_path(img_path)
|
||
|
|
limgs.append(ImageReader._read(img_path))
|
||
|
|
except (FileNotFoundError, NotImplementedError, IsADirectoryError) as e:
|
||
|
|
print("[ERROR]: ", e)
|
||
|
|
print("[INFO]: Skipping image {}".format(img_path))
|
||
|
|
return limgs
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def from_pdf(pdf_path: str, start_page: int = 0, end_page: int = 0) -> List[np.ndarray]:
|
||
|
|
pdf_file = convert_from_path(pdf_path)
|
||
|
|
if end_page is not None:
|
||
|
|
end_page = min(len(pdf_file), end_page + 1)
|
||
|
|
limgs = [np.array(pdf_page) for pdf_page in pdf_file[start_page:end_page]]
|
||
|
|
return limgs
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _read(img: Union[str, np.ndarray, Image.Image]) -> np.ndarray:
|
||
|
|
if isinstance(img, str):
|
||
|
|
return ImageReader.from_str(img)
|
||
|
|
elif isinstance(img, Image.Image):
|
||
|
|
return ImageReader.from_PIL(img)
|
||
|
|
elif isinstance(img, np.ndarray):
|
||
|
|
return ImageReader.from_np(img)
|
||
|
|
else:
|
||
|
|
raise ValueError("Invalid img argument type: ", type(img))
|
||
|
|
|
||
|
|
|
||
|
|
def get_name(file_path, ext: bool = True):
|
||
|
|
file_path_ = os.path.basename(file_path)
|
||
|
|
return file_path_ if ext else os.path.splitext(file_path_)[0]
|
||
|
|
|
||
|
|
|
||
|
|
def construct_file_path(dir, file_path, ext=''):
|
||
|
|
'''
|
||
|
|
args:
|
||
|
|
dir: /path/to/dir
|
||
|
|
file_path /example_path/to/file.txt
|
||
|
|
ext = '.json'
|
||
|
|
return
|
||
|
|
/path/to/dir/file.json
|
||
|
|
'''
|
||
|
|
return os.path.join(
|
||
|
|
dir, get_name(file_path,
|
||
|
|
True)) if ext == '' else os.path.join(
|
||
|
|
dir, get_name(file_path,
|
||
|
|
False)) + ext
|
||
|
|
|
||
|
|
|
||
|
|
def chunks(lst: list, n: int) -> Generator:
|
||
|
|
"""
|
||
|
|
Yield successive n-sized chunks from lst.
|
||
|
|
https://stackoverflow.com/questions/312443/how-do-i-split-a-list-into-equally-sized-chunks
|
||
|
|
"""
|
||
|
|
for i in range(0, len(lst), n):
|
||
|
|
yield lst[i:i + n]
|
||
|
|
|
||
|
|
|
||
|
|
def read_ocr_result_from_txt(file_path: str) -> Tuple[list, list]:
|
||
|
|
'''
|
||
|
|
return list of bounding boxes, list of words
|
||
|
|
'''
|
||
|
|
with open(file_path, 'r') as f:
|
||
|
|
lines = f.read().splitlines()
|
||
|
|
boxes, words = [], []
|
||
|
|
for line in lines:
|
||
|
|
if line == "":
|
||
|
|
continue
|
||
|
|
x1, y1, x2, y2, text = line.split("\t")
|
||
|
|
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
|
||
|
|
if text and text != " ":
|
||
|
|
words.append(text)
|
||
|
|
boxes.append((x1, y1, x2, y2))
|
||
|
|
return boxes, words
|
||
|
|
|
||
|
|
|
||
|
|
def get_xyxywh_base_on_format(bbox, format):
|
||
|
|
if format == "xywh":
|
||
|
|
x1, y1, w, h = bbox[0], bbox[1], bbox[2], bbox[3]
|
||
|
|
x2, y2 = x1 + w, y1 + h
|
||
|
|
elif format == "xyxy":
|
||
|
|
x1, y1, x2, y2 = bbox
|
||
|
|
w, h = x2 - x1, y2 - y1
|
||
|
|
else:
|
||
|
|
raise NotImplementedError("Invalid format {}".format(format))
|
||
|
|
return (x1, y1, x2, y2, w, h)
|
||
|
|
|
||
|
|
|
||
|
|
def get_dynamic_params_for_bbox_of_label(text, x1, y1, w, h, img_h, img_w, font, font_scale_offset=1):
|
||
|
|
font_scale_factor = img_h / (img_w + img_h) * font_scale_offset
|
||
|
|
font_scale = w / (w + h) * font_scale_factor # adjust font scale by width height
|
||
|
|
thickness = int(font_scale_factor) + 1
|
||
|
|
(text_width, text_height) = cv2.getTextSize(text, font, fontScale=font_scale, thickness=thickness)[0]
|
||
|
|
text_offset_x = x1
|
||
|
|
text_offset_y = y1 - thickness
|
||
|
|
box_coords = ((text_offset_x, text_offset_y + 1), (text_offset_x + text_width - 2, text_offset_y - text_height - 2))
|
||
|
|
return (font_scale, thickness, text_height, box_coords)
|
||
|
|
|
||
|
|
|
||
|
|
def visualize_bbox_and_label(
|
||
|
|
img, bboxes, texts, bbox_color=(200, 180, 60),
|
||
|
|
text_color=(0, 0, 0),
|
||
|
|
format="xyxy", is_vnese=False, draw_text=True):
|
||
|
|
ori_img_type = type(img)
|
||
|
|
if is_vnese:
|
||
|
|
img = Image.fromarray(img) if ori_img_type is np.ndarray else img
|
||
|
|
draw = ImageDraw.Draw(img)
|
||
|
|
img_w, img_h = img.size
|
||
|
|
font_pil_str = "fonts/arial.ttf"
|
||
|
|
font_cv2 = cv2.FONT_HERSHEY_SIMPLEX
|
||
|
|
else:
|
||
|
|
img_h, img_w = img.shape[0], img.shape[1]
|
||
|
|
font_cv2 = cv2.FONT_HERSHEY_SIMPLEX
|
||
|
|
for i in range(len(bboxes)):
|
||
|
|
text = texts[i] # text = "{}: {:.0f}%".format(LABELS[classIDs[i]], confidences[i]*100)
|
||
|
|
x1, y1, x2, y2, w, h = get_xyxywh_base_on_format(bboxes[i], format)
|
||
|
|
font_scale, thickness, text_height, box_coords = get_dynamic_params_for_bbox_of_label(
|
||
|
|
text, x1, y1, w, h, img_h, img_w, font=font_cv2)
|
||
|
|
if is_vnese:
|
||
|
|
font_pil = ImageFont.truetype(font_pil_str, size=text_height) # type: ignore
|
||
|
|
fdraw_text = draw.text # type: ignore
|
||
|
|
fdraw_bbox = draw.rectangle # type: ignore
|
||
|
|
# Pil use different coordinate => y = y+thickness = y-thickness + 2*thickness
|
||
|
|
arg_text = ((box_coords[0][0], box_coords[1][1]), text)
|
||
|
|
kwarg_text = {"font": font_pil, "fill": text_color, "width": thickness}
|
||
|
|
arg_rec = ((x1, y1, x2, y2),)
|
||
|
|
kwarg_rec = {"outline": bbox_color, "width": thickness}
|
||
|
|
arg_rec_text = ((box_coords[0], box_coords[1]),)
|
||
|
|
kwarg_rec_text = {"fill": bbox_color, "width": thickness}
|
||
|
|
else:
|
||
|
|
# cv2.rectangle(img, box_coords[0], box_coords[1], color, cv2.FILLED)
|
||
|
|
# cv2.putText(img, text, (text_offset_x, text_offset_y), font, fontScale=font_scale, color=(50, 0,0), thickness=thickness)
|
||
|
|
# cv2.rectangle(img, (x1, y1), (x2, y2), color, thickness)
|
||
|
|
fdraw_text = cv2.putText
|
||
|
|
fdraw_bbox = cv2.rectangle
|
||
|
|
arg_text = (img, text, box_coords[0])
|
||
|
|
kwarg_text = {"fontFace": font_cv2, "fontScale": font_scale, "color": text_color, "thickness": thickness}
|
||
|
|
arg_rec = (img, (x1, y1), (x2, y2))
|
||
|
|
kwarg_rec = {"color": bbox_color, "thickness": thickness}
|
||
|
|
arg_rec_text = (img, box_coords[0], box_coords[1])
|
||
|
|
kwarg_rec_text = {"color": bbox_color, "thickness": cv2.FILLED}
|
||
|
|
# draw a bounding box rectangle and label on the img
|
||
|
|
fdraw_bbox(*arg_rec, **kwarg_rec) # type: ignore
|
||
|
|
if draw_text:
|
||
|
|
fdraw_bbox(*arg_rec_text, **kwarg_rec_text) # type: ignore
|
||
|
|
fdraw_text(*arg_text, **kwarg_text) # type: ignore # text have to put in front of rec_text
|
||
|
|
return np.array(img) if ori_img_type is np.ndarray and is_vnese else img
|