netron/source/message.js

// Experimental

var message = message || {};

message.ModelFactory = class {

    match(context) {
        const obj = context.open('json');
        if (obj && obj.signature && obj.signature.startsWith('netron:')) {
            return obj.signature;
        }
        return '';
    }

    open(context) {
        return Promise.resolve().then(() => {
            const obj = context.open('json');
            return new message.Model(obj);
        });
    }
};

message.Model = class {

    constructor(data) {
        this._format = data.format || '';
        this._producer = data.producer || '';
        this._version = data.version || '';
        this._description = data.description || '';
        this._metadata = (data.metadata || []).map((entry) => { return { name: entry.name, value: entry.value }; });
        this._graphs = (data.graphs || []).map((graph) => new message.Graph(graph));
    }

    get format() {
        return this._format;
    }

    get producer() {
        return this._producer;
    }

    get version() {
        return this._version;
    }

    get description() {
        return this._description;
    }

    get metadata() {
        return this._metadata;
    }

    get graphs() {
        return this._graphs;
    }
};

message.Graph = class {

    constructor(data) {
        this._inputs = [];
        this._outputs = [];
        this._nodes = [];
        const args = (data.arguments || []).map((argument) => new message.Argument(argument));
        for (const parameter of data.inputs || []) {
            parameter.arguments = parameter.arguments.map((index) => args[index]).filter((argument) => !argument.initializer);
            if (parameter.arguments.filter((argument) => !argument.initializer).length > 0) {
                this._inputs.push(new message.Parameter(parameter));
            }
        }
        for (const parameter of data.outputs || []) {
            parameter.arguments = parameter.arguments.map((index) => args[index]);
            if (parameter.arguments.filter((argument) => !argument.initializer).length > 0) {
                this._inputs.push(new message.Parameter(parameter));
            }
        }
        for (const node of data.nodes || []) {
            for (const parameter of node.inputs || []) {
                parameter.arguments = parameter.arguments.map((index) => args[index]);
            }
            for (const parameter of node.outputs || []) {
                parameter.arguments = parameter.arguments.map((index) => args[index]);
            }
            this._nodes.push(new message.Node(node));
        }
    }

    get inputs() {
        return this._inputs;
    }

    get outputs() {
        return this._outputs;
    }

    get nodes() {
        return this._nodes;
    }
};

message.Parameter = class {

    constructor(data) {
        this._name = data.name || '';
        this._arguments = (data.arguments || []);
    }

    get name() {
        return this._name;
    }

    get arguments() {
        return this._arguments;
    }

    get visible() {
        return true;
    }
};

message.Argument = class {

    constructor(data) {
        this._name= data.name || '';
        this._type = data.type ? new message.TensorType(data.type) : null;
        this._initializer = data.initializer ? new message.Tensor(data.initializer) : null;
    }

    get name() {
        return this._name;
    }

    get type() {
        if (this._initializer && this._initializer.type) {
            return this._initializer.type;
        }
        return this._type;
    }

    get initializer() {
        return this._initializer;
    }
};

message.Node = class {

    constructor(data) {
        this._type = { name: data.type.name, category: data.type.category };
        this._name = data.name;
        this._inputs = (data.inputs || []).map((input) => new message.Parameter(input));
        this._outputs = (data.outputs || []).map((output) => new message.Parameter(output));
        this._attributes = (data.attributes || []).map((attribute) => new message.Attribute(attribute));
    }

    get type() {
        return this._type;
    }

    get name() {
        return this._name;
    }

    get inputs() {
        return this._inputs;
    }

    get outputs() {
        return this._outputs;
    }

    get attributes() {
        return this._attributes;
    }
};

message.Attribute = class {

    constructor(data) {
        this._type = data.type || '';
        this._name = data.name;
        this._value = data.value;
    }

    get name() {
        return this._name;
    }

    get value() {
        return this._value;
    }

    get type() {
        return this._type;
    }
};

message.TensorType = class {

    constructor(data) {
        this._dataType = data.dataType;
        this._shape = new message.TensorShape(data.shape);
    }

    get dataType() {
        return this._dataType;
    }

    get shape() {
        return this._shape;
    }

    toString() {
        return this._dataType + this._shape.toString();
    }
};

message.TensorShape = class {

    constructor(data) {
        this._dimensions = data.dimensions;
    }

    get dimensions() {
        return this._dimensions;
    }

    toString() {
        return '[' + this._dimensions.toString() + ']';
    }
};

message.Tensor = class {

    constructor() {
    }
};

message.Error = class extends Error {
    constructor(message) {
        super(message);
        this.name = 'Message Error';
    }
};

if (typeof module !== 'undefined' && typeof module.exports === 'object') {
    module.exports.ModelFactory = message.ModelFactory;
}