deepdoc-lib 0.2.0__py3-none-any.whl

deepdoc/vision/recognizer.py

@@ -0,0 +1,451 @@
+#
+#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+import logging
+import math
+import numpy as np
+import cv2
+from functools import cmp_to_key
+
+
+from ..common.model_store import resolve_vision_model_dir
+from .operators import *  # noqa: F403
+from .operators import preprocess
+from . import operators
+from .ocr import load_model
+
+class Recognizer:
+    def __init__(
+        self,
+        label_list,
+        task_name,
+        model_dir=None,
+        model_home: str | None = None,
+        model_provider: str | None = None,
+        offline: bool | None = None,
+    ):
+        if not model_dir:
+            model_dir = resolve_vision_model_dir(
+                model_home=model_home,
+                provider=model_provider,
+                offline=offline,
+            )
+        self.ort_sess, self.run_options = load_model(model_dir, task_name)
+        self.input_names = [node.name for node in self.ort_sess.get_inputs()]
+        self.output_names = [node.name for node in self.ort_sess.get_outputs()]
+        self.input_shape = self.ort_sess.get_inputs()[0].shape[2:4]
+        self.label_list = label_list
+
+    @staticmethod
+    def sort_Y_firstly(arr, threshold):
+        def cmp(c1, c2):
+            diff = c1["top"] - c2["top"]
+            if abs(diff) < threshold:
+                diff = c1["x0"] - c2["x0"]
+            return diff
+        arr = sorted(arr, key=cmp_to_key(cmp))
+        return arr
+
+    @staticmethod
+    def sort_X_firstly(arr, threshold):
+        def cmp(c1, c2):
+            diff = c1["x0"] - c2["x0"]
+            if abs(diff) < threshold:
+                diff = c1["top"] - c2["top"]
+            return diff
+        arr = sorted(arr, key=cmp_to_key(cmp))
+        return arr
+
+    @staticmethod
+    def sort_C_firstly(arr, thr=0):
+        # sort using y1 first and then x1
+        # sorted(arr, key=lambda r: (r["x0"], r["top"]))
+        arr = Recognizer.sort_X_firstly(arr, thr)
+        for i in range(len(arr) - 1):
+            for j in range(i, -1, -1):
+                # restore the order using th
+                if "C" not in arr[j] or "C" not in arr[j + 1]:
+                    continue
+                if arr[j + 1]["C"] < arr[j]["C"] \
+                        or (
+                        arr[j + 1]["C"] == arr[j]["C"]
+                        and arr[j + 1]["top"] < arr[j]["top"]
+                ):
+                    tmp = arr[j]
+                    arr[j] = arr[j + 1]
+                    arr[j + 1] = tmp
+        return arr
+
+    @staticmethod
+    def sort_R_firstly(arr, thr=0):
+        # sort using y1 first and then x1
+        # sorted(arr, key=lambda r: (r["top"], r["x0"]))
+        arr = Recognizer.sort_Y_firstly(arr, thr)
+        for i in range(len(arr) - 1):
+            for j in range(i, -1, -1):
+                if "R" not in arr[j] or "R" not in arr[j + 1]:
+                    continue
+                if arr[j + 1]["R"] < arr[j]["R"] \
+                        or (
+                        arr[j + 1]["R"] == arr[j]["R"]
+                        and arr[j + 1]["x0"] < arr[j]["x0"]
+                ):
+                    tmp = arr[j]
+                    arr[j] = arr[j + 1]
+                    arr[j + 1] = tmp
+        return arr
+
+    @staticmethod
+    def overlapped_area(a, b, ratio=True):
+        tp, btm, x0, x1 = a["top"], a["bottom"], a["x0"], a["x1"]
+        if b["x0"] > x1 or b["x1"] < x0:
+            return 0
+        if b["bottom"] < tp or b["top"] > btm:
+            return 0
+        x0_ = max(b["x0"], x0)
+        x1_ = min(b["x1"], x1)
+        assert x0_ <= x1_, "Bbox mismatch! T:{},B:{},X0:{},X1:{} ==> {}".format(
+            tp, btm, x0, x1, b)
+        tp_ = max(b["top"], tp)
+        btm_ = min(b["bottom"], btm)
+        assert tp_ <= btm_, "Bbox mismatch! T:{},B:{},X0:{},X1:{} => {}".format(
+            tp, btm, x0, x1, b)
+        ov = (btm_ - tp_) * (x1_ - x0_) if x1 - \
+                                           x0 != 0 and btm - tp != 0 else 0
+        if ov > 0 and ratio:
+            ov /= (x1 - x0) * (btm - tp)
+        return ov
+
+    @staticmethod
+    def layouts_cleanup(boxes, layouts, far=2, thr=0.7):
+        def not_overlapped(a, b):
+            return any([a["x1"] < b["x0"],
+                        a["x0"] > b["x1"],
+                        a["bottom"] < b["top"],
+                        a["top"] > b["bottom"]])
+
+        i = 0
+        while i + 1 < len(layouts):
+            j = i + 1
+            while j < min(i + far, len(layouts)) \
+                    and (layouts[i].get("type", "") != layouts[j].get("type", "")
+                         or not_overlapped(layouts[i], layouts[j])):
+                j += 1
+            if j >= min(i + far, len(layouts)):
+                i += 1
+                continue
+            if Recognizer.overlapped_area(layouts[i], layouts[j]) < thr \
+                    and Recognizer.overlapped_area(layouts[j], layouts[i]) < thr:
+                i += 1
+                continue
+
+            if layouts[i].get("score") and layouts[j].get("score"):
+                if layouts[i]["score"] > layouts[j]["score"]:
+                    layouts.pop(j)
+                else:
+                    layouts.pop(i)
+                continue
+
+            area_i, area_i_1 = 0, 0
+            for b in boxes:
+                if not not_overlapped(b, layouts[i]):
+                    area_i += Recognizer.overlapped_area(b, layouts[i], False)
+                if not not_overlapped(b, layouts[j]):
+                    area_i_1 += Recognizer.overlapped_area(b, layouts[j], False)
+
+            if area_i > area_i_1:
+                layouts.pop(j)
+            else:
+                layouts.pop(i)
+
+        return layouts
+
+    def create_inputs(self, imgs, im_info):
+        """generate input for different model type
+        Args:
+            imgs (list(numpy)): list of images (np.ndarray)
+            im_info (list(dict)): list of image info
+        Returns:
+            inputs (dict): input of model
+        """
+        inputs = {}
+
+        im_shape = []
+        scale_factor = []
+        if len(imgs) == 1:
+            inputs['image'] = np.array((imgs[0],)).astype('float32')
+            inputs['im_shape'] = np.array(
+                (im_info[0]['im_shape'],)).astype('float32')
+            inputs['scale_factor'] = np.array(
+                (im_info[0]['scale_factor'],)).astype('float32')
+            return inputs
+        
+        im_shape = np.array([info['im_shape'] for info in im_info], dtype='float32')
+        scale_factor = np.array([info['scale_factor'] for info in im_info], dtype='float32')
+
+        inputs['im_shape'] = np.concatenate(im_shape, axis=0)
+        inputs['scale_factor'] = np.concatenate(scale_factor, axis=0)
+
+        imgs_shape = [[e.shape[1], e.shape[2]] for e in imgs]
+        max_shape_h = max([e[0] for e in imgs_shape])
+        max_shape_w = max([e[1] for e in imgs_shape])
+        padding_imgs = []
+        for img in imgs:
+            im_c, im_h, im_w = img.shape[:]
+            padding_im = np.zeros(
+                (im_c, max_shape_h, max_shape_w), dtype=np.float32)
+            padding_im[:, :im_h, :im_w] = img
+            padding_imgs.append(padding_im)
+        inputs['image'] = np.stack(padding_imgs, axis=0)
+        return inputs
+
+    @staticmethod
+    def find_overlapped(box, boxes_sorted_by_y, naive=False):
+        if not boxes_sorted_by_y:
+            return
+        bxs = boxes_sorted_by_y
+        s, e, ii = 0, len(bxs), 0
+        while s < e and not naive:
+            ii = (e + s) // 2
+            pv = bxs[ii]
+            if box["bottom"] < pv["top"]:
+                e = ii
+                continue
+            if box["top"] > pv["bottom"]:
+                s = ii + 1
+                continue
+            break
+        while s < ii:
+            if box["top"] > bxs[s]["bottom"]:
+                s += 1
+            break
+        while e - 1 > ii:
+            if box["bottom"] < bxs[e - 1]["top"]:
+                e -= 1
+            break
+
+        max_overlapped_i, max_overlapped = None, 0
+        for i in range(s, e):
+            ov = Recognizer.overlapped_area(bxs[i], box)
+            if ov <= max_overlapped:
+                continue
+            max_overlapped_i = i
+            max_overlapped = ov
+
+        return max_overlapped_i
+
+    @staticmethod
+    def find_horizontally_tightest_fit(box, boxes):
+        if not boxes:
+            return
+        min_dis, min_i = 1000000, None
+        for i,b in enumerate(boxes):
+            if box.get("layoutno", "0") != b.get("layoutno", "0"):
+                continue
+            dis = min(abs(box["x0"] - b["x0"]), abs(box["x1"] - b["x1"]), abs(box["x0"]+box["x1"] - b["x1"] - b["x0"])/2)
+            if dis < min_dis:
+                min_i = i
+                min_dis = dis
+        return min_i
+
+    @staticmethod
+    def find_overlapped_with_threshold(box, boxes, thr=0.3):
+        if not boxes:
+            return
+        max_overlapped_i, max_overlapped, _max_overlapped = None, thr, 0
+        s, e = 0, len(boxes)
+        for i in range(s, e):
+            ov = Recognizer.overlapped_area(box, boxes[i])
+            _ov = Recognizer.overlapped_area(boxes[i], box)
+            if (ov, _ov) < (max_overlapped, _max_overlapped):
+                continue
+            max_overlapped_i = i
+            max_overlapped = ov
+            _max_overlapped = _ov
+
+        return max_overlapped_i
+
+    def preprocess(self, image_list):
+        inputs = []
+        if "scale_factor" in self.input_names:
+            preprocess_ops = []
+            for op_info in [
+                {'interp': 2, 'keep_ratio': False, 'target_size': [800, 608], 'type': 'LinearResize'},
+                {'is_scale': True, 'mean': [0.485, 0.456, 0.406], 'std': [0.229, 0.224, 0.225], 'type': 'StandardizeImage'},
+                {'type': 'Permute'},
+                {'stride': 32, 'type': 'PadStride'}
+            ]:
+                new_op_info = op_info.copy()
+                op_type = new_op_info.pop('type')
+                preprocess_ops.append(getattr(operators, op_type)(**new_op_info))
+
+            for im_path in image_list:
+                im, im_info = preprocess(im_path, preprocess_ops)
+                inputs.append({"image": np.array((im,)).astype('float32'),
+                               "scale_factor": np.array((im_info["scale_factor"],)).astype('float32')})
+        else:
+            hh, ww = self.input_shape
+            for img in image_list:
+                h, w = img.shape[:2]
+                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+                img = cv2.resize(np.array(img).astype('float32'), (ww, hh))
+                # Scale input pixel values to 0 to 1
+                img /= 255.0
+                img = img.transpose(2, 0, 1)
+                img = img[np.newaxis, :, :, :].astype(np.float32)
+                inputs.append({self.input_names[0]: img, "scale_factor": [w/ww, h/hh]})
+        return inputs
+
+    def postprocess(self, boxes, inputs, thr):
+        if "scale_factor" in self.input_names:
+            bb = []
+            for b in boxes:
+                clsid, bbox, score = int(b[0]), b[2:], b[1]
+                if score < thr:
+                    continue
+                if clsid >= len(self.label_list):
+                    continue
+                bb.append({
+                    "type": self.label_list[clsid].lower(),
+                    "bbox": [float(t) for t in bbox.tolist()],
+                    "score": float(score)
+                })
+            return bb
+
+        def xywh2xyxy(x):
+            # [x, y, w, h] to [x1, y1, x2, y2]
+            y = np.copy(x)
+            y[:, 0] = x[:, 0] - x[:, 2] / 2
+            y[:, 1] = x[:, 1] - x[:, 3] / 2
+            y[:, 2] = x[:, 0] + x[:, 2] / 2
+            y[:, 3] = x[:, 1] + x[:, 3] / 2
+            return y
+
+        def compute_iou(box, boxes):
+            # Compute xmin, ymin, xmax, ymax for both boxes
+            xmin = np.maximum(box[0], boxes[:, 0])
+            ymin = np.maximum(box[1], boxes[:, 1])
+            xmax = np.minimum(box[2], boxes[:, 2])
+            ymax = np.minimum(box[3], boxes[:, 3])
+
+            # Compute intersection area
+            intersection_area = np.maximum(0, xmax - xmin) * np.maximum(0, ymax - ymin)
+
+            # Compute union area
+            box_area = (box[2] - box[0]) * (box[3] - box[1])
+            boxes_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
+            union_area = box_area + boxes_area - intersection_area
+
+            # Compute IoU
+            iou = intersection_area / union_area
+
+            return iou
+
+        def iou_filter(boxes, scores, iou_threshold):
+            sorted_indices = np.argsort(scores)[::-1]
+
+            keep_boxes = []
+            while sorted_indices.size > 0:
+                # Pick the last box
+                box_id = sorted_indices[0]
+                keep_boxes.append(box_id)
+
+                # Compute IoU of the picked box with the rest
+                ious = compute_iou(boxes[box_id, :], boxes[sorted_indices[1:], :])
+
+                # Remove boxes with IoU over the threshold
+                keep_indices = np.where(ious < iou_threshold)[0]
+
+                # print(keep_indices.shape, sorted_indices.shape)
+                sorted_indices = sorted_indices[keep_indices + 1]
+
+            return keep_boxes
+
+        boxes = np.squeeze(boxes).T
+        # Filter out object confidence scores below threshold
+        scores = np.max(boxes[:, 4:], axis=1)
+        boxes = boxes[scores > thr, :]
+        scores = scores[scores > thr]
+        if len(boxes) == 0:
+            return []
+
+        # Get the class with the highest confidence
+        class_ids = np.argmax(boxes[:, 4:], axis=1)
+        boxes = boxes[:, :4]
+        input_shape = np.array([inputs["scale_factor"][0], inputs["scale_factor"][1], inputs["scale_factor"][0], inputs["scale_factor"][1]])
+        boxes = np.multiply(boxes, input_shape, dtype=np.float32)
+        boxes = xywh2xyxy(boxes)
+
+        unique_class_ids = np.unique(class_ids)
+        indices = []
+        for class_id in unique_class_ids:
+            class_indices = np.where(class_ids == class_id)[0]
+            class_boxes = boxes[class_indices, :]
+            class_scores = scores[class_indices]
+            class_keep_boxes = iou_filter(class_boxes, class_scores, 0.2)
+            indices.extend(class_indices[class_keep_boxes])
+
+        return [{
+            "type": self.label_list[class_ids[i]].lower(),
+            "bbox": [float(t) for t in boxes[i].tolist()],
+            "score": float(scores[i])
+        } for i in indices]
+
+    def close(self):
+        # NOTE: `__del__` can run during interpreter shutdown when module
+        # globals (including `logging`/`gc`) may already be cleared to None.
+        try:
+            import logging as _logging
+            _logging.info("Close recognizer.")
+        except Exception:
+            pass
+        if hasattr(self, "ort_sess"):
+            del self.ort_sess
+        try:
+            import gc as _gc
+            _gc.collect()
+        except Exception:
+            pass
+
+    def __call__(self, image_list, thr=0.7, batch_size=16):
+        res = []
+        images = []
+        for i in range(len(image_list)):
+            if not isinstance(image_list[i], np.ndarray):
+                images.append(np.array(image_list[i]))
+            else:
+                images.append(image_list[i])
+
+        batch_loop_cnt = math.ceil(float(len(images)) / batch_size)
+        for i in range(batch_loop_cnt):
+            start_index = i * batch_size
+            end_index = min((i + 1) * batch_size, len(images))
+            batch_image_list = images[start_index:end_index]
+            inputs = self.preprocess(batch_image_list)
+            logging.debug("preprocess")
+            for ins in inputs:
+                bb = self.postprocess(self.ort_sess.run(None, {k:v for k,v in ins.items() if k in self.input_names}, self.run_options)[0], ins, thr)
+                res.append(bb)
+
+        #seeit.save_results(image_list, res, self.label_list, threshold=thr)
+
+        return res
+
+    def __del__(self):
+        try:
+            self.close()
+        except Exception:
+            # Destructors must never raise.
+            pass

deepdoc/vision/seeit.py

@@ -0,0 +1,87 @@
+#
+#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+import logging
+import os
+import PIL
+from PIL import ImageDraw
+
+
+def save_results(image_list, results, labels, output_dir='output/', threshold=0.5):
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+    for idx, im in enumerate(image_list):
+        im = draw_box(im, results[idx], labels, threshold=threshold)
+
+        out_path = os.path.join(output_dir, f"{idx}.jpg")
+        im.save(out_path, quality=95)
+        logging.debug("save result to: " + out_path)
+
+
+def draw_box(im, result, labels, threshold=0.5):
+    draw_thickness = min(im.size) // 320
+    draw = ImageDraw.Draw(im)
+    color_list = get_color_map_list(len(labels))
+    clsid2color = {n.lower():color_list[i] for i,n in enumerate(labels)}
+    result = [r for r in result if r["score"] >= threshold]
+
+    for dt in result:
+        color = tuple(clsid2color[dt["type"]])
+        xmin, ymin, xmax, ymax = dt["bbox"]
+        draw.line(
+            [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin),
+             (xmin, ymin)],
+            width=draw_thickness,
+            fill=color)
+
+        # draw label
+        text = "{} {:.4f}".format(dt["type"], dt["score"])
+        tw, th = imagedraw_textsize_c(draw, text)
+        draw.rectangle(
+            [(xmin + 1, ymin - th), (xmin + tw + 1, ymin)], fill=color)
+        draw.text((xmin + 1, ymin - th), text, fill=(255, 255, 255))
+    return im
+
+
+def get_color_map_list(num_classes):
+    """
+    Args:
+        num_classes (int): number of class
+    Returns:
+        color_map (list): RGB color list
+    """
+    color_map = num_classes * [0, 0, 0]
+    for i in range(0, num_classes):
+        j = 0
+        lab = i
+        while lab:
+            color_map[i * 3] |= (((lab >> 0) & 1) << (7 - j))
+            color_map[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j))
+            color_map[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j))
+            j += 1
+            lab >>= 3
+    color_map = [color_map[i:i + 3] for i in range(0, len(color_map), 3)]
+    return color_map
+
+
+def imagedraw_textsize_c(draw, text):
+    if int(PIL.__version__.split('.')[0]) < 10:
+        tw, th = draw.textsize(text)
+    else:
+        left, top, right, bottom = draw.textbbox((0, 0), text)
+        tw, th = right - left, bottom - top
+
+    return tw, th

deepdoc/vision/t_ocr.py

@@ -0,0 +1,101 @@
+#
+#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+import asyncio
+import logging
+import os
+import sys
+sys.path.insert(
+    0,
+    os.path.abspath(
+        os.path.join(
+            os.path.dirname(
+                os.path.abspath(__file__)),
+            '../../')))
+
+from deepdoc.vision.seeit import draw_box
+from deepdoc.vision import OCR, init_in_out
+import argparse
+import numpy as np
+
+# os.environ['CUDA_VISIBLE_DEVICES'] = '0,2' #2 gpus, uncontinuous
+os.environ['CUDA_VISIBLE_DEVICES'] = '0' #1 gpu
+# os.environ['CUDA_VISIBLE_DEVICES'] = '' #cpu
+
+
+def main(args):
+    import torch.cuda
+
+    cuda_devices = torch.cuda.device_count()
+    limiter = [asyncio.Semaphore(1) for _ in range(cuda_devices)] if cuda_devices > 1 else None
+    ocr = OCR()
+    images, outputs = init_in_out(args)
+
+    def __ocr(i, id, img):
+        print("Task {} start".format(i))
+        bxs = ocr(np.array(img), id)
+        bxs = [(line[0], line[1][0]) for line in bxs]
+        bxs = [{
+            "text": t,
+            "bbox": [b[0][0], b[0][1], b[1][0], b[-1][1]],
+            "type": "ocr",
+            "score": 1} for b, t in bxs if b[0][0] <= b[1][0] and b[0][1] <= b[-1][1]]
+        img = draw_box(images[i], bxs, ["ocr"], 1.)
+        img.save(outputs[i], quality=95)
+        with open(outputs[i] + ".txt", "w+", encoding='utf-8') as f:
+            f.write("\n".join([o["text"] for o in bxs]))
+
+        print("Task {} done".format(i))
+
+    async def __ocr_thread(i, id, img, limiter = None):
+        if limiter:
+            async with limiter:
+                print(f"Task {i} use device {id}")
+                await asyncio.to_thread(__ocr, i, id, img)
+        else:
+            await asyncio.to_thread(__ocr, i, id, img)
+
+
+    async def __ocr_launcher():
+        tasks = []
+        for i, img in enumerate(images):
+            dev_id = i % cuda_devices if cuda_devices > 1 else 0
+            semaphore = limiter[dev_id] if limiter else None
+            tasks.append(asyncio.create_task(__ocr_thread(i, dev_id, img, semaphore)))
+
+        try:
+            await asyncio.gather(*tasks, return_exceptions=False)
+        except Exception as e:
+            logging.error("OCR tasks failed: {}".format(e))
+            for t in tasks:
+                t.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+            raise
+
+    asyncio.run(__ocr_launcher())
+
+    print("OCR tasks are all done")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--inputs',
+                        help="Directory where to store images or PDFs, or a file path to a single image or PDF",
+                        required=True)
+    parser.add_argument('--output_dir', help="Directory where to store the output images. Default: './ocr_outputs'",
+                        default="./ocr_outputs")
+    args = parser.parse_args()
+    main(args)