onnxruntime

1
// Copyright (c) Microsoft Corporation. All rights reserved.
2
// Licensed under the MIT License.
3

4
import { Logger } from '../../instrument';
5

6
export declare namespace Encoder {
7
  export interface DataTypeMap {
8
    float: Float32Array;
9
    byte: Uint8Array;
10
    int: Uint32Array;
11
  }
12
  export type DataType = keyof DataTypeMap;
13
  type DataArrayType = DataTypeMap[DataType];
14
}
15

16
/* eslint-disable @typescript-eslint/naming-convention */
17
export const enum EncoderUsage {
18
  Default = 0,
19
  UploadOnly,
20
  Download4BytesAsFloat32,
21
}
22
/* eslint-enable @typescript-eslint/naming-convention */
23

24
/**
25
 * Abstraction for mapping data types to texture texlets
26
 * Encoding means how a Float32 is mapped to 1 or 4 channels for each texlet
27
 * Decoding means how a texlet's channels are mapped to a resulting Float32
28
 */
29
export interface DataEncoder {
30
  internalFormat: number;
31
  format: number;
32
  textureType: number;
33
  channelSize: number;
34
  encode(src: Encoder.DataArrayType, textureSize: number): Encoder.DataArrayType;
35
  allocate(size: number): Encoder.DataArrayType;
36
  decode(buffer: Encoder.DataArrayType, dataSize: number): Encoder.DataArrayType;
37
}
38
/**
39
 * WebGL2 data encoder
40
 * Uses R32F as the format for texlet
41
 */
42
export class RedFloat32DataEncoder implements DataEncoder {
43
  internalFormat: number;
44
  format: number;
45
  textureType: number;
46
  channelSize: number;
47
  constructor(gl: WebGL2RenderingContext, channels = 1) {
48
    if (channels === 1) {
49
      this.internalFormat = gl.R32F;
50
      this.format = gl.RED;
51
      this.textureType = gl.FLOAT;
52
      this.channelSize = channels;
53
    } else if (channels === 4) {
54
      this.internalFormat = gl.RGBA32F;
55
      this.format = gl.RGBA;
56
      this.textureType = gl.FLOAT;
57
      this.channelSize = channels;
58
    } else {
59
      throw new Error(`Invalid number of channels: ${channels}`);
60
    }
61
  }
62
  encode(src: Encoder.DataArrayType, textureSize: number): Encoder.DataArrayType {
63
    let result: Float32Array;
64
    let source: Float32Array;
65
    if (src.constructor !== Float32Array) {
66
      Logger.warning('Encoder', 'data was not of type Float32; creating new Float32Array');
67
      source = new Float32Array(src);
68
    }
69
    if (textureSize * this.channelSize > src.length) {
70
      Logger.warning('Encoder', 'Source data too small. Allocating larger array');
71
      source = src as Float32Array;
72
      result = this.allocate(textureSize * this.channelSize) as Float32Array;
73
      source.forEach((v, i) => (result[i] = v));
74
    } else {
75
      source = src as Float32Array;
76
      result = source;
77
    }
78
    return result;
79
  }
80
  allocate(size: number): Encoder.DataArrayType {
81
    return new Float32Array(size * 4);
82
  }
83
  decode(buffer: Encoder.DataArrayType, dataSize: number): Float32Array {
84
    if (this.channelSize === 1) {
85
      const filteredData = (buffer as Float32Array).filter((_value, index) => index % 4 === 0).subarray(0, dataSize);
86
      return filteredData;
87
    }
88
    return buffer.subarray(0, dataSize) as Float32Array;
89
  }
90
}
91
/**
92
 * Data encoder for WebGL 1 with support for floating point texture
93
 */
94
export class RGBAFloatDataEncoder implements DataEncoder {
95
  internalFormat: number;
96
  format: number;
97
  textureType: number;
98
  channelSize: number;
99
  constructor(gl: WebGLRenderingContext, channels = 1, textureType?: number) {
100
    if (channels !== 1 && channels !== 4) {
101
      throw new Error(`Invalid number of channels: ${channels}`);
102
    }
103
    this.internalFormat = gl.RGBA;
104
    this.format = gl.RGBA;
105
    this.channelSize = channels;
106
    this.textureType = textureType || gl.FLOAT;
107
  }
108
  encode(src: Float32Array, textureSize: number): Encoder.DataArrayType {
109
    let dest = src;
110
    if (this.channelSize === 1) {
111
      Logger.verbose('Encoder', 'Exploding into a larger array');
112
      dest = this.allocate(textureSize) as Float32Array;
113
      src.forEach((v, i) => (dest[i * 4] = v));
114
    }
115
    return dest;
116
  }
117
  allocate(size: number): Encoder.DataArrayType {
118
    return new Float32Array(size * 4);
119
  }
120
  decode(buffer: Encoder.DataArrayType, dataSize: number): Float32Array {
121
    if (this.channelSize === 1) {
122
      const filteredData = (buffer as Float32Array).filter((_value, index) => index % 4 === 0).subarray(0, dataSize);
123
      return filteredData;
124
    }
125
    return buffer.subarray(0, dataSize) as Float32Array;
126
  }
127
}
128

129
export class Uint8DataEncoder implements DataEncoder {
130
  internalFormat: number;
131
  format: number;
132
  textureType: number;
133
  channelSize = 4;
134
  constructor(gl: WebGLRenderingContext, channels = 1) {
135
    if (channels === 1) {
136
      this.internalFormat = gl.ALPHA;
137
      this.format = gl.ALPHA; // not tested
138
      this.textureType = gl.UNSIGNED_BYTE;
139
      this.channelSize = channels;
140
    } else if (channels === 4) {
141
      this.internalFormat = gl.RGBA;
142
      this.format = gl.RGBA;
143
      this.textureType = gl.UNSIGNED_BYTE;
144
      this.channelSize = channels;
145
    } else {
146
      throw new Error(`Invalid number of channels: ${channels}`);
147
    }
148
  }
149
  encode(src: Uint8Array, _textureSize: number): Encoder.DataArrayType {
150
    return new Uint8Array(src.buffer, src.byteOffset, src.byteLength);
151
  }
152
  allocate(size: number): Encoder.DataArrayType {
153
    return new Uint8Array(size * this.channelSize);
154
  }
155
  decode(buffer: Encoder.DataArrayType, dataSize: number): Uint8Array {
156
    if (buffer instanceof Uint8Array) {
157
      return buffer.subarray(0, dataSize);
158
    }
159
    throw new Error(`Invalid array type: ${buffer.constructor}`);
160
  }
161
}
162