openocr-python 0.0.2__py3-none-any.whl

openocr/tools/utils/e2e_utils/pgnet_pp_utils.py

@@ -0,0 +1,159 @@
+import torch
+import os
+import sys
+
+__dir__ = os.path.dirname(__file__)
+sys.path.append(__dir__)
+sys.path.append(os.path.join(__dir__, ".."))
+from extract_textpoint_slow import *
+from extract_textpoint_fast import generate_pivot_list_fast, restore_poly
+
+
+class PGNet_PostProcess(object):
+    # two different post-process
+    def __init__(
+            self,
+            character_dict_path,
+            valid_set,
+            score_thresh,
+            outs_dict,
+            shape_list,
+            point_gather_mode=None, ):
+        self.Lexicon_Table = get_dict(character_dict_path)
+        self.valid_set = valid_set
+        self.score_thresh = score_thresh
+        self.outs_dict = outs_dict
+        self.shape_list = shape_list
+        self.point_gather_mode = point_gather_mode
+
+    def pg_postprocess_fast(self):
+        p_score = self.outs_dict["f_score"]
+        p_border = self.outs_dict["f_border"]
+        p_char = self.outs_dict["f_char"]
+        p_direction = self.outs_dict["f_direction"]
+        if isinstance(p_score, torch.Tensor):
+            p_score = p_score[0].numpy()
+            p_border = p_border[0].numpy()
+            p_direction = p_direction[0].numpy()
+            p_char = p_char[0].numpy()
+        else:
+            p_score = p_score[0]
+            p_border = p_border[0]
+            p_direction = p_direction[0]
+            p_char = p_char[0]
+
+        src_h, src_w, ratio_h, ratio_w = self.shape_list[0]
+        instance_yxs_list, seq_strs = generate_pivot_list_fast(
+            p_score,
+            p_char,
+            p_direction,
+            self.Lexicon_Table,
+            score_thresh=self.score_thresh,
+            point_gather_mode=self.point_gather_mode, )
+        poly_list, keep_str_list = restore_poly(
+            instance_yxs_list,
+            seq_strs,
+            p_border,
+            ratio_w,
+            ratio_h,
+            src_w,
+            src_h,
+            self.valid_set, )
+        data = {
+            "points": poly_list,
+            "texts": keep_str_list,
+        }
+        return data
+
+    def pg_postprocess_slow(self):
+        p_score = self.outs_dict["f_score"]
+        p_border = self.outs_dict["f_border"]
+        p_char = self.outs_dict["f_char"]
+        p_direction = self.outs_dict["f_direction"]
+        if isinstance(p_score, torch.Tensor):
+            p_score = p_score[0].numpy()
+            p_border = p_border[0].numpy()
+            p_direction = p_direction[0].numpy()
+            p_char = p_char[0].numpy()
+        else:
+            p_score = p_score[0]
+            p_border = p_border[0]
+            p_direction = p_direction[0]
+            p_char = p_char[0]
+        src_h, src_w, ratio_h, ratio_w = self.shape_list[0]
+        is_curved = self.valid_set == "totaltext"
+        char_seq_idx_set, instance_yxs_list = generate_pivot_list_slow(
+            p_score,
+            p_char,
+            p_direction,
+            score_thresh=self.score_thresh,
+            is_backbone=True,
+            is_curved=is_curved, )
+        seq_strs = []
+        for char_idx_set in char_seq_idx_set:
+            pr_str = "".join([self.Lexicon_Table[pos] for pos in char_idx_set])
+            seq_strs.append(pr_str)
+        poly_list = []
+        keep_str_list = []
+        all_point_list = []
+        all_point_pair_list = []
+        for yx_center_line, keep_str in zip(instance_yxs_list, seq_strs):
+            if len(yx_center_line) == 1:
+                yx_center_line.append(yx_center_line[-1])
+
+            offset_expand = 1.0
+            if self.valid_set == "totaltext":
+                offset_expand = 1.2
+
+            point_pair_list = []
+            for batch_id, y, x in yx_center_line:
+                offset = p_border[:, y, x].reshape(2, 2)
+                if offset_expand != 1.0:
+                    offset_length = np.linalg.norm(
+                        offset, axis=1, keepdims=True)
+                    expand_length = np.clip(
+                        offset_length * (offset_expand - 1),
+                        a_min=0.5,
+                        a_max=3.0)
+                    offset_detal = offset / offset_length * expand_length
+                    offset = offset + offset_detal
+                ori_yx = np.array([y, x], dtype=np.float32)
+                point_pair = ((ori_yx + offset)[:, ::-1] * 4.0 /
+                              np.array([ratio_w, ratio_h]).reshape(-1, 2))
+                point_pair_list.append(point_pair)
+
+                all_point_list.append([
+                    int(round(x * 4.0 / ratio_w)),
+                    int(round(y * 4.0 / ratio_h))
+                ])
+                all_point_pair_list.append(point_pair.round().astype(np.int32)
+                                           .tolist())
+
+            detected_poly, pair_length_info = point_pair2poly(point_pair_list)
+            detected_poly = expand_poly_along_width(
+                detected_poly, shrink_ratio_of_width=0.2)
+            detected_poly[:, 0] = np.clip(
+                detected_poly[:, 0], a_min=0, a_max=src_w)
+            detected_poly[:, 1] = np.clip(
+                detected_poly[:, 1], a_min=0, a_max=src_h)
+
+            if len(keep_str) < 2:
+                continue
+
+            keep_str_list.append(keep_str)
+            detected_poly = np.round(detected_poly).astype("int32")
+            if self.valid_set == "partvgg":
+                middle_point = len(detected_poly) // 2
+                detected_poly = detected_poly[
+                    [0, middle_point - 1, middle_point, -1], :]
+                poly_list.append(detected_poly)
+            elif self.valid_set == "totaltext":
+                poly_list.append(detected_poly)
+            else:
+                print("--> Not supported format.")
+                exit(-1)
+        data = {
+            "points": poly_list,
+            "texts": keep_str_list,
+        }
+        return data

openocr/tools/utils/e2e_utils/visual.py

@@ -0,0 +1,152 @@
+import numpy as np
+import cv2
+import time
+
+
+def resize_image(im, max_side_len=512):
+    """
+    resize image to a size multiple of max_stride which is required by the network
+    :param im: the resized image
+    :param max_side_len: limit of max image size to avoid out of memory in gpu
+    :return: the resized image and the resize ratio
+    """
+    h, w, _ = im.shape
+
+    resize_w = w
+    resize_h = h
+
+    if resize_h > resize_w:
+        ratio = float(max_side_len) / resize_h
+    else:
+        ratio = float(max_side_len) / resize_w
+
+    resize_h = int(resize_h * ratio)
+    resize_w = int(resize_w * ratio)
+
+    max_stride = 128
+    resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
+    resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
+    im = cv2.resize(im, (int(resize_w), int(resize_h)))
+    ratio_h = resize_h / float(h)
+    ratio_w = resize_w / float(w)
+
+    return im, (ratio_h, ratio_w)
+
+
+def resize_image_min(im, max_side_len=512):
+    """ """
+    h, w, _ = im.shape
+
+    resize_w = w
+    resize_h = h
+
+    if resize_h < resize_w:
+        ratio = float(max_side_len) / resize_h
+    else:
+        ratio = float(max_side_len) / resize_w
+
+    resize_h = int(resize_h * ratio)
+    resize_w = int(resize_w * ratio)
+
+    max_stride = 128
+    resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
+    resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
+    im = cv2.resize(im, (int(resize_w), int(resize_h)))
+    ratio_h = resize_h / float(h)
+    ratio_w = resize_w / float(w)
+    return im, (ratio_h, ratio_w)
+
+
+def resize_image_for_totaltext(im, max_side_len=512):
+    """ """
+    h, w, _ = im.shape
+
+    resize_w = w
+    resize_h = h
+    ratio = 1.25
+    if h * ratio > max_side_len:
+        ratio = float(max_side_len) / resize_h
+
+    resize_h = int(resize_h * ratio)
+    resize_w = int(resize_w * ratio)
+
+    max_stride = 128
+    resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
+    resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
+    im = cv2.resize(im, (int(resize_w), int(resize_h)))
+    ratio_h = resize_h / float(h)
+    ratio_w = resize_w / float(w)
+    return im, (ratio_h, ratio_w)
+
+
+def point_pair2poly(point_pair_list):
+    """
+    Transfer vertical point_pairs into poly point in clockwise.
+    """
+    pair_length_list = []
+    for point_pair in point_pair_list:
+        pair_length = np.linalg.norm(point_pair[0] - point_pair[1])
+        pair_length_list.append(pair_length)
+    pair_length_list = np.array(pair_length_list)
+    pair_info = (
+        pair_length_list.max(),
+        pair_length_list.min(),
+        pair_length_list.mean(), )
+
+    point_num = len(point_pair_list) * 2
+    point_list = [0] * point_num
+    for idx, point_pair in enumerate(point_pair_list):
+        point_list[idx] = point_pair[0]
+        point_list[point_num - 1 - idx] = point_pair[1]
+    return np.array(point_list).reshape(-1, 2), pair_info
+
+
+def shrink_quad_along_width(quad, begin_width_ratio=0.0, end_width_ratio=1.0):
+    """
+    Generate shrink_quad_along_width.
+    """
+    ratio_pair = np.array(
+        [[begin_width_ratio], [end_width_ratio]], dtype=np.float32)
+    p0_1 = quad[0] + (quad[1] - quad[0]) * ratio_pair
+    p3_2 = quad[3] + (quad[2] - quad[3]) * ratio_pair
+    return np.array([p0_1[0], p0_1[1], p3_2[1], p3_2[0]])
+
+
+def expand_poly_along_width(poly, shrink_ratio_of_width=0.3):
+    """
+    expand poly along width.
+    """
+    point_num = poly.shape[0]
+    left_quad = np.array(
+        [poly[0], poly[1], poly[-2], poly[-1]], dtype=np.float32)
+    left_ratio = (-shrink_ratio_of_width *
+                  np.linalg.norm(left_quad[0] - left_quad[3]) /
+                  (np.linalg.norm(left_quad[0] - left_quad[1]) + 1e-6))
+    left_quad_expand = shrink_quad_along_width(left_quad, left_ratio, 1.0)
+    right_quad = np.array(
+        [
+            poly[point_num // 2 - 2],
+            poly[point_num // 2 - 1],
+            poly[point_num // 2],
+            poly[point_num // 2 + 1],
+        ],
+        dtype=np.float32, )
+    right_ratio = 1.0 + shrink_ratio_of_width * np.linalg.norm(right_quad[
+        0] - right_quad[3]) / (np.linalg.norm(right_quad[0] - right_quad[1]) +
+                               1e-6)
+    right_quad_expand = shrink_quad_along_width(right_quad, 0.0, right_ratio)
+    poly[0] = left_quad_expand[0]
+    poly[-1] = left_quad_expand[-1]
+    poly[point_num // 2 - 1] = right_quad_expand[1]
+    poly[point_num // 2] = right_quad_expand[2]
+    return poly
+
+
+def norm2(x, axis=None):
+    if axis:
+        return np.sqrt(np.sum(x**2, axis=axis))
+    return np.sqrt(np.sum(x**2))
+
+
+def cos(p1, p2):
+    return (p1 * p2).sum() / (norm2(p1) * norm2(p2))

openocr/tools/utils/en_dict.txt

@@ -0,0 +1,95 @@
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+:
+;
+<
+=
+>
+?
+@
+A
+B
+C
+D
+E
+F
+G
+H
+I
+J
+K
+L
+M
+N
+O
+P
+Q
+R
+S
+T
+U
+V
+W
+X
+Y
+Z
+[
+\
+]
+^
+_
+`
+a
+b
+c
+d
+e
+f
+g
+h
+i
+j
+k
+l
+m
+n
+o
+p
+q
+r
+s
+t
+u
+v
+w
+x
+y
+z
+{
+|
+}
+~
+!
+"
+#
+$
+%
+&
+'
+(
+)
+*
++
+,
+-
+.
+/
+ 

openocr/tools/utils/gen_label.py

@@ -0,0 +1,68 @@
+import os
+import argparse
+import json
+
+
+def gen_rec_label(input_path, out_label):
+    with open(out_label, "w") as out_file:
+        with open(input_path, "r") as f:
+            for line in f.readlines():
+                tmp = line.strip("\n").replace(" ", "").split(",")
+                img_path, label = tmp[0], tmp[1]
+                label = label.replace('"', "")
+                out_file.write(img_path + "\t" + label + "\n")
+
+
+def gen_det_label(root_path, input_dir, out_label):
+    with open(out_label, "w") as out_file:
+        for label_file in os.listdir(input_dir):
+            img_path = os.path.join(root_path, label_file[3:-4] + ".jpg")
+            label = []
+            with open(
+                    os.path.join(input_dir, label_file), "r",
+                    encoding="utf-8-sig") as f:
+                for line in f.readlines():
+                    tmp = line.strip("\n\r").replace("\xef\xbb\xbf",
+                                                     "").split(",")
+                    points = tmp[:8]
+                    s = []
+                    for i in range(0, len(points), 2):
+                        b = points[i:i + 2]
+                        b = [int(t) for t in b]
+                        s.append(b)
+                    result = {"transcription": tmp[8], "points": s}
+                    label.append(result)
+
+            out_file.write(img_path + "\t" + json.dumps(
+                label, ensure_ascii=False) + "\n")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--mode",
+        type=str,
+        default="rec",
+        help="Generate rec_label or det_label, can be set rec or det", )
+    parser.add_argument(
+        "--root_path",
+        type=str,
+        default=".",
+        help="The root directory of images.Only takes effect when mode=det ", )
+    parser.add_argument(
+        "--input_path",
+        type=str,
+        default=".",
+        help="Input_label or input path to be converted", )
+    parser.add_argument(
+        "--output_label",
+        type=str,
+        default="out_label.txt",
+        help="Output file name")
+
+    args = parser.parse_args()
+    if args.mode == "rec":
+        print("Generate rec label")
+        gen_rec_label(args.input_path, args.output_label)
+    elif args.mode == "det":
+        gen_det_label(args.root_path, args.input_path, args.output_label)

openocr/tools/utils/ic15_dict.txt

@@ -0,0 +1,36 @@
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+a
+b
+c
+d
+e
+f
+g
+h
+i
+j
+k
+l
+m
+n
+o
+p
+q
+r
+s
+t
+u
+v
+w
+x
+y
+z

openocr/tools/utils/logging.py

@@ -0,0 +1,56 @@
+import os
+import sys
+import logging
+import functools
+import torch
+import torch.distributed as dist
+
+logger_initialized = {}
+
+
+@functools.lru_cache()
+def get_logger(name="openrec", log_file=None, log_level=logging.DEBUG):
+    """Initialize and get a logger by name.
+    If the logger has not been initialized, this method will initialize the
+    logger by adding one or two handlers, otherwise the initialized logger will
+    be directly returned. During initialization, a StreamHandler will always be
+    added. If `log_file` is specified a FileHandler will also be added.
+    Args:
+        name (str): Logger name.
+        log_file (str | None): The log filename. If specified, a FileHandler
+            will be added to the logger.
+        log_level (int): The logger level. Note that only the process of
+            rank 0 is affected, and other processes will set the level to
+            "Error" thus be silent most of the time.
+    Returns:
+        logging.Logger: The expected logger.
+    """
+    logger = logging.getLogger(name)
+    if name in logger_initialized:
+        return logger
+    for logger_name in logger_initialized:
+        if name.startswith(logger_name):
+            return logger
+
+    formatter = logging.Formatter(
+        "[%(asctime)s] %(name)s %(levelname)s: %(message)s",
+        datefmt="%Y/%m/%d %H:%M:%S")
+
+    stream_handler = logging.StreamHandler(stream=sys.stdout)
+    stream_handler.setFormatter(formatter)
+    logger.addHandler(stream_handler)
+
+    rank = int(os.environ["LOCAL_RANK"]) if "LOCAL_RANK" in os.environ else 0
+    if log_file is not None and rank == 0:
+        log_file_folder = os.path.split(log_file)[0]
+        os.makedirs(log_file_folder, exist_ok=True)
+        file_handler = logging.FileHandler(log_file, "a")
+        file_handler.setFormatter(formatter)
+        logger.addHandler(file_handler)
+    if rank == 0:
+        logger.setLevel(log_level)
+    else:
+        logger.setLevel(logging.ERROR)
+    logger_initialized[name] = True
+    logger.propagate = False
+    return logger

openocr/tools/utils/poly_nms.py

@@ -0,0 +1,132 @@
+import numpy as np
+from shapely.geometry import Polygon
+
+
+def points2polygon(points):
+    """Convert k points to 1 polygon.
+
+    Args:
+        points (ndarray or list): A ndarray or a list of shape (2k)
+            that indicates k points.
+
+    Returns:
+        polygon (Polygon): A polygon object.
+    """
+    if isinstance(points, list):
+        points = np.array(points)
+
+    assert isinstance(points, np.ndarray)
+    assert (points.size % 2 == 0) and (points.size >= 8)
+
+    point_mat = points.reshape([-1, 2])
+    return Polygon(point_mat)
+
+
+def poly_intersection(poly_det, poly_gt, buffer=0.0001):
+    """Calculate the intersection area between two polygon.
+
+    Args:
+        poly_det (Polygon): A polygon predicted by detector.
+        poly_gt (Polygon): A gt polygon.
+
+    Returns:
+        intersection_area (float): The intersection area between two polygons.
+    """
+    assert isinstance(poly_det, Polygon)
+    assert isinstance(poly_gt, Polygon)
+
+    if buffer == 0:
+        poly_inter = poly_det & poly_gt
+    else:
+        poly_inter = poly_det.buffer(buffer) & poly_gt.buffer(buffer)
+    return poly_inter.area, poly_inter
+
+
+def poly_union(poly_det, poly_gt):
+    """Calculate the union area between two polygon.
+
+    Args:
+        poly_det (Polygon): A polygon predicted by detector.
+        poly_gt (Polygon): A gt polygon.
+
+    Returns:
+        union_area (float): The union area between two polygons.
+    """
+    assert isinstance(poly_det, Polygon)
+    assert isinstance(poly_gt, Polygon)
+
+    area_det = poly_det.area
+    area_gt = poly_gt.area
+    area_inters, _ = poly_intersection(poly_det, poly_gt)
+    return area_det + area_gt - area_inters
+
+
+def valid_boundary(x, with_score=True):
+    num = len(x)
+    if num < 8:
+        return False
+    if num % 2 == 0 and (not with_score):
+        return True
+    if num % 2 == 1 and with_score:
+        return True
+
+    return False
+
+
+def boundary_iou(src, target):
+    """Calculate the IOU between two boundaries.
+
+    Args:
+       src (list): Source boundary.
+       target (list): Target boundary.
+
+    Returns:
+       iou (float): The iou between two boundaries.
+    """
+    assert valid_boundary(src, False)
+    assert valid_boundary(target, False)
+    src_poly = points2polygon(src)
+    target_poly = points2polygon(target)
+
+    return poly_iou(src_poly, target_poly)
+
+
+def poly_iou(poly_det, poly_gt):
+    """Calculate the IOU between two polygons.
+
+    Args:
+        poly_det (Polygon): A polygon predicted by detector.
+        poly_gt (Polygon): A gt polygon.
+
+    Returns:
+        iou (float): The IOU between two polygons.
+    """
+    assert isinstance(poly_det, Polygon)
+    assert isinstance(poly_gt, Polygon)
+    area_inters, _ = poly_intersection(poly_det, poly_gt)
+    area_union = poly_union(poly_det, poly_gt)
+    if area_union == 0:
+        return 0.0
+    return area_inters / area_union
+
+
+def poly_nms(polygons, threshold):
+    assert isinstance(polygons, list)
+
+    polygons = np.array(sorted(polygons, key=lambda x: x[-1]))
+
+    keep_poly = []
+    index = [i for i in range(polygons.shape[0])]
+
+    while len(index) > 0:
+        keep_poly.append(polygons[index[-1]].tolist())
+        A = polygons[index[-1]][:-1]
+        index = np.delete(index, -1)
+        iou_list = np.zeros((len(index), ))
+        for i in range(len(index)):
+            B = polygons[index[i]][:-1]
+            iou_list[i] = boundary_iou(A, B)
+        remove_index = np.where(iou_list > threshold)
+        index = np.delete(index, remove_index)
+
+    return keep_poly