openocr-python 0.0.2__py3-none-any.whl

openocr/opendet/postprocess/db_postprocess.py

@@ -0,0 +1,273 @@
+import numpy as np
+import cv2
+import torch
+from shapely.geometry import Polygon
+import pyclipper
+"""
+This code is refered from:
+https://github.com/WenmuZhou/DBNet.pytorch/blob/master/post_processing/seg_detector_representer.py
+"""
+
+
+class DBPostProcess(object):
+    """
+    The post process for Differentiable Binarization (DB).
+    """
+
+    def __init__(
+        self,
+        thresh=0.3,
+        box_thresh=0.7,
+        max_candidates=1000,
+        unclip_ratio=2.0,
+        use_dilation=False,
+        score_mode='fast',
+        box_type='quad',
+        **kwargs,
+    ):
+        self.thresh = thresh
+        self.box_thresh = box_thresh
+        self.max_candidates = max_candidates
+        self.unclip_ratio = unclip_ratio
+        self.min_size = 3
+        self.score_mode = score_mode
+        self.box_type = box_type
+        assert score_mode in [
+            'slow',
+            'fast',
+        ], 'Score mode must be in [slow, fast] but got: {}'.format(score_mode)
+
+        self.dilation_kernel = None if not use_dilation else np.array([[1, 1],
+                                                                       [1, 1]])
+
+    def polygons_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
+        """
+        _bitmap: single map with shape (1, H, W),
+            whose values are binarized as {0, 1}
+        """
+
+        bitmap = _bitmap
+        height, width = bitmap.shape
+
+        boxes = []
+        scores = []
+
+        contours, _ = cv2.findContours((bitmap * 255).astype(np.uint8),
+                                       cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
+
+        for contour in contours[:self.max_candidates]:
+            epsilon = 0.002 * cv2.arcLength(contour, True)
+            approx = cv2.approxPolyDP(contour, epsilon, True)
+            points = approx.reshape((-1, 2))
+            if points.shape[0] < 4:
+                continue
+
+            score = self.box_score_fast(pred, points.reshape(-1, 2))
+            if self.box_thresh > score:
+                continue
+
+            if points.shape[0] > 2:
+                box = self.unclip(points, self.unclip_ratio)
+                if len(box) > 1:
+                    continue
+            else:
+                continue
+            box = np.array(box).reshape(-1, 2)
+            if len(box) == 0:
+                continue
+
+            _, sside = self.get_mini_boxes(box.reshape((-1, 1, 2)))
+            if sside < self.min_size + 2:
+                continue
+
+            box = np.array(box)
+            box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0,
+                                dest_width)
+            box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0,
+                                dest_height)
+            boxes.append(box.tolist())
+            scores.append(score)
+        return boxes, scores
+
+    def boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
+        """
+        _bitmap: single map with shape (1, H, W),
+                whose values are binarized as {0, 1}
+        """
+
+        bitmap = _bitmap
+        height, width = bitmap.shape
+
+        outs = cv2.findContours((bitmap * 255).astype(np.uint8), cv2.RETR_LIST,
+                                cv2.CHAIN_APPROX_SIMPLE)
+        if len(outs) == 3:
+            img, contours, _ = outs[0], outs[1], outs[2]
+        elif len(outs) == 2:
+            contours, _ = outs[0], outs[1]
+
+        num_contours = min(len(contours), self.max_candidates)
+
+        boxes = []
+        scores = []
+        for index in range(num_contours):
+            contour = contours[index]
+            points, sside = self.get_mini_boxes(contour)
+            if sside < self.min_size:
+                continue
+            points = np.array(points)
+            if self.score_mode == 'fast':
+                score = self.box_score_fast(pred, points.reshape(-1, 2))
+            else:
+                score = self.box_score_slow(pred, contour)
+            if self.box_thresh > score:
+                continue
+
+            box = self.unclip(points, self.unclip_ratio)
+            if len(box) > 1:
+                continue
+            box = np.array(box).reshape(-1, 1, 2)
+            box, sside = self.get_mini_boxes(box)
+            if sside < self.min_size + 2:
+                continue
+            box = np.array(box)
+
+            box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0,
+                                dest_width)
+            box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0,
+                                dest_height)
+            boxes.append(box.astype('int32'))
+            scores.append(score)
+        return np.array(boxes, dtype='int32'), scores
+
+    def unclip(self, box, unclip_ratio):
+        poly = Polygon(box)
+        distance = poly.area * unclip_ratio / poly.length
+        offset = pyclipper.PyclipperOffset()
+        offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
+        expanded = offset.Execute(distance)
+        return expanded
+
+    def get_mini_boxes(self, contour):
+        bounding_box = cv2.minAreaRect(contour)
+        points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])
+
+        index_1, index_2, index_3, index_4 = 0, 1, 2, 3
+        if points[1][1] > points[0][1]:
+            index_1 = 0
+            index_4 = 1
+        else:
+            index_1 = 1
+            index_4 = 0
+        if points[3][1] > points[2][1]:
+            index_2 = 2
+            index_3 = 3
+        else:
+            index_2 = 3
+            index_3 = 2
+
+        box = [
+            points[index_1], points[index_2], points[index_3], points[index_4]
+        ]
+        return box, min(bounding_box[1])
+
+    def box_score_fast(self, bitmap, _box):
+        """
+        box_score_fast: use bbox mean score as the mean score
+        """
+        h, w = bitmap.shape[:2]
+        box = _box.copy()
+        xmin = np.clip(np.floor(box[:, 0].min()).astype('int32'), 0, w - 1)
+        xmax = np.clip(np.ceil(box[:, 0].max()).astype('int32'), 0, w - 1)
+        ymin = np.clip(np.floor(box[:, 1].min()).astype('int32'), 0, h - 1)
+        ymax = np.clip(np.ceil(box[:, 1].max()).astype('int32'), 0, h - 1)
+
+        mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
+        box[:, 0] = box[:, 0] - xmin
+        box[:, 1] = box[:, 1] - ymin
+        cv2.fillPoly(mask, box.reshape(1, -1, 2).astype('int32'), 1)
+        return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]
+
+    def box_score_slow(self, bitmap, contour):
+        """
+        box_score_slow: use polyon mean score as the mean score
+        """
+        h, w = bitmap.shape[:2]
+        contour = contour.copy()
+        contour = np.reshape(contour, (-1, 2))
+
+        xmin = np.clip(np.min(contour[:, 0]), 0, w - 1)
+        xmax = np.clip(np.max(contour[:, 0]), 0, w - 1)
+        ymin = np.clip(np.min(contour[:, 1]), 0, h - 1)
+        ymax = np.clip(np.max(contour[:, 1]), 0, h - 1)
+
+        mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
+
+        contour[:, 0] = contour[:, 0] - xmin
+        contour[:, 1] = contour[:, 1] - ymin
+
+        cv2.fillPoly(mask, contour.reshape(1, -1, 2).astype('int32'), 1)
+        return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]
+
+    def __call__(self, outs_dict, shape_list):
+        pred = outs_dict['maps']
+        if isinstance(pred, torch.Tensor):
+            pred = pred.detach().cpu().numpy()
+        pred = pred[:, 0, :, :]
+        segmentation = pred > self.thresh
+
+        boxes_batch = []
+        for batch_index in range(pred.shape[0]):
+            src_h, src_w, ratio_h, ratio_w = shape_list[batch_index]
+            if self.dilation_kernel is not None:
+                mask = cv2.dilate(
+                    np.array(segmentation[batch_index]).astype(np.uint8),
+                    self.dilation_kernel,
+                )
+            else:
+                mask = segmentation[batch_index]
+            if self.box_type == 'poly':
+                boxes, scores = self.polygons_from_bitmap(
+                    pred[batch_index], mask, src_w, src_h)
+            elif self.box_type == 'quad':
+                boxes, scores = self.boxes_from_bitmap(pred[batch_index], mask,
+                                                       src_w, src_h)
+            else:
+                raise ValueError(
+                    "box_type can only be one of ['quad', 'poly']")
+
+            boxes_batch.append({'points': boxes})
+        return boxes_batch
+
+
+class DistillationDBPostProcess(object):
+
+    def __init__(
+        self,
+        model_name=['student'],
+        key=None,
+        thresh=0.3,
+        box_thresh=0.6,
+        max_candidates=1000,
+        unclip_ratio=1.5,
+        use_dilation=False,
+        score_mode='fast',
+        box_type='quad',
+        **kwargs,
+    ):
+        self.model_name = model_name
+        self.key = key
+        self.post_process = DBPostProcess(
+            thresh=thresh,
+            box_thresh=box_thresh,
+            max_candidates=max_candidates,
+            unclip_ratio=unclip_ratio,
+            use_dilation=use_dilation,
+            score_mode=score_mode,
+            box_type=box_type,
+        )
+
+    def __call__(self, predicts, shape_list):
+        results = {}
+        for k in self.model_name:
+            results[k] = self.post_process(predicts[k], shape_list=shape_list)
+        return results

openocr/opendet/preprocess/__init__.py

@@ -0,0 +1,154 @@
+import io
+
+import cv2
+import numpy as np
+from PIL import Image
+
+from .db_resize_for_test import DetResizeForTest
+
+
+class NormalizeImage(object):
+    """normalize image such as substract mean, divide std"""
+
+    def __init__(self, scale=None, mean=None, std=None, order='chw', **kwargs):
+        if isinstance(scale, str):
+            scale = eval(scale)
+        self.scale = np.float32(scale if scale is not None else 1.0 / 255.0)
+        mean = mean if mean is not None else [0.485, 0.456, 0.406]
+        std = std if std is not None else [0.229, 0.224, 0.225]
+
+        shape = (3, 1, 1) if order == 'chw' else (1, 1, 3)
+        self.mean = np.array(mean).reshape(shape).astype('float32')
+        self.std = np.array(std).reshape(shape).astype('float32')
+
+    def __call__(self, data):
+        img = data['image']
+        from PIL import Image
+
+        if isinstance(img, Image.Image):
+            img = np.array(img)
+        assert isinstance(img,
+                          np.ndarray), "invalid input 'img' in NormalizeImage"
+        data['image'] = (img.astype('float32') * self.scale -
+                         self.mean) / self.std
+        return data
+
+
+class ToCHWImage(object):
+    """convert hwc image to chw image"""
+
+    def __init__(self, **kwargs):
+        pass
+
+    def __call__(self, data):
+        img = data['image']
+        from PIL import Image
+
+        if isinstance(img, Image.Image):
+            img = np.array(img)
+        data['image'] = img.transpose((2, 0, 1))
+        return data
+
+
+class KeepKeys(object):
+
+    def __init__(self, keep_keys, **kwargs):
+        self.keep_keys = keep_keys
+
+    def __call__(self, data):
+        data_list = []
+        for key in self.keep_keys:
+            data_list.append(data[key])
+        return data_list
+
+
+def transform(data, ops=None):
+    """transform."""
+    if ops is None:
+        ops = []
+    for op in ops:
+        data = op(data)
+        if data is None:
+            return None
+    return data
+
+
+class DecodeImage(object):
+    """decode image."""
+
+    def __init__(self,
+                 img_mode='RGB',
+                 channel_first=False,
+                 ignore_orientation=False,
+                 **kwargs):
+        self.img_mode = img_mode
+        self.channel_first = channel_first
+        self.ignore_orientation = ignore_orientation
+
+    def __call__(self, data):
+        img = data['image']
+
+        assert type(img) is bytes and len(
+            img) > 0, "invalid input 'img' in DecodeImage"
+        img = np.frombuffer(img, dtype='uint8')
+        if self.ignore_orientation:
+            img = cv2.imdecode(
+                img, cv2.IMREAD_IGNORE_ORIENTATION | cv2.IMREAD_COLOR)
+        else:
+            img = cv2.imdecode(img, 1)
+        if img is None:
+            return None
+        if self.img_mode == 'GRAY':
+            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+        elif self.img_mode == 'RGB':
+            assert img.shape[2] == 3, 'invalid shape of image[%s]' % (
+                img.shape)
+            img = img[:, :, ::-1]
+
+        if self.channel_first:
+            img = img.transpose((2, 0, 1))
+
+        data['image'] = img
+        return data
+
+
+class DecodeImagePIL(object):
+    """decode image."""
+
+    def __init__(self, img_mode='RGB', **kwargs):
+        self.img_mode = img_mode
+
+    def __call__(self, data):
+        img = data['image']
+        assert type(img) is bytes and len(
+            img) > 0, "invalid input 'img' in DecodeImage"
+        img = data['image']
+        buf = io.BytesIO(img)
+        img = Image.open(buf).convert('RGB')
+        if self.img_mode == 'Gray':
+            img = img.convert('L')
+        elif self.img_mode == 'BGR':
+            img = np.array(img)[:, :, ::-1]  # 将图片转为numpy格式，并将最后一维通道倒序
+            img = Image.fromarray(np.uint8(img))
+        data['image'] = img
+        return data
+
+
+def create_operators(op_param_list, global_config=None):
+    """create operators based on the config.
+
+    Args:
+        params(list): a dict list, used to create some operators
+    """
+    assert isinstance(op_param_list, list), 'operator config should be a list'
+    ops = []
+    for operator in op_param_list:
+        assert isinstance(operator,
+                          dict) and len(operator) == 1, 'yaml format error'
+        op_name = list(operator)[0]
+        param = {} if operator[op_name] is None else operator[op_name]
+        if global_config is not None:
+            param.update(global_config)
+        op = eval(op_name)(**param)
+        ops.append(op)
+    return ops

openocr/opendet/preprocess/crop_resize.py

@@ -0,0 +1,121 @@
+import cv2
+
+
+def padding_image(img, size=(640, 640)):
+    """
+    Padding an image using OpenCV:
+    - If the image is smaller than the target size, pad it to 640x640.
+    - If the image is larger than the target size, split it into multiple 640x640 images and record positions.
+
+    :param image_path: Path to the input image.
+    :param output_dir: Directory to save the output images.
+    :param size: The target size for padding or splitting (default 640x640).
+    :return: List of tuples containing the coordinates of the top-left corner of each cropped 640x640 image.
+    """
+
+    img_height, img_width = img.shape[:2]
+    target_width, target_height = size
+
+    # If image is smaller than target size, pad the image to 640x640
+
+    # Calculate padding amounts (top, bottom, left, right)
+    pad_top = 0
+    pad_bottom = target_height - img_height
+    pad_left = 0
+    pad_right = target_width - img_width
+
+    # Pad the image (white padding, border type: constant)
+    padded_img = cv2.copyMakeBorder(img,
+                                    pad_top,
+                                    pad_bottom,
+                                    pad_left,
+                                    pad_right,
+                                    cv2.BORDER_CONSTANT,
+                                    value=[0, 0, 0])
+
+    # Return the padded area positions (top-left and bottom-right coordinates of the original image)
+    return padded_img
+
+
+class CropResize(object):
+
+    def __init__(self, size=(640, 640), interpolation=cv2.INTER_LINEAR):
+        self.size = size
+        self.interpolation = interpolation
+
+    def __call__(self, data):
+        """
+        Resize an image using OpenCV:
+        - If the image is smaller than the target size, pad it to 640x640.
+        - If the image is larger than the target size, split it into multiple 640x640 images and record positions.
+
+        :param image_path: Path to the input image.
+        :param output_dir: Directory to save the output images.
+        :param size: The target size for padding or splitting (default 640x640).
+        :return: List of tuples containing the coordinates of the top-left corner of each cropped 640x640 image.
+        """
+        img = data['image']
+        img_height, img_width = img.shape[:2]
+        target_width, target_height = self.size
+
+        # If image is smaller than target size, pad the image to 640x640
+        if img_width <= target_width and img_height <= target_height:
+            # Calculate padding amounts (top, bottom, left, right)
+            if img_width == target_width and img_height == target_height:
+                return [img], [[0, 0, img_width, img_height]]
+            padded_img = padding_image(img, self.size)
+
+            # Return the padded area positions (top-left and bottom-right coordinates of the original image)
+            return [padded_img], [[0, 0, img_width, img_height]]
+
+        if img_width < target_width:
+            img = cv2.copyMakeBorder(img,
+                                     0,
+                                     0,
+                                     0,
+                                     target_width - img_width,
+                                     cv2.BORDER_CONSTANT,
+                                     value=[0, 0, 0])
+
+        if img_height < target_height:
+            img = cv2.copyMakeBorder(img,
+                                     0,
+                                     target_height - img_height,
+                                     0,
+                                     0,
+                                     cv2.BORDER_CONSTANT,
+                                     value=[0, 0, 0])
+            # raise ValueError("Image dimensions must be greater than or equal to target size")
+
+        img_height, img_width = img.shape[:2]
+        # If image is larger than or equal to target size, crop it into 640x640 tiles
+        crop_positions = []
+        count = 0
+        cropped_img_list = []
+        for top in range(0, img_height - target_height // 2,
+                         target_height // 2):
+            for left in range(0, img_width - target_height // 2,
+                              target_width // 2):
+                # Calculate the bottom and right boundaries for the crop
+                right = min(left + target_width, img_width)
+                bottom = min(top + target_height, img_height)
+                if right > img_width:
+                    right = img_width
+                    left = max(0, right - target_width)
+                if bottom > img_height:
+                    bottom = img_height
+                    top = max(0, bottom - target_height)
+                # Crop the image
+                cropped_img = img[top:bottom, left:right]
+                if bottom - top < target_height or right - left < target_width:
+                    cropped_img = padding_image(cropped_img, self.size)
+
+                count += 1
+                cropped_img_list.append(cropped_img)
+
+                # Record the position of the cropped image
+                crop_positions.append([left, top, right, bottom])
+
+        # print(f"Images cropped and saved at {output_dir}.")
+
+        return cropped_img_list, crop_positions

openocr/opendet/preprocess/db_resize_for_test.py

@@ -0,0 +1,135 @@
+import math
+import sys
+import cv2
+import numpy as np
+
+
+class DetResizeForTest(object):
+
+    def __init__(self, **kwargs):
+        super(DetResizeForTest, self).__init__()
+        self.resize_type = 0
+        self.keep_ratio = False
+        if 'image_shape' in kwargs:
+            self.image_shape = kwargs['image_shape']
+            self.resize_type = 1
+            if 'keep_ratio' in kwargs:
+                self.keep_ratio = kwargs['keep_ratio']
+        elif 'limit_side_len' in kwargs:
+            self.limit_side_len = kwargs['limit_side_len']
+            self.limit_type = kwargs.get('limit_type', 'min')
+        elif 'resize_long' in kwargs:
+            self.resize_type = 2
+            self.resize_long = kwargs.get('resize_long', 960)
+        else:
+            self.limit_side_len = 736
+            self.limit_type = 'min'
+
+    def __call__(self, data):
+        img = data['image']
+        src_h, src_w, _ = img.shape
+        if sum([src_h, src_w]) < 64:
+            img = self.image_padding(img)
+
+        if self.resize_type == 0:
+            # img, shape = self.resize_image_type0(img)
+            img, [ratio_h, ratio_w] = self.resize_image_type0(img)
+        elif self.resize_type == 2:
+            img, [ratio_h, ratio_w] = self.resize_image_type2(img)
+        else:
+            # img, shape = self.resize_image_type1(img)
+            img, [ratio_h, ratio_w] = self.resize_image_type1(img)
+        data['image'] = img
+        data['shape'] = np.array([src_h, src_w, ratio_h, ratio_w])
+        return data
+
+    def image_padding(self, im, value=0):
+        h, w, c = im.shape
+        im_pad = np.zeros((max(32, h), max(32, w), c), np.uint8) + value
+        im_pad[:h, :w, :] = im
+        return im_pad
+
+    def resize_image_type1(self, img):
+        resize_h, resize_w = self.image_shape
+        ori_h, ori_w = img.shape[:2]  # (h, w, c)
+        if self.keep_ratio is True:
+            resize_w = ori_w * resize_h / ori_h
+            N = math.ceil(resize_w / 32)
+            resize_w = N * 32
+        ratio_h = float(resize_h) / ori_h
+        ratio_w = float(resize_w) / ori_w
+        img = cv2.resize(img, (int(resize_w), int(resize_h)))
+        # return img, np.array([ori_h, ori_w])
+        return img, [ratio_h, ratio_w]
+
+    def resize_image_type0(self, img):
+        """
+        resize image to a size multiple of 32 which is required by the network
+        args:
+            img(array): array with shape [h, w, c]
+        return(tuple):
+            img, (ratio_h, ratio_w)
+        """
+        limit_side_len = self.limit_side_len
+        h, w, c = img.shape
+
+        # limit the max side
+        if self.limit_type == 'max':
+            if max(h, w) > limit_side_len:
+                if h > w:
+                    ratio = float(limit_side_len) / h
+                else:
+                    ratio = float(limit_side_len) / w
+            else:
+                ratio = 1.0
+        elif self.limit_type == 'min':
+            if min(h, w) < limit_side_len:
+                if h < w:
+                    ratio = float(limit_side_len) / h
+                else:
+                    ratio = float(limit_side_len) / w
+            else:
+                ratio = 1.0
+        elif self.limit_type == 'resize_long':
+            ratio = float(limit_side_len) / max(h, w)
+        else:
+            raise Exception('not support limit type, image ')
+        resize_h = int(h * ratio)
+        resize_w = int(w * ratio)
+
+        resize_h = max(int(round(resize_h / 32) * 32), 32)
+        resize_w = max(int(round(resize_w / 32) * 32), 32)
+
+        try:
+            if int(resize_w) <= 0 or int(resize_h) <= 0:
+                return None, (None, None)
+            img = cv2.resize(img, (int(resize_w), int(resize_h)))
+        except:
+            print(img.shape, resize_w, resize_h)
+            sys.exit(0)
+        ratio_h = resize_h / float(h)
+        ratio_w = resize_w / float(w)
+        return img, [ratio_h, ratio_w]
+
+    def resize_image_type2(self, img):
+        h, w, _ = img.shape
+
+        resize_w = w
+        resize_h = h
+
+        if resize_h > resize_w:
+            ratio = float(self.resize_long) / resize_h
+        else:
+            ratio = float(self.resize_long) / 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
+        img = cv2.resize(img, (int(resize_w), int(resize_h)))
+        ratio_h = resize_h / float(h)
+        ratio_w = resize_w / float(w)
+
+        return img, [ratio_h, ratio_w]