tfjs-text-detection-ctpn/main.js at master · BadMachine/tfjs-text-detection-ctpn · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
import * as tf from '@tensorflow/tfjs-node-gpu';
import * as fs from 'fs';
import * as PImage from 'pureimage';
import { proposal_layer } from './rpn_msr/proposal_layer_tf.js';
import { resize_im } from "./utils/resize.js";
import { TextDetector } from './text_connector/detectors.js';
import { RGB2BGR } from './utils/RGB2BGR.js';
import { _get_blobs } from "./fast_rcnn/inference_blob.js";

export default class CTPN{
    constructor(config) {
        this.model = tf.loadGraphModel('https://cdn.jsdelivr.net/gh/BadMachine/tfjs-text-detection-ctpn/ctpn_web/model.json'); //tf.loadGraphModel('file://./ctpn_web/model.json');
        this.cfg = config;
    }

    async predict(image_path){
        tf.engine().startScope()
        const image = RGB2BGR(tf.node.decodeImage(fs.readFileSync(image_path)).cast('float32'));
        const [img, scale] = resize_im(image, 600, 1200);

        const [blobs, im_scales] = _get_blobs(img, null, this.cfg);
        if (this.cfg.HAS_RPN){
            const im_blob = blobs.data;
            blobs.im_info = tf.tensor( [[im_blob.shape[1], im_blob.shape[2], im_scales[0]]]);
        }
        const model = await this.model;

        const raw = await model.executeAsync(img.expandDims());
        const [cls_prob, box_pred] = raw;

        let [scores, proposals, bbox_deltas] = await proposal_layer(this.cfg, cls_prob, box_pred, blobs.im_info,'TEST');
        const boxes = tf.div(proposals, im_scales[0]);
        const textDetector = new TextDetector(this.cfg);
        const _boxes = await textDetector.detect(boxes, scores.reshape([scores.shape[0],1]), img.shape.slice(0,2));
        return [_boxes, scale];
        tf.engine().endScope()
    }

    async draw(image_name, writeTo, _boxes, scale, color){
        const image = await PImage.decodeJPEGFromStream(fs.createReadStream(`./${image_name}`));
        const boxes = _boxes.arraySync();

        for(let box of boxes){
            const ctx = image.getContext('2d');
            ctx.beginPath();
            ctx.strokeStyle = color;
            ctx.lineWidth = 4;
            ctx.moveTo(box[0]/ scale, box[1]/ scale);
            ctx.lineTo(box[2] / scale, box[3] / scale);

            ctx.lineTo(box[0] / scale, box[1] / scale);
            ctx.lineTo(box[4] / scale, box[5] / scale);

            ctx.lineTo(box[6] / scale, box[7] / scale);
            ctx.lineTo(box[2] / scale, box[3] / scale);

            ctx.stroke();
            ctx.closePath();

        }

        PImage.encodeJPEGToStream(image,fs.createWriteStream(writeTo), 50).then(() => {
            console.log("done writing");
        });

    }

}

// (async ()=> {
//     const cfg = {
//         NMS_FUNCTION: 'TF',
//         ANCHOR_SCALES: [16],
//         PIXEL_MEANS: tf.tensor([[[102.9801, 115.9465, 122.7717]]]),
//         SCALES: [600,] ,
//         MAX_SIZE:  1000,
//         HAS_RPN: true,
//         DETECT_MODE: 'O',
//         pre_nms_topN: 12000,
//         post_nms_topN: 2000,
//         nms_thresh:0.7,
//         min_size: 8,
//     };
//     const ctpn = new CTPN(cfg);
//     const image = './test/007.jpg';
//     const predicted = await ctpn.predict(image);
//     console.log(predicted);
//     ctpn.draw(image,'res.jpg',...predicted, 'red')
// })();