FastDeploy/examples/application/js/package/packages/paddlejs-models/ocrdetection/src/dbPostprocess.ts

import clipper from 'js-clipper';
import { divide, enableBoundaryChecking, plus } from 'number-precision';
import CV from '@paddlejs-mediapipe/opencv/library/opencv_ocr';

import * as d3Polygon from 'd3-polygon';
import { BOX, POINT, POINTS } from "./type";

export default class DBPostprocess {
    private readonly thresh: number;
    private readonly box_thresh: number;
    private readonly max_candidates: number;
    private readonly unclip_ratio: number;
    private readonly min_size: number;
    private readonly pred: number[];
    private readonly segmentation: number[];
    private readonly width: number;
    private readonly height: number;

    constructor(result: number[], shape: number[], thresh=0.3, box_thresh=0.6, unclip_ratio=1.5) {
        enableBoundaryChecking(false);
        this.thresh = thresh ? thresh : 0.3;
        this.box_thresh = box_thresh ? box_thresh : 0.6;
        this.max_candidates = 1000;
        this.unclip_ratio = unclip_ratio ? unclip_ratio:1.5;
        this.min_size = 3;
        this.width = shape[0];
        this.height = shape[1];
        this.pred = result;
        this.segmentation = [];
        this.pred.forEach((item: number) => {
            this.segmentation.push(item > this.thresh ? 255 : 0);
        });
    }

    public outputBox() {
        // eslint-disable-next-line new-cap
        const src = new CV.matFromArray(960, 960, CV.CV_8UC1, this.segmentation);
        const contours = new CV.MatVector();
        const hierarchy = new CV.Mat();
        // 获取轮廓
        CV.findContours(src, contours, hierarchy, CV.RETR_LIST, CV.CHAIN_APPROX_SIMPLE);
        const num_contours = Math.min(contours.size(), this.max_candidates);
        const boxes: BOX = [];
        const scores: number[] = [];
        const arr: number[] = [];
        for (let i = 0; i < num_contours; i++) {
            const contour = contours.get(i);
            const minBox = this.get_mini_boxes(contour);
            const points = minBox.points;
            let side = minBox.side;
            if (side < this.min_size) {
                continue;
            }
            const score = this.box_score_fast(this.pred, points);
            if (this.box_thresh > score) {
                continue;
            }
            let box = this.unclip(points);
            // eslint-disable-next-line new-cap
            const boxMap = new CV.matFromArray(box.length / 2, 1, CV.CV_32SC2, box);
            const resultObj = this.get_mini_boxes(boxMap);
            box = resultObj.points;
            side = resultObj.side;
            if (side < this.min_size + 2) {
                continue;
            }
            box.forEach(item => {
                item[0] = this.clip(Math.round(item[0]), 0, this.width);
                item[1] = this.clip(Math.round(item[1]), 0, this.height);
            });
            boxes.push(box);
            scores.push(score);
            arr.push(i);
            boxMap.delete();
        }
        src.delete();
        contours.delete();
        hierarchy.delete();
        return { boxes, scores };
    }

    private get_mini_boxes(contour) {
        // 生成最小外接矩形
        const bounding_box = CV.minAreaRect(contour);
        const points: POINTS = [];
        const mat = new CV.Mat();
        // 获取矩形的四个顶点坐标
        CV.boxPoints(bounding_box, mat);
        for (let i = 0; i < mat.data32F.length; i += 2) {
            const arr: POINT = [mat.data32F[i], mat.data32F[i + 1]];
            points.push(arr);
        }
        function sortNumber(a: POINT, b: POINT) {
            return a[0] - b[0];
        }
        points.sort(sortNumber);

        let index_1: number;
        let index_2: number;
        let index_3: number;
        let index_4: number;
        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;
        }

        const box = [
            points[index_1],
            points[index_2],
            points[index_3],
            points[index_4]
        ];

        const side = Math.min(bounding_box.size.height, bounding_box.size.width);
        mat.delete();

        return { points: box, side };
    }

    private box_score_fast(bitmap: number[], _box: POINTS) {
        const h = this.height;
        const w = this.width;
        const box = JSON.parse(JSON.stringify(_box));
        const x = [] as number[];
        const y = [] as number[];
        box.forEach((item: POINT) => {
            x.push(item[0]);
            y.push(item[1]);
        });
        // clip这个函数将将数组中的元素限制在a_min, a_max之间，大于a_max的就使得它等于 a_max，小于a_min,的就使得它等于a_min。
        const xmin = this.clip(Math.floor(Math.min(...x)), 0, w - 1);
        const xmax = this.clip(Math.ceil(Math.max(...x)), 0, w - 1);
        const ymin = this.clip(Math.floor(Math.min(...y)), 0, h - 1);
        const ymax = this.clip(Math.ceil(Math.max(...y)), 0, h - 1);
        // eslint-disable-next-line new-cap
        const mask = new CV.Mat.zeros(ymax - ymin + 1, xmax - xmin + 1, CV.CV_8UC1);
        box.forEach((item: POINT) => {
            item[0] = Math.max(item[0] - xmin, 0);
            item[1] = Math.max(item[1] - ymin, 0);
        });
        const npts = 4;
        const point_data = new Uint8Array(box.flat());
        const points = CV.matFromArray(npts, 1, CV.CV_32SC2, point_data);
        const pts = new CV.MatVector();
        pts.push_back(points);
        const color = new CV.Scalar(255);
        // 多个多边形填充
        CV.fillPoly(mask, pts, color, 1);
        const sliceArr = [];
        for (let i = ymin; i < ymax + 1; i++) {
            sliceArr.push(...bitmap.slice(960 * i + xmin, 960 * i + xmax + 1) as []);
        }
        const mean = this.mean(sliceArr, mask.data);
        mask.delete();
        points.delete();
        pts.delete();
        return mean;
    }

    private clip(data: number, min: number, max: number) {
        return data < min ? min : data > max ? max : data;
    }

    private unclip(box: POINTS) {
        const unclip_ratio = this.unclip_ratio;
        const area = Math.abs(d3Polygon.polygonArea(box));
        const length = d3Polygon.polygonLength(box);
        const distance = area * unclip_ratio / length;
        const tmpArr: { X: number; Y: number; }[] = [];
        box.forEach(item => {
            const obj = {
                X: 0,
                Y: 0
            };
            obj.X = item[0];
            obj.Y = item[1];
            tmpArr.push(obj);
        });
        const offset = new clipper.ClipperOffset();
        offset.AddPath(tmpArr, clipper.JoinType.jtRound, clipper.EndType.etClosedPolygon);
        const expanded: { X: number; Y: number; }[][] = [];
        offset.Execute(expanded, distance);
        let expandedArr: POINTS = [];
        expanded[0] && expanded[0].forEach(item => {
            expandedArr.push([item.X, item.Y]);
        });
        expandedArr = [].concat(...expandedArr as []);
        return expandedArr;
    }

    private mean(data: number[], mask: number[]) {
        let sum = 0;
        let length = 0;
        for (let i = 0; i < data.length; i++) {
            if (mask[i]) {
                sum = plus(sum, data[i]);
                length++;
            }
        }
        const num = divide(sum, length);
        return num;
    }
}