@huggingface/transformers 3.0.0-alpha.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (96) hide show
  1. package/LICENSE +202 -0
  2. package/README.md +376 -0
  3. package/dist/ort-wasm-simd-threaded.jsep.wasm +0 -0
  4. package/dist/transformers.cjs +30741 -0
  5. package/dist/transformers.cjs.map +1 -0
  6. package/dist/transformers.js +33858 -0
  7. package/dist/transformers.js.map +1 -0
  8. package/dist/transformers.min.cjs +173 -0
  9. package/dist/transformers.min.cjs.map +1 -0
  10. package/dist/transformers.min.js +231 -0
  11. package/dist/transformers.min.js.map +1 -0
  12. package/package.json +92 -0
  13. package/src/backends/onnx.js +151 -0
  14. package/src/configs.js +360 -0
  15. package/src/env.js +152 -0
  16. package/src/generation/configuration_utils.js +381 -0
  17. package/src/generation/logits_process.js +716 -0
  18. package/src/generation/logits_sampler.js +204 -0
  19. package/src/generation/parameters.js +35 -0
  20. package/src/generation/stopping_criteria.js +156 -0
  21. package/src/generation/streamers.js +212 -0
  22. package/src/models/whisper/common_whisper.js +151 -0
  23. package/src/models/whisper/generation_whisper.js +89 -0
  24. package/src/models.js +7028 -0
  25. package/src/ops/registry.js +92 -0
  26. package/src/pipelines.js +3341 -0
  27. package/src/processors.js +2614 -0
  28. package/src/tokenizers.js +4395 -0
  29. package/src/transformers.js +28 -0
  30. package/src/utils/audio.js +704 -0
  31. package/src/utils/constants.js +2 -0
  32. package/src/utils/core.js +149 -0
  33. package/src/utils/data-structures.js +445 -0
  34. package/src/utils/devices.js +11 -0
  35. package/src/utils/dtypes.js +62 -0
  36. package/src/utils/generic.js +35 -0
  37. package/src/utils/hub.js +671 -0
  38. package/src/utils/image.js +745 -0
  39. package/src/utils/maths.js +1050 -0
  40. package/src/utils/tensor.js +1378 -0
  41. package/types/backends/onnx.d.ts +26 -0
  42. package/types/backends/onnx.d.ts.map +1 -0
  43. package/types/configs.d.ts +59 -0
  44. package/types/configs.d.ts.map +1 -0
  45. package/types/env.d.ts +106 -0
  46. package/types/env.d.ts.map +1 -0
  47. package/types/generation/configuration_utils.d.ts +320 -0
  48. package/types/generation/configuration_utils.d.ts.map +1 -0
  49. package/types/generation/logits_process.d.ts +354 -0
  50. package/types/generation/logits_process.d.ts.map +1 -0
  51. package/types/generation/logits_sampler.d.ts +51 -0
  52. package/types/generation/logits_sampler.d.ts.map +1 -0
  53. package/types/generation/parameters.d.ts +47 -0
  54. package/types/generation/parameters.d.ts.map +1 -0
  55. package/types/generation/stopping_criteria.d.ts +81 -0
  56. package/types/generation/stopping_criteria.d.ts.map +1 -0
  57. package/types/generation/streamers.d.ts +81 -0
  58. package/types/generation/streamers.d.ts.map +1 -0
  59. package/types/models/whisper/common_whisper.d.ts +8 -0
  60. package/types/models/whisper/common_whisper.d.ts.map +1 -0
  61. package/types/models/whisper/generation_whisper.d.ts +76 -0
  62. package/types/models/whisper/generation_whisper.d.ts.map +1 -0
  63. package/types/models.d.ts +3845 -0
  64. package/types/models.d.ts.map +1 -0
  65. package/types/ops/registry.d.ts +11 -0
  66. package/types/ops/registry.d.ts.map +1 -0
  67. package/types/pipelines.d.ts +2403 -0
  68. package/types/pipelines.d.ts.map +1 -0
  69. package/types/processors.d.ts +917 -0
  70. package/types/processors.d.ts.map +1 -0
  71. package/types/tokenizers.d.ts +999 -0
  72. package/types/tokenizers.d.ts.map +1 -0
  73. package/types/transformers.d.ts +13 -0
  74. package/types/transformers.d.ts.map +1 -0
  75. package/types/utils/audio.d.ts +130 -0
  76. package/types/utils/audio.d.ts.map +1 -0
  77. package/types/utils/constants.d.ts +2 -0
  78. package/types/utils/constants.d.ts.map +1 -0
  79. package/types/utils/core.d.ts +91 -0
  80. package/types/utils/core.d.ts.map +1 -0
  81. package/types/utils/data-structures.d.ts +236 -0
  82. package/types/utils/data-structures.d.ts.map +1 -0
  83. package/types/utils/devices.d.ts +8 -0
  84. package/types/utils/devices.d.ts.map +1 -0
  85. package/types/utils/dtypes.d.ts +22 -0
  86. package/types/utils/dtypes.d.ts.map +1 -0
  87. package/types/utils/generic.d.ts +11 -0
  88. package/types/utils/generic.d.ts.map +1 -0
  89. package/types/utils/hub.d.ts +191 -0
  90. package/types/utils/hub.d.ts.map +1 -0
  91. package/types/utils/image.d.ts +119 -0
  92. package/types/utils/image.d.ts.map +1 -0
  93. package/types/utils/maths.d.ts +280 -0
  94. package/types/utils/maths.d.ts.map +1 -0
  95. package/types/utils/tensor.d.ts +392 -0
  96. package/types/utils/tensor.d.ts.map +1 -0
@@ -0,0 +1,2614 @@
1
+
2
+ /**
3
+ * @file Processors are used to prepare non-textual inputs (e.g., image or audio) for a model.
4
+ *
5
+ * **Example:** Using a `WhisperProcessor` to prepare an audio input for a model.
6
+ * ```javascript
7
+ * import { AutoProcessor, read_audio } from '@huggingface/transformers';
8
+ *
9
+ * let processor = await AutoProcessor.from_pretrained('openai/whisper-tiny.en');
10
+ * let audio = await read_audio('https://huggingface.co/datasets/Narsil/asr_dummy/resolve/main/mlk.flac', 16000);
11
+ * let { input_features } = await processor(audio);
12
+ * // Tensor {
13
+ * // data: Float32Array(240000) [0.4752984642982483, 0.5597258806228638, 0.56434166431427, ...],
14
+ * // dims: [1, 80, 3000],
15
+ * // type: 'float32',
16
+ * // size: 240000,
17
+ * // }
18
+ * ```
19
+ *
20
+ * @module processors
21
+ */
22
+ import {
23
+ Callable,
24
+ } from './utils/generic.js';
25
+
26
+ import {
27
+ calculateDimensions,
28
+ calculateReflectOffset,
29
+ } from './utils/core.js';
30
+
31
+ import {
32
+ getModelJSON,
33
+ } from './utils/hub.js';
34
+
35
+ import {
36
+ min,
37
+ max,
38
+ softmax,
39
+ bankers_round,
40
+ } from './utils/maths.js';
41
+
42
+
43
+ import { Tensor, cat, interpolate, stack, interpolate_4d } from './utils/tensor.js';
44
+
45
+ import { RawImage } from './utils/image.js';
46
+ import {
47
+ window_function,
48
+ spectrogram,
49
+ mel_filter_bank,
50
+ } from './utils/audio.js';
51
+
52
+
53
+ // Helper functions
54
+
55
+ /**
56
+ * Converts bounding boxes from center format to corners format.
57
+ *
58
+ * @param {number[]} arr The coordinate for the center of the box and its width, height dimensions (center_x, center_y, width, height)
59
+ * @returns {number[]} The coodinates for the top-left and bottom-right corners of the box (top_left_x, top_left_y, bottom_right_x, bottom_right_y)
60
+ */
61
+ function center_to_corners_format([centerX, centerY, width, height]) {
62
+ return [
63
+ centerX - width / 2,
64
+ centerY - height / 2,
65
+ centerX + width / 2,
66
+ centerY + height / 2
67
+ ];
68
+ }
69
+
70
+ /**
71
+ * Post-processes the outputs of the model (for object detection).
72
+ * @param {Object} outputs The outputs of the model that must be post-processed
73
+ * @param {Tensor} outputs.logits The logits
74
+ * @param {Tensor} outputs.pred_boxes The predicted boxes.
75
+ * @param {number} [threshold=0.5] The threshold to use for the scores.
76
+ * @param {number[][]} [target_sizes=null] The sizes of the original images.
77
+ * @param {boolean} [is_zero_shot=false] Whether zero-shot object detection was performed.
78
+ * @return {Object[]} An array of objects containing the post-processed outputs.
79
+ * @private
80
+ */
81
+ function post_process_object_detection(outputs, threshold = 0.5, target_sizes = null, is_zero_shot = false) {
82
+ const out_logits = outputs.logits;
83
+ const out_bbox = outputs.pred_boxes;
84
+ const [batch_size, num_boxes, num_classes] = out_logits.dims;
85
+
86
+ if (target_sizes !== null && target_sizes.length !== batch_size) {
87
+ throw Error("Make sure that you pass in as many target sizes as the batch dimension of the logits")
88
+ }
89
+ let toReturn = [];
90
+ for (let i = 0; i < batch_size; ++i) {
91
+ let target_size = target_sizes !== null ? target_sizes[i] : null;
92
+ let info = {
93
+ boxes: [],
94
+ classes: [],
95
+ scores: []
96
+ }
97
+ let logits = out_logits[i];
98
+ let bbox = out_bbox[i];
99
+
100
+ for (let j = 0; j < num_boxes; ++j) {
101
+ let logit = logits[j];
102
+
103
+ let indices = [];
104
+ let probs;
105
+ if (is_zero_shot) {
106
+ // Get indices of classes with high enough probability
107
+ probs = logit.sigmoid().data;
108
+ for (let k = 0; k < probs.length; ++k) {
109
+ if (probs[k] > threshold) {
110
+ indices.push(k);
111
+ }
112
+ }
113
+
114
+ } else {
115
+ // Get most probable class
116
+ let maxIndex = max(logit.data)[1];
117
+
118
+ if (maxIndex === num_classes - 1) {
119
+ // This is the background class, skip it
120
+ continue;
121
+ }
122
+ // Compute softmax over classes
123
+ probs = softmax(logit.data);
124
+
125
+ if (probs[maxIndex] < threshold) {
126
+ continue;
127
+ }
128
+ indices.push(maxIndex);
129
+ }
130
+
131
+ for (const index of indices) {
132
+
133
+ // Some class has a high enough probability
134
+ /** @type {number[]} */
135
+ let box = bbox[j].data;
136
+
137
+ // convert to [x0, y0, x1, y1] format
138
+ box = center_to_corners_format(box)
139
+ if (target_size !== null) {
140
+ box = box.map((x, i) => x * target_size[(i + 1) % 2])
141
+ }
142
+
143
+ info.boxes.push(box);
144
+ info.classes.push(index);
145
+ info.scores.push(probs[index]);
146
+ }
147
+ }
148
+ toReturn.push(info);
149
+ }
150
+ return toReturn;
151
+ }
152
+
153
+ /**
154
+ * Named tuple to indicate the order we are using is (height x width), even though
155
+ * the Graphics’ industry standard is (width x height).
156
+ * @typedef {[height: number, width: number]} HeightWidth
157
+ */
158
+
159
+ /**
160
+ * Helper function to validate audio inputs.
161
+ * @param {any} audio The audio data.
162
+ * @param {string} feature_extractor The name of the feature extractor.
163
+ * @private
164
+ */
165
+ function validate_audio_inputs(audio, feature_extractor) {
166
+ if (!(audio instanceof Float32Array || audio instanceof Float64Array)) {
167
+ throw new Error(
168
+ `${feature_extractor} expects input to be a Float32Array or a Float64Array, but got ${audio?.constructor?.name ?? typeof audio} instead. ` +
169
+ `If using the feature extractor directly, remember to use \`read_audio(url, sampling_rate)\` to obtain the raw audio data of the file/url.`
170
+ )
171
+ }
172
+ }
173
+
174
+ /**
175
+ * Helper function to constrain a value to be a multiple of a number.
176
+ * @param {number} val The value to constrain.
177
+ * @param {number} multiple The number to constrain to.
178
+ * @param {number} [minVal=0] The minimum value to constrain to.
179
+ * @param {number} [maxVal=null] The maximum value to constrain to.
180
+ * @returns {number} The constrained value.
181
+ * @private
182
+ */
183
+ function constraint_to_multiple_of(val, multiple, minVal = 0, maxVal = null) {
184
+ const a = val / multiple;
185
+ let x = bankers_round(a) * multiple;
186
+
187
+ if (maxVal !== null && x > maxVal) {
188
+ x = Math.floor(a) * multiple;
189
+ }
190
+
191
+ if (x < minVal) {
192
+ x = Math.ceil(a) * multiple;
193
+ }
194
+
195
+ return x;
196
+ }
197
+
198
+ /**
199
+ * Rounds the height and width down to the closest multiple of size_divisibility
200
+ * @param {[number, number]} size The size of the image
201
+ * @param {number} divisor The divisor to use.
202
+ * @returns {[number, number]} The rounded size.
203
+ */
204
+ function enforce_size_divisibility([width, height], divisor) {
205
+ return [
206
+ Math.max(Math.floor(width / divisor), 1) * divisor,
207
+ Math.max(Math.floor(height / divisor), 1) * divisor
208
+ ];
209
+ }
210
+
211
+
212
+ /**
213
+ * Base class for feature extractors.
214
+ *
215
+ * @extends Callable
216
+ */
217
+ export class FeatureExtractor extends Callable {
218
+ /**
219
+ * Constructs a new FeatureExtractor instance.
220
+ *
221
+ * @param {Object} config The configuration for the feature extractor.
222
+ */
223
+ constructor(config) {
224
+ super();
225
+ this.config = config
226
+ }
227
+ }
228
+
229
+ /**
230
+ * @typedef {object} ImageFeatureExtractorResult
231
+ * @property {Tensor} pixel_values The pixel values of the batched preprocessed images.
232
+ * @property {HeightWidth[]} original_sizes Array of two-dimensional tuples like [[480, 640]].
233
+ * @property {HeightWidth[]} reshaped_input_sizes Array of two-dimensional tuples like [[1000, 1330]].
234
+ */
235
+
236
+ /**
237
+ * Feature extractor for image models.
238
+ *
239
+ * @extends FeatureExtractor
240
+ */
241
+ export class ImageFeatureExtractor extends FeatureExtractor {
242
+
243
+ /**
244
+ * Constructs a new ImageFeatureExtractor instance.
245
+ *
246
+ * @param {Object} config The configuration for the feature extractor.
247
+ * @param {number[]} config.image_mean The mean values for image normalization.
248
+ * @param {number[]} config.image_std The standard deviation values for image normalization.
249
+ * @param {boolean} config.do_rescale Whether to rescale the image pixel values to the [0,1] range.
250
+ * @param {number} config.rescale_factor The factor to use for rescaling the image pixel values.
251
+ * @param {boolean} config.do_normalize Whether to normalize the image pixel values.
252
+ * @param {boolean} config.do_resize Whether to resize the image.
253
+ * @param {number} config.resample What method to use for resampling.
254
+ * @param {number|Object} config.size The size to resize the image to.
255
+ * @param {boolean} [config.do_flip_channel_order=false] Whether to flip the color channels from RGB to BGR.
256
+ * Can be overridden by the `do_flip_channel_order` parameter in the `preprocess` method.
257
+ */
258
+ constructor(config) {
259
+ super(config);
260
+
261
+ this.image_mean = this.config.image_mean ?? this.config.mean;
262
+ this.image_std = this.config.image_std ?? this.config.std;
263
+
264
+ this.resample = this.config.resample ?? 2; // 2 => bilinear
265
+ this.do_rescale = this.config.do_rescale ?? true;
266
+ this.rescale_factor = this.config.rescale_factor ?? (1 / 255);
267
+ this.do_normalize = this.config.do_normalize;
268
+
269
+ this.do_resize = this.config.do_resize;
270
+ this.do_thumbnail = this.config.do_thumbnail;
271
+ this.size = this.config.size;
272
+ this.size_divisibility = this.config.size_divisibility ?? this.config.size_divisor;
273
+
274
+ this.do_center_crop = this.config.do_center_crop;
275
+ this.crop_size = this.config.crop_size;
276
+ this.do_convert_rgb = this.config.do_convert_rgb ?? true;
277
+ this.do_crop_margin = this.config.do_crop_margin;
278
+
279
+ this.pad_size = this.config.pad_size;
280
+ this.do_pad = this.config.do_pad;
281
+
282
+ if (this.do_pad && !this.pad_size && this.size && this.size.width !== undefined && this.size.height !== undefined) {
283
+ // Should pad, but no pad size specified
284
+ // We infer the pad size from the resize size
285
+ this.pad_size = this.size
286
+ }
287
+
288
+ this.do_flip_channel_order = this.config.do_flip_channel_order ?? false;
289
+ }
290
+
291
+ /**
292
+ * Resize the image to make a thumbnail. The image is resized so that no dimension is larger than any
293
+ * corresponding dimension of the specified size.
294
+ * @param {RawImage} image The image to be resized.
295
+ * @param {{height:number, width:number}} size The size `{"height": h, "width": w}` to resize the image to.
296
+ * @param {string | 0 | 1 | 2 | 3 | 4 | 5} [resample=2] The resampling filter to use.
297
+ * @returns {Promise<RawImage>} The resized image.
298
+ */
299
+ async thumbnail(image, size, resample = 2) {
300
+ const input_height = image.height;
301
+ const input_width = image.width;
302
+
303
+ const output_height = size.height;
304
+ const output_width = size.width;
305
+
306
+ // We always resize to the smallest of either the input or output size.
307
+ let height = Math.min(input_height, output_height)
308
+ let width = Math.min(input_width, output_width)
309
+
310
+ if (height === input_height && width === input_width) {
311
+ return image;
312
+ }
313
+ if (input_height > input_width) {
314
+ width = Math.floor(input_width * height / input_height);
315
+ } else if (input_width > input_height) {
316
+ height = Math.floor(input_height * width / input_width);
317
+ }
318
+ return await image.resize(width, height, { resample });
319
+ }
320
+
321
+
322
+ /**
323
+ * Crops the margin of the image. Gray pixels are considered margin (i.e., pixels with a value below the threshold).
324
+ * @param {RawImage} image The image to be cropped.
325
+ * @param {number} gray_threshold Value below which pixels are considered to be gray.
326
+ * @returns {Promise<RawImage>} The cropped image.
327
+ */
328
+ async crop_margin(image, gray_threshold = 200) {
329
+
330
+ const gray_image = image.clone().grayscale();
331
+
332
+ const minValue = min(gray_image.data)[0];
333
+ const maxValue = max(gray_image.data)[0];
334
+ const diff = maxValue - minValue;
335
+
336
+ if (diff === 0) {
337
+ return image;
338
+ }
339
+
340
+ const threshold = gray_threshold / 255;
341
+
342
+ let x_min = gray_image.width, y_min = gray_image.height, x_max = 0, y_max = 0;
343
+ const gray_image_data = gray_image.data;
344
+ for (let j = 0; j < gray_image.height; ++j) {
345
+ const row = j * gray_image.width;
346
+ for (let i = 0; i < gray_image.width; ++i) {
347
+ if ((gray_image_data[row + i] - minValue) / diff < threshold) {
348
+ // We have a non-zero pixel, so we update the min/max values accordingly
349
+ x_min = Math.min(x_min, i);
350
+ y_min = Math.min(y_min, j);
351
+ x_max = Math.max(x_max, i);
352
+ y_max = Math.max(y_max, j);
353
+ }
354
+ }
355
+ }
356
+
357
+ image = await image.crop([x_min, y_min, x_max, y_max]);
358
+ return image;
359
+ }
360
+
361
+ /**
362
+ * Pad the image by a certain amount.
363
+ * @param {Float32Array} pixelData The pixel data to pad.
364
+ * @param {number[]} imgDims The dimensions of the image (height, width, channels).
365
+ * @param {{width:number; height:number}|number} padSize The dimensions of the padded image.
366
+ * @param {Object} options The options for padding.
367
+ * @param {'constant'|'symmetric'} [options.mode='constant'] The type of padding to add.
368
+ * @param {boolean} [options.center=false] Whether to center the image.
369
+ * @param {number} [options.constant_values=0] The constant value to use for padding.
370
+ * @returns {[Float32Array, number[]]} The padded pixel data and image dimensions.
371
+ */
372
+ pad_image(pixelData, imgDims, padSize, {
373
+ mode = 'constant',
374
+ center = false,
375
+ constant_values = 0,
376
+ } = {}) {
377
+ const [imageHeight, imageWidth, imageChannels] = imgDims;
378
+
379
+ let paddedImageWidth, paddedImageHeight;
380
+ if (typeof padSize === 'number') {
381
+ paddedImageWidth = padSize;
382
+ paddedImageHeight = padSize;
383
+ } else {
384
+ paddedImageWidth = padSize.width;
385
+ paddedImageHeight = padSize.height;
386
+ }
387
+
388
+ // Only add padding if there is a difference in size
389
+ if (paddedImageWidth !== imageWidth || paddedImageHeight !== imageHeight) {
390
+ const paddedPixelData = new Float32Array(paddedImageWidth * paddedImageHeight * imageChannels);
391
+ if (Array.isArray(constant_values)) {
392
+ // Fill with constant values, cycling through the array
393
+ for (let i = 0; i < paddedPixelData.length; ++i) {
394
+ paddedPixelData[i] = constant_values[i % imageChannels];
395
+ }
396
+ } else if (constant_values !== 0) {
397
+ paddedPixelData.fill(constant_values);
398
+ }
399
+
400
+ const [left, top] = center
401
+ ? [Math.floor((paddedImageWidth - imageWidth) / 2), Math.floor((paddedImageHeight - imageHeight) / 2)]
402
+ : [0, 0];
403
+
404
+ // Copy the original image into the padded image
405
+ for (let i = 0; i < imageHeight; ++i) {
406
+ const a = (i + top) * paddedImageWidth;
407
+ const b = i * imageWidth;
408
+ for (let j = 0; j < imageWidth; ++j) {
409
+ const c = (a + j + left) * imageChannels;
410
+ const d = (b + j) * imageChannels;
411
+ for (let k = 0; k < imageChannels; ++k) {
412
+ paddedPixelData[c + k] = pixelData[d + k];
413
+ }
414
+ }
415
+ }
416
+
417
+ if (mode === 'symmetric') {
418
+ if (center) {
419
+ throw new Error('`center` padding is not supported when `mode` is set to `symmetric`.');
420
+ // TODO: Implement this
421
+ }
422
+ const h1 = imageHeight - 1;
423
+ const w1 = imageWidth - 1;
424
+ for (let i = 0; i < paddedImageHeight; ++i) {
425
+ const a = i * paddedImageWidth;
426
+ const b = calculateReflectOffset(i, h1) * imageWidth;
427
+
428
+ for (let j = 0; j < paddedImageWidth; ++j) {
429
+ if (i < imageHeight && j < imageWidth) continue; // Do not overwrite original image
430
+ const c = (a + j) * imageChannels;
431
+ const d = (b + calculateReflectOffset(j, w1)) * imageChannels;
432
+
433
+ // Copy channel-wise
434
+ for (let k = 0; k < imageChannels; ++k) {
435
+ paddedPixelData[c + k] = pixelData[d + k];
436
+ }
437
+ }
438
+ }
439
+ }
440
+
441
+
442
+ // Update pixel data and image dimensions
443
+ pixelData = paddedPixelData;
444
+ imgDims = [paddedImageHeight, paddedImageWidth, imageChannels]
445
+ }
446
+ return [pixelData, imgDims];
447
+ }
448
+
449
+ /**
450
+ * Rescale the image' pixel values by `this.rescale_factor`.
451
+ * @param {Float32Array} pixelData The pixel data to rescale.
452
+ * @returns {void}
453
+ */
454
+ rescale(pixelData) {
455
+ for (let i = 0; i < pixelData.length; ++i) {
456
+ pixelData[i] = this.rescale_factor * pixelData[i];
457
+ }
458
+ }
459
+
460
+ /**
461
+ * Find the target (width, height) dimension of the output image after
462
+ * resizing given the input image and the desired size.
463
+ * @param {RawImage} image The image to resize.
464
+ * @param {any} size The size to use for resizing the image.
465
+ * @returns {[number, number]} The target (width, height) dimension of the output image after resizing.
466
+ */
467
+ get_resize_output_image_size(image, size) {
468
+ // `size` comes in many forms, so we need to handle them all here:
469
+ // 1. `size` is an integer, in which case we resize the image to be a square
470
+
471
+ const [srcWidth, srcHeight] = image.size;
472
+
473
+ let shortest_edge;
474
+ let longest_edge;
475
+
476
+ if (this.do_thumbnail) {
477
+ // NOTE: custom logic for `Donut` models
478
+ const { height, width } = size;
479
+ shortest_edge = Math.min(height, width)
480
+ }
481
+ // Support both formats for backwards compatibility
482
+ else if (Number.isInteger(size)) {
483
+ shortest_edge = size;
484
+ longest_edge = this.config.max_size ?? shortest_edge;
485
+
486
+ } else if (size !== undefined) {
487
+ // Extract known properties from `size`
488
+ shortest_edge = size.shortest_edge;
489
+ longest_edge = size.longest_edge;
490
+ }
491
+
492
+ // If `longest_edge` and `shortest_edge` are set, maintain aspect ratio and resize to `shortest_edge`
493
+ // while keeping the largest dimension <= `longest_edge`
494
+ if (shortest_edge !== undefined || longest_edge !== undefined) {
495
+ // http://opensourcehacker.com/2011/12/01/calculate-aspect-ratio-conserving-resize-for-images-in-javascript/
496
+ // Try resize so that shortest edge is `shortest_edge` (target)
497
+ const shortResizeFactor = shortest_edge === undefined
498
+ ? 1 // If `shortest_edge` is not set, don't upscale
499
+ : Math.max(shortest_edge / srcWidth, shortest_edge / srcHeight);
500
+
501
+ const newWidth = srcWidth * shortResizeFactor;
502
+ const newHeight = srcHeight * shortResizeFactor;
503
+
504
+ // The new width and height might be greater than `longest_edge`, so
505
+ // we downscale again to ensure the largest dimension is `longest_edge`
506
+ const longResizeFactor = longest_edge === undefined
507
+ ? 1 // If `longest_edge` is not set, don't downscale
508
+ : Math.min(longest_edge / newWidth, longest_edge / newHeight);
509
+
510
+ // To avoid certain floating point precision issues, we round to 2 decimal places
511
+ let finalWidth = Math.floor(Number((newWidth * longResizeFactor).toFixed(2)));
512
+ let finalHeight = Math.floor(Number((newHeight * longResizeFactor).toFixed(2)));
513
+
514
+ if (this.size_divisibility !== undefined) {
515
+ [finalWidth, finalHeight] = enforce_size_divisibility([finalWidth, finalHeight], this.size_divisibility)
516
+ }
517
+ return [finalWidth, finalHeight];
518
+
519
+ } else if (size !== undefined && size.width !== undefined && size.height !== undefined) {
520
+ // If `width` and `height` are set, resize to those dimensions
521
+
522
+ let newWidth = size.width;
523
+ let newHeight = size.height;
524
+
525
+ // Custom for DPT models
526
+ if (this.config.keep_aspect_ratio && this.config.ensure_multiple_of) {
527
+
528
+ // determine new height and width
529
+ let scale_height = newHeight / srcHeight;
530
+ let scale_width = newWidth / srcWidth;
531
+
532
+ // scale as little as possible
533
+ if (Math.abs(1 - scale_width) < Math.abs(1 - scale_height)) {
534
+ // fit width
535
+ scale_height = scale_width;
536
+ } else {
537
+ // fit height
538
+ scale_width = scale_height;
539
+ }
540
+
541
+ newHeight = constraint_to_multiple_of(scale_height * srcHeight, this.config.ensure_multiple_of);
542
+ newWidth = constraint_to_multiple_of(scale_width * srcWidth, this.config.ensure_multiple_of);
543
+ }
544
+
545
+ return [newWidth, newHeight];
546
+
547
+ } else if (this.size_divisibility !== undefined) {
548
+ return enforce_size_divisibility([srcWidth, srcHeight], this.size_divisibility);
549
+ } else {
550
+ throw new Error(`Could not resize image due to unsupported \`this.size\` option in config: ${JSON.stringify(size)}`);
551
+ }
552
+ }
553
+
554
+ /**
555
+ * Resizes the image.
556
+ * @param {RawImage} image The image to resize.
557
+ * @returns {Promise<RawImage>} The resized image.
558
+ */
559
+ async resize(image) {
560
+ const [newWidth, newHeight] = this.get_resize_output_image_size(image, this.size);
561
+ return await image.resize(newWidth, newHeight, {
562
+ resample: this.resample,
563
+ });
564
+ }
565
+
566
+ /**
567
+ * @typedef {object} PreprocessedImage
568
+ * @property {HeightWidth} original_size The original size of the image.
569
+ * @property {HeightWidth} reshaped_input_size The reshaped input size of the image.
570
+ * @property {Tensor} pixel_values The pixel values of the preprocessed image.
571
+ */
572
+
573
+ /**
574
+ * Preprocesses the given image.
575
+ *
576
+ * @param {RawImage} image The image to preprocess.
577
+ * @param {Object} overrides The overrides for the preprocessing options.
578
+ * @returns {Promise<PreprocessedImage>} The preprocessed image.
579
+ */
580
+ async preprocess(image, {
581
+ do_normalize = null,
582
+ do_pad = null,
583
+ do_convert_rgb = null,
584
+ do_convert_grayscale = null,
585
+ do_flip_channel_order = null,
586
+ } = {}) {
587
+ if (this.do_crop_margin) {
588
+ // NOTE: Specific to nougat processors. This is done before resizing,
589
+ // and can be interpreted as a pre-preprocessing step.
590
+ image = await this.crop_margin(image);
591
+ }
592
+
593
+ const [srcWidth, srcHeight] = image.size; // original image size
594
+
595
+ // Convert image to RGB if specified in config.
596
+ if (do_convert_rgb ?? this.do_convert_rgb) {
597
+ image = image.rgb();
598
+ } else if (do_convert_grayscale) {
599
+ image = image.grayscale();
600
+ }
601
+
602
+ // TODO:
603
+ // For efficiency reasons, it might be best to merge the resize and center crop operations into one.
604
+
605
+ // Resize all images
606
+ if (this.do_resize) {
607
+ image = await this.resize(image);
608
+ }
609
+
610
+ // Resize the image using thumbnail method.
611
+ if (this.do_thumbnail) {
612
+ image = await this.thumbnail(image, this.size, this.resample);
613
+ }
614
+
615
+ if (this.do_center_crop) {
616
+
617
+ let crop_width;
618
+ let crop_height;
619
+ if (Number.isInteger(this.crop_size)) {
620
+ crop_width = this.crop_size;
621
+ crop_height = this.crop_size;
622
+ } else {
623
+ crop_width = this.crop_size.width;
624
+ crop_height = this.crop_size.height;
625
+ }
626
+
627
+ image = await image.center_crop(crop_width, crop_height);
628
+ }
629
+
630
+ /** @type {HeightWidth} */
631
+ const reshaped_input_size = [image.height, image.width];
632
+
633
+ // NOTE: All pixel-level manipulation (i.e., modifying `pixelData`)
634
+ // occurs with data in the hwc format (height, width, channels),
635
+ // to emulate the behavior of the original Python code (w/ numpy).
636
+ let pixelData = Float32Array.from(image.data);
637
+ let imgDims = [image.height, image.width, image.channels];
638
+
639
+ if (this.do_rescale) {
640
+ this.rescale(pixelData);
641
+ }
642
+
643
+ if (do_normalize ?? this.do_normalize) {
644
+ let image_mean = this.image_mean;
645
+ if (!Array.isArray(this.image_mean)) {
646
+ image_mean = new Array(image.channels).fill(image_mean);
647
+ }
648
+
649
+ let image_std = this.image_std;
650
+ if (!Array.isArray(this.image_std)) {
651
+ image_std = new Array(image.channels).fill(image_mean);
652
+ }
653
+
654
+ if (image_mean.length !== image.channels || image_std.length !== image.channels) {
655
+ throw new Error(`When set to arrays, the length of \`image_mean\` (${image_mean.length}) and \`image_std\` (${image_std.length}) must match the number of channels in the image (${image.channels}).`);
656
+ }
657
+
658
+ for (let i = 0; i < pixelData.length; i += image.channels) {
659
+ for (let j = 0; j < image.channels; ++j) {
660
+ pixelData[i + j] = (pixelData[i + j] - image_mean[j]) / image_std[j];
661
+ }
662
+ }
663
+ }
664
+
665
+ // do padding after rescaling/normalizing
666
+ if (do_pad ?? this.do_pad) {
667
+ if (this.pad_size) {
668
+ const padded = this.pad_image(pixelData, [image.height, image.width, image.channels], this.pad_size);
669
+ [pixelData, imgDims] = padded; // Update pixel data and image dimensions
670
+ } else if (this.size_divisibility) {
671
+ const [paddedWidth, paddedHeight] = enforce_size_divisibility([imgDims[1], imgDims[0]], this.size_divisibility);
672
+ [pixelData, imgDims] = this.pad_image(pixelData, imgDims, { width: paddedWidth, height: paddedHeight });
673
+ }
674
+ }
675
+
676
+ if (do_flip_channel_order ?? this.do_flip_channel_order) {
677
+ if (imgDims[2] !== 3) {
678
+ throw new Error('Flipping channel order is only supported for RGB images.');
679
+ }
680
+ // Convert RGB to BGR
681
+ for (let i = 0; i < pixelData.length; i += 3) {
682
+ const temp = pixelData[i];
683
+ pixelData[i] = pixelData[i + 2];
684
+ pixelData[i + 2] = temp;
685
+ }
686
+ }
687
+
688
+ const pixel_values = new Tensor('float32', pixelData, imgDims)
689
+ .permute(2, 0, 1); // convert to channel dimension format (hwc -> chw)
690
+
691
+ return {
692
+ original_size: [srcHeight, srcWidth],
693
+ reshaped_input_size: reshaped_input_size,
694
+ pixel_values,
695
+ }
696
+ }
697
+
698
+ /**
699
+ * Calls the feature extraction process on an array of images,
700
+ * preprocesses each image, and concatenates the resulting
701
+ * features into a single Tensor.
702
+ * @param {RawImage[]} images The image(s) to extract features from.
703
+ * @param {...any} args Additional arguments.
704
+ * @returns {Promise<ImageFeatureExtractorResult>} An object containing the concatenated pixel values (and other metadata) of the preprocessed images.
705
+ */
706
+ async _call(images, ...args) {
707
+ if (!Array.isArray(images)) {
708
+ images = [images];
709
+ }
710
+ /** @type {PreprocessedImage[]} */
711
+ const imageData = await Promise.all(images.map(x => this.preprocess(x)));
712
+
713
+ // Stack pixel values
714
+ const pixel_values = stack(imageData.map(x => x.pixel_values), 0);
715
+
716
+ return {
717
+ pixel_values,
718
+
719
+ // Original sizes of images
720
+ original_sizes: imageData.map(x => x.original_size),
721
+
722
+ // Reshaped sizes of images, before padding or cropping
723
+ reshaped_input_sizes: imageData.map(x => x.reshaped_input_size),
724
+ }
725
+ }
726
+
727
+ }
728
+
729
+ export class SegformerFeatureExtractor extends ImageFeatureExtractor {
730
+
731
+ /**
732
+ * Converts the output of `SegformerForSemanticSegmentation` into semantic segmentation maps.
733
+ * @param {*} outputs Raw outputs of the model.
734
+ * @param {number[][]} [target_sizes=null] List of tuples corresponding to the requested final size
735
+ * (height, width) of each prediction. If unset, predictions will not be resized.
736
+ * @returns {{segmentation: Tensor; labels: number[]}[]} The semantic segmentation maps.
737
+ */
738
+ post_process_semantic_segmentation(outputs, target_sizes = null) {
739
+
740
+ const logits = outputs.logits;
741
+ const batch_size = logits.dims[0];
742
+
743
+ if (target_sizes !== null && target_sizes.length !== batch_size) {
744
+ throw Error("Make sure that you pass in as many target sizes as the batch dimension of the logits")
745
+ }
746
+
747
+ const toReturn = [];
748
+ for (let i = 0; i < batch_size; ++i) {
749
+ const target_size = target_sizes !== null ? target_sizes[i] : null;
750
+
751
+ let data = logits[i];
752
+
753
+ // 1. If target_size is not null, we need to resize the masks to the target size
754
+ if (target_size !== null) {
755
+ // resize the masks to the target size
756
+ data = interpolate(data, target_size, 'bilinear', false);
757
+ }
758
+ const [height, width] = target_size ?? data.dims.slice(-2);
759
+
760
+ const segmentation = new Tensor(
761
+ 'int32',
762
+ new Int32Array(height * width),
763
+ [height, width]
764
+ );
765
+
766
+ // Buffer to store current largest value
767
+ const buffer = data[0].data;
768
+ const segmentation_data = segmentation.data;
769
+ for (let j = 1; j < data.dims[0]; ++j) {
770
+ const row = data[j].data;
771
+ for (let k = 0; k < row.length; ++k) {
772
+ if (row[k] > buffer[k]) {
773
+ buffer[k] = row[k];
774
+ segmentation_data[k] = j;
775
+ }
776
+ }
777
+ }
778
+
779
+ // Store which objects have labels
780
+ // This is much more efficient that creating a set of the final values
781
+ const hasLabel = new Array(data.dims[0]);
782
+ const out = segmentation.data;
783
+ for (let j = 0; j < out.length; ++j) {
784
+ const index = out[j];
785
+ hasLabel[index] = index;
786
+ }
787
+ /** @type {number[]} The unique list of labels that were detected */
788
+ const labels = hasLabel.filter(x => x !== undefined);
789
+
790
+ toReturn.push({ segmentation, labels });
791
+ }
792
+ return toReturn;
793
+ }
794
+ }
795
+ export class DPTFeatureExtractor extends ImageFeatureExtractor { }
796
+ export class DPTImageProcessor extends DPTFeatureExtractor { } // NOTE: extends DPTFeatureExtractor
797
+ export class BitImageProcessor extends ImageFeatureExtractor { }
798
+ export class GLPNFeatureExtractor extends ImageFeatureExtractor { }
799
+ export class CLIPFeatureExtractor extends ImageFeatureExtractor { }
800
+ export class CLIPImageProcessor extends CLIPFeatureExtractor { } // NOTE: extends CLIPFeatureExtractor
801
+ export class ChineseCLIPFeatureExtractor extends ImageFeatureExtractor { }
802
+ export class SiglipImageProcessor extends ImageFeatureExtractor { }
803
+ export class ConvNextFeatureExtractor extends ImageFeatureExtractor {
804
+ constructor(config) {
805
+ super(config);
806
+
807
+ /**
808
+ * Percentage of the image to crop. Only has an effect if this.size < 384.
809
+ */
810
+ this.crop_pct = this.config.crop_pct ?? (224 / 256);
811
+ }
812
+
813
+ async resize(image) {
814
+ const shortest_edge = this.size?.shortest_edge;
815
+ if (shortest_edge === undefined) {
816
+ throw new Error(`Size dictionary must contain 'shortest_edge' key.`);
817
+ }
818
+
819
+ if (shortest_edge < 384) {
820
+ // maintain same ratio, resizing shortest edge to shortest_edge/crop_pct
821
+ const resize_shortest_edge = Math.floor(shortest_edge / this.crop_pct);
822
+
823
+ const [newWidth, newHeight] = this.get_resize_output_image_size(image, {
824
+ shortest_edge: resize_shortest_edge,
825
+ });
826
+
827
+ image = await image.resize(newWidth, newHeight, {
828
+ resample: this.resample,
829
+ });
830
+
831
+ // then crop to (shortest_edge, shortest_edge)
832
+ image = await image.center_crop(shortest_edge, shortest_edge);
833
+ } else {
834
+ // warping (no cropping) when evaluated at 384 or larger
835
+ image = await image.resize(shortest_edge, shortest_edge, {
836
+ resample: this.resample,
837
+ });
838
+ }
839
+
840
+ return image;
841
+ }
842
+ }
843
+ export class ConvNextImageProcessor extends ConvNextFeatureExtractor { } // NOTE extends ConvNextFeatureExtractor
844
+ export class ViTFeatureExtractor extends ImageFeatureExtractor { }
845
+ export class ViTImageProcessor extends ImageFeatureExtractor { }
846
+
847
+ export class EfficientNetImageProcessor extends ImageFeatureExtractor {
848
+ constructor(config) {
849
+ super(config);
850
+ this.include_top = this.config.include_top ?? true;
851
+ if (this.include_top) {
852
+ this.image_std = this.image_std.map(x => x * x);
853
+ }
854
+ }
855
+ }
856
+
857
+ export class MobileNetV1FeatureExtractor extends ImageFeatureExtractor { }
858
+ export class MobileNetV2FeatureExtractor extends ImageFeatureExtractor { }
859
+ export class MobileNetV3FeatureExtractor extends ImageFeatureExtractor { }
860
+ export class MobileNetV4FeatureExtractor extends ImageFeatureExtractor { }
861
+
862
+ export class MobileViTFeatureExtractor extends ImageFeatureExtractor { }
863
+ export class MobileViTImageProcessor extends MobileViTFeatureExtractor { } // NOTE extends MobileViTFeatureExtractor
864
+ export class OwlViTFeatureExtractor extends ImageFeatureExtractor {
865
+ /** @type {post_process_object_detection} */
866
+ post_process_object_detection(...args) {
867
+ return post_process_object_detection(...args);
868
+ }
869
+ }
870
+ export class Owlv2ImageProcessor extends OwlViTFeatureExtractor { } // NOTE extends OwlViTFeatureExtractor
871
+
872
+ export class RTDetrImageProcessor extends ImageFeatureExtractor {
873
+ /** @type {post_process_object_detection} */
874
+ post_process_object_detection(...args) {
875
+ return post_process_object_detection(...args);
876
+ }
877
+ }
878
+
879
+ export class DeiTFeatureExtractor extends ImageFeatureExtractor { }
880
+ export class BeitFeatureExtractor extends ImageFeatureExtractor { }
881
+ export class DonutFeatureExtractor extends ImageFeatureExtractor {
882
+ pad_image(pixelData, imgDims, padSize, options = {}) {
883
+ const [imageHeight, imageWidth, imageChannels] = imgDims;
884
+
885
+ let image_mean = this.image_mean;
886
+ if (!Array.isArray(this.image_mean)) {
887
+ image_mean = new Array(imageChannels).fill(image_mean);
888
+ }
889
+
890
+ let image_std = this.image_std;
891
+ if (!Array.isArray(image_std)) {
892
+ image_std = new Array(imageChannels).fill(image_mean);
893
+ }
894
+
895
+ const constant_values = image_mean.map((x, i) => - x / image_std[i]);
896
+
897
+ return super.pad_image(pixelData, imgDims, padSize, {
898
+ center: true,
899
+
900
+ // Since normalization is done after padding, we need to use certain constant values to ensure the same behaviour is observed.
901
+ // For more information, see https://github.com/huggingface/transformers/blob/main/src/transformers/models/donut/image_processing_donut.py#L433-L451
902
+ constant_values: constant_values,
903
+ ...options,
904
+ });
905
+ }
906
+ }
907
+ export class NougatImageProcessor extends DonutFeatureExtractor { } // NOTE extends DonutFeatureExtractor
908
+
909
+ /**
910
+ * @typedef {object} DetrFeatureExtractorResultProps
911
+ * @property {Tensor} pixel_mask
912
+ * @typedef {ImageFeatureExtractorResult & DetrFeatureExtractorResultProps} DetrFeatureExtractorResult
913
+ */
914
+
915
+ /**
916
+ * Detr Feature Extractor.
917
+ *
918
+ * @extends ImageFeatureExtractor
919
+ */
920
+ export class DetrFeatureExtractor extends ImageFeatureExtractor {
921
+ /**
922
+ * Calls the feature extraction process on an array of images, preprocesses
923
+ * each image, and concatenates the resulting features into a single Tensor.
924
+ * @param {RawImage[]} images The image(s) to extract features from.
925
+ * @returns {Promise<DetrFeatureExtractorResult>} An object containing the concatenated pixel values of the preprocessed images.
926
+ */
927
+ async _call(images) {
928
+ const result = await super._call(images);
929
+
930
+ // TODO support differently-sized images, for now assume all images are the same size.
931
+ // TODO support different mask sizes (not just 64x64)
932
+ // Currently, just fill pixel mask with 1s
933
+ const maskSize = [result.pixel_values.dims[0], 64, 64];
934
+ const pixel_mask = new Tensor(
935
+ 'int64',
936
+ new BigInt64Array(maskSize.reduce((a, b) => a * b)).fill(1n),
937
+ maskSize
938
+ );
939
+
940
+ return { ...result, pixel_mask };
941
+ }
942
+
943
+ /**
944
+ * Post-processes the outputs of the model (for object detection).
945
+ * @param {Object} outputs The outputs of the model that must be post-processed
946
+ * @param {Tensor} outputs.logits The logits
947
+ * @param {Tensor} outputs.pred_boxes The predicted boxes.
948
+ * @return {Object[]} An array of objects containing the post-processed outputs.
949
+ */
950
+
951
+ /** @type {post_process_object_detection} */
952
+ post_process_object_detection(...args) {
953
+ return post_process_object_detection(...args);
954
+ }
955
+
956
+ /**
957
+ * Binarize the given masks using `object_mask_threshold`, it returns the associated values of `masks`, `scores` and `labels`.
958
+ * @param {Tensor} class_logits The class logits.
959
+ * @param {Tensor} mask_logits The mask logits.
960
+ * @param {number} object_mask_threshold A number between 0 and 1 used to binarize the masks.
961
+ * @param {number} num_labels The number of labels.
962
+ * @returns {[Tensor[], number[], number[]]} The binarized masks, the scores, and the labels.
963
+ */
964
+ remove_low_and_no_objects(class_logits, mask_logits, object_mask_threshold, num_labels) {
965
+
966
+ let mask_probs_item = [];
967
+ let pred_scores_item = [];
968
+ let pred_labels_item = [];
969
+
970
+ for (let j = 0; j < class_logits.dims[0]; ++j) {
971
+ let cls = class_logits[j];
972
+ let mask = mask_logits[j];
973
+
974
+ let pred_label = max(cls.data)[1];
975
+ if (pred_label === num_labels) {
976
+ // Is the background, so we ignore it
977
+ continue;
978
+ }
979
+
980
+ let scores = softmax(cls.data);
981
+ let pred_score = scores[pred_label];
982
+ if (pred_score > object_mask_threshold) {
983
+ mask_probs_item.push(mask);
984
+ pred_scores_item.push(pred_score);
985
+ pred_labels_item.push(pred_label);
986
+ }
987
+ }
988
+
989
+ return [mask_probs_item, pred_scores_item, pred_labels_item];
990
+
991
+ }
992
+
993
+ /**
994
+ * Checks whether the segment is valid or not.
995
+ * @param {Int32Array} mask_labels Labels for each pixel in the mask.
996
+ * @param {Tensor[]} mask_probs Probabilities for each pixel in the masks.
997
+ * @param {number} k The class id of the segment.
998
+ * @param {number} mask_threshold The mask threshold.
999
+ * @param {number} overlap_mask_area_threshold The overlap mask area threshold.
1000
+ * @returns {[boolean, number[]]} Whether the segment is valid or not, and the indices of the valid labels.
1001
+ */
1002
+ check_segment_validity(
1003
+ mask_labels,
1004
+ mask_probs,
1005
+ k,
1006
+ mask_threshold = 0.5,
1007
+ overlap_mask_area_threshold = 0.8
1008
+ ) {
1009
+ // mask_k is a 1D array of indices, indicating where the mask is equal to k
1010
+ let mask_k = [];
1011
+ let mask_k_area = 0;
1012
+ let original_area = 0;
1013
+
1014
+ const mask_probs_k_data = mask_probs[k].data;
1015
+
1016
+ // Compute the area of all the stuff in query k
1017
+ for (let i = 0; i < mask_labels.length; ++i) {
1018
+ if (mask_labels[i] === k) {
1019
+ mask_k.push(i);
1020
+ ++mask_k_area;
1021
+ }
1022
+
1023
+ if (mask_probs_k_data[i] >= mask_threshold) {
1024
+ ++original_area;
1025
+ }
1026
+ }
1027
+ let mask_exists = mask_k_area > 0 && original_area > 0;
1028
+
1029
+ // Eliminate disconnected tiny segments
1030
+ if (mask_exists) {
1031
+ // Perform additional check
1032
+ let area_ratio = mask_k_area / original_area;
1033
+ mask_exists = area_ratio > overlap_mask_area_threshold;
1034
+ }
1035
+
1036
+ return [mask_exists, mask_k]
1037
+ }
1038
+
1039
+ /**
1040
+ * Computes the segments.
1041
+ * @param {Tensor[]} mask_probs The mask probabilities.
1042
+ * @param {number[]} pred_scores The predicted scores.
1043
+ * @param {number[]} pred_labels The predicted labels.
1044
+ * @param {number} mask_threshold The mask threshold.
1045
+ * @param {number} overlap_mask_area_threshold The overlap mask area threshold.
1046
+ * @param {Set<number>} label_ids_to_fuse The label ids to fuse.
1047
+ * @param {number[]} target_size The target size of the image.
1048
+ * @returns {[Tensor, Array<{id: number, label_id: number, score: number}>]} The computed segments.
1049
+ */
1050
+ compute_segments(
1051
+ mask_probs,
1052
+ pred_scores,
1053
+ pred_labels,
1054
+ mask_threshold,
1055
+ overlap_mask_area_threshold,
1056
+ label_ids_to_fuse = null,
1057
+ target_size = null,
1058
+ ) {
1059
+ let [height, width] = target_size ?? mask_probs[0].dims;
1060
+
1061
+ let segmentation = new Tensor(
1062
+ 'int32',
1063
+ new Int32Array(height * width),
1064
+ [height, width]
1065
+ );
1066
+ let segments = [];
1067
+
1068
+ // 1. If target_size is not null, we need to resize the masks to the target size
1069
+ if (target_size !== null) {
1070
+ // resize the masks to the target size
1071
+ for (let i = 0; i < mask_probs.length; ++i) {
1072
+ mask_probs[i] = interpolate(mask_probs[i], target_size, 'bilinear', false);
1073
+ }
1074
+ }
1075
+
1076
+ // 2. Weigh each mask by its prediction score
1077
+ // NOTE: `mask_probs` is updated in-place
1078
+ //
1079
+ // Temporary storage for the best label/scores for each pixel ([height, width]):
1080
+ let mask_labels = new Int32Array(mask_probs[0].data.length);
1081
+ let bestScores = new Float32Array(mask_probs[0].data.length);
1082
+
1083
+ for (let i = 0; i < mask_probs.length; ++i) {
1084
+ let score = pred_scores[i];
1085
+
1086
+ const mask_probs_i_data = mask_probs[i].data;
1087
+
1088
+ for (let j = 0; j < mask_probs_i_data.length; ++j) {
1089
+ mask_probs_i_data[j] *= score
1090
+ if (mask_probs_i_data[j] > bestScores[j]) {
1091
+ mask_labels[j] = i;
1092
+ bestScores[j] = mask_probs_i_data[j];
1093
+ }
1094
+ }
1095
+ }
1096
+
1097
+ let current_segment_id = 0;
1098
+
1099
+ // let stuff_memory_list = {}
1100
+ const segmentation_data = segmentation.data;
1101
+ for (let k = 0; k < pred_labels.length; ++k) {
1102
+ let pred_class = pred_labels[k];
1103
+
1104
+ // TODO add `should_fuse`
1105
+ // let should_fuse = pred_class in label_ids_to_fuse
1106
+
1107
+ // Check if mask exists and large enough to be a segment
1108
+ let [mask_exists, mask_k] = this.check_segment_validity(
1109
+ mask_labels,
1110
+ mask_probs,
1111
+ k,
1112
+ mask_threshold,
1113
+ overlap_mask_area_threshold
1114
+ )
1115
+
1116
+ if (!mask_exists) {
1117
+ // Nothing to see here
1118
+ continue;
1119
+ }
1120
+
1121
+ // TODO
1122
+ // if (pred_class in stuff_memory_list) {
1123
+ // current_segment_id = stuff_memory_list[pred_class]
1124
+ // } else {
1125
+ // current_segment_id += 1;
1126
+ // }
1127
+ ++current_segment_id;
1128
+
1129
+
1130
+ // Add current object segment to final segmentation map
1131
+ for (let index of mask_k) {
1132
+ segmentation_data[index] = current_segment_id;
1133
+ }
1134
+
1135
+ segments.push({
1136
+ id: current_segment_id,
1137
+ label_id: pred_class,
1138
+ // was_fused: should_fuse, TODO
1139
+ score: pred_scores[k],
1140
+ })
1141
+
1142
+ // TODO
1143
+ // if(should_fuse){
1144
+ // stuff_memory_list[pred_class] = current_segment_id
1145
+ // }
1146
+ }
1147
+
1148
+ return [segmentation, segments];
1149
+ }
1150
+
1151
+ /**
1152
+ * Post-process the model output to generate the final panoptic segmentation.
1153
+ * @param {*} outputs The model output to post process
1154
+ * @param {number} [threshold=0.5] The probability score threshold to keep predicted instance masks.
1155
+ * @param {number} [mask_threshold=0.5] Threshold to use when turning the predicted masks into binary values.
1156
+ * @param {number} [overlap_mask_area_threshold=0.8] The overlap mask area threshold to merge or discard small disconnected parts within each binary instance mask.
1157
+ * @param {Set<number>} [label_ids_to_fuse=null] The labels in this state will have all their instances be fused together.
1158
+ * @param {number[][]} [target_sizes=null] The target sizes to resize the masks to.
1159
+ * @returns {Array<{ segmentation: Tensor, segments_info: Array<{id: number, label_id: number, score: number}>}>}
1160
+ */
1161
+ post_process_panoptic_segmentation(
1162
+ outputs,
1163
+ threshold = 0.5,
1164
+ mask_threshold = 0.5,
1165
+ overlap_mask_area_threshold = 0.8,
1166
+ label_ids_to_fuse = null,
1167
+ target_sizes = null,
1168
+ ) {
1169
+ if (label_ids_to_fuse === null) {
1170
+ console.warn("`label_ids_to_fuse` unset. No instance will be fused.")
1171
+ label_ids_to_fuse = new Set();
1172
+ }
1173
+
1174
+ const class_queries_logits = outputs.logits; // [batch_size, num_queries, num_classes+1]
1175
+ const masks_queries_logits = outputs.pred_masks; // [batch_size, num_queries, height, width]
1176
+
1177
+ const mask_probs = masks_queries_logits.sigmoid() // [batch_size, num_queries, height, width]
1178
+
1179
+ let [batch_size, num_queries, num_labels] = class_queries_logits.dims;
1180
+ num_labels -= 1; // Remove last class (background)
1181
+
1182
+ if (target_sizes !== null && target_sizes.length !== batch_size) {
1183
+ throw Error("Make sure that you pass in as many target sizes as the batch dimension of the logits")
1184
+ }
1185
+
1186
+ let toReturn = [];
1187
+ for (let i = 0; i < batch_size; ++i) {
1188
+ let target_size = target_sizes !== null ? target_sizes[i] : null;
1189
+
1190
+ let class_logits = class_queries_logits[i];
1191
+ let mask_logits = mask_probs[i];
1192
+
1193
+ let [mask_probs_item, pred_scores_item, pred_labels_item] = this.remove_low_and_no_objects(class_logits, mask_logits, threshold, num_labels);
1194
+
1195
+ if (pred_labels_item.length === 0) {
1196
+ // No mask found
1197
+ let [height, width] = target_size ?? mask_logits.dims.slice(-2);
1198
+
1199
+ let segmentation = new Tensor(
1200
+ 'int32',
1201
+ new Int32Array(height * width).fill(-1),
1202
+ [height, width]
1203
+ )
1204
+ toReturn.push({
1205
+ segmentation: segmentation,
1206
+ segments_info: []
1207
+ });
1208
+ continue;
1209
+ }
1210
+
1211
+
1212
+ // Get segmentation map and segment information of batch item
1213
+ let [segmentation, segments] = this.compute_segments(
1214
+ mask_probs_item,
1215
+ pred_scores_item,
1216
+ pred_labels_item,
1217
+ mask_threshold,
1218
+ overlap_mask_area_threshold,
1219
+ label_ids_to_fuse,
1220
+ target_size,
1221
+ )
1222
+
1223
+ toReturn.push({
1224
+ segmentation: segmentation,
1225
+ segments_info: segments
1226
+ })
1227
+ }
1228
+
1229
+ return toReturn;
1230
+ }
1231
+
1232
+ post_process_instance_segmentation() {
1233
+ // TODO
1234
+ throw Error("Not implemented yet");
1235
+ }
1236
+ }
1237
+
1238
+ export class YolosFeatureExtractor extends ImageFeatureExtractor {
1239
+ /** @type {post_process_object_detection} */
1240
+ post_process_object_detection(...args) {
1241
+ return post_process_object_detection(...args);
1242
+ }
1243
+ }
1244
+
1245
+ /**
1246
+ * @typedef {object} SamImageProcessorResult
1247
+ * @property {Tensor} pixel_values
1248
+ * @property {HeightWidth[]} original_sizes
1249
+ * @property {HeightWidth[]} reshaped_input_sizes
1250
+ * @property {Tensor} [input_points]
1251
+ * @property {Tensor} [input_labels]
1252
+ * @property {Tensor} [input_boxes]
1253
+ */
1254
+
1255
+ export class SamImageProcessor extends ImageFeatureExtractor {
1256
+
1257
+ /**
1258
+ *
1259
+ * @param {any} input_points
1260
+ * @param {HeightWidth[]} original_sizes
1261
+ * @param {HeightWidth[]} reshaped_input_sizes
1262
+ * @returns {Tensor}
1263
+ */
1264
+ reshape_input_points(input_points, original_sizes, reshaped_input_sizes, is_bounding_box = false) {
1265
+
1266
+ // Make deep copy to avoid altering user's input
1267
+ input_points = structuredClone(input_points);
1268
+ let shape = calculateDimensions(input_points);
1269
+
1270
+ // TODO: add support for 2D input_points
1271
+ if (shape.length === 3) {
1272
+ // Correct user's input
1273
+ if (!is_bounding_box) {
1274
+ shape = [1, ...shape];
1275
+ }
1276
+ input_points = [input_points];
1277
+ } else if (shape.length !== 4) {
1278
+ throw Error("The input_points must be a 4D tensor of shape `batch_size`, `point_batch_size`, `nb_points_per_image`, `2`.")
1279
+ }
1280
+
1281
+ // Reshape input points
1282
+ for (let i = 0; i < input_points.length; ++i) { // batch_size
1283
+ let originalImageSize = original_sizes[i];
1284
+ let reshapedImageSize = reshaped_input_sizes[i];
1285
+
1286
+ let resizeFactors = [
1287
+ reshapedImageSize[0] / originalImageSize[0],
1288
+ reshapedImageSize[1] / originalImageSize[1]
1289
+ ]
1290
+
1291
+ for (let j = 0; j < input_points[i].length; ++j) { // point_batch_size
1292
+ for (let k = 0; k < input_points[i][j].length; ++k) { // nb_points_per_image
1293
+ for (let w = 0; w < input_points[i][j][k].length; ++w) { // 2 or 4
1294
+ input_points[i][j][k][w] *= resizeFactors[w % 2];
1295
+ }
1296
+ }
1297
+ }
1298
+ }
1299
+
1300
+ return new Tensor(
1301
+ 'float32',
1302
+ Float32Array.from(input_points.flat(Infinity)),
1303
+ shape
1304
+ )
1305
+
1306
+ }
1307
+
1308
+ /**
1309
+ *
1310
+ * @param {any} input_labels
1311
+ * @param {Tensor} input_points
1312
+ * @returns {Tensor}
1313
+ */
1314
+ add_input_labels(input_labels, input_points) {
1315
+ let shape = calculateDimensions(input_labels);
1316
+ if (shape.length === 2) {
1317
+ // Correct user's input
1318
+ shape = [1, ...shape];
1319
+ input_labels = [input_labels];
1320
+ } else if (shape.length !== 3) {
1321
+ throw Error("The input_points must be a 4D tensor of shape `batch_size`, `point_batch_size`, `nb_points_per_image`, `2`.")
1322
+ }
1323
+
1324
+ if (shape.some((x, i) => x !== input_points.dims[i])) {
1325
+ throw Error(`The first ${shape.length} dimensions of 'input_points' and 'input_labels' must be the same.`)
1326
+ }
1327
+ return new Tensor(
1328
+ 'int64',
1329
+ input_labels.flat(Infinity).map(BigInt),
1330
+ shape,
1331
+ )
1332
+ }
1333
+ /**
1334
+ * @param {any[]} images The URL(s) of the image(s) to extract features from.
1335
+ * @param {Object} [options] Additional options for the processor.
1336
+ * @param {any} [options.input_points=null] A 3D or 4D array, representing the input points provided by the user.
1337
+ * - 3D: `[point_batch_size, nb_points_per_image, 2]`. In this case, `batch_size` is assumed to be 1.
1338
+ * - 4D: `[batch_size, point_batch_size, nb_points_per_image, 2]`.
1339
+ * @param {any} [options.input_labels=null] A 2D or 3D array, representing the input labels for the points, used by the prompt encoder to encode the prompt.
1340
+ * - 2D: `[point_batch_size, nb_points_per_image]`. In this case, `batch_size` is assumed to be 1.
1341
+ * - 3D: `[batch_size, point_batch_size, nb_points_per_image]`.
1342
+ * @param {number[][][]} [options.input_boxes=null] A 3D array of shape `(batch_size, num_boxes, 4)`, representing the input boxes provided by the user.
1343
+ * This is used by the prompt encoder to encode the prompt. Generally yields to much better generated masks.
1344
+ * The processor will generate a tensor, with each dimension corresponding respectively to the image batch size,
1345
+ * the number of boxes per image and the coordinates of the top left and botton right point of the box.
1346
+ * In the order (`x1`, `y1`, `x2`, `y2`):
1347
+ * - `x1`: the x coordinate of the top left point of the input box
1348
+ * - `y1`: the y coordinate of the top left point of the input box
1349
+ * - `x2`: the x coordinate of the bottom right point of the input box
1350
+ * - `y2`: the y coordinate of the bottom right point of the input box
1351
+ * @returns {Promise<SamImageProcessorResult>}
1352
+ */
1353
+ async _call(images, {
1354
+ input_points = null,
1355
+ input_labels = null,
1356
+ input_boxes = null
1357
+ } = {}) {
1358
+ // TODO allow user to use preprocessed images
1359
+ /** @type {SamImageProcessorResult} */
1360
+ const processed = await super._call(images);
1361
+
1362
+ if (input_points) {
1363
+ processed.input_points = this.reshape_input_points(
1364
+ input_points, processed.original_sizes, processed.reshaped_input_sizes
1365
+ );
1366
+ }
1367
+
1368
+ if (input_labels) {
1369
+ if (!processed.input_points) {
1370
+ throw Error("`input_points` must be provided if `input_labels` are provided.")
1371
+ }
1372
+ processed.input_labels = this.add_input_labels(input_labels, processed.input_points);
1373
+ }
1374
+
1375
+ if (input_boxes) {
1376
+ processed.input_boxes = this.reshape_input_points(
1377
+ input_boxes, processed.original_sizes, processed.reshaped_input_sizes, true,
1378
+ );
1379
+ }
1380
+
1381
+ return processed;
1382
+ }
1383
+
1384
+ /**
1385
+ * Remove padding and upscale masks to the original image size.
1386
+ * @param {Tensor} masks Batched masks from the mask_decoder in (batch_size, num_channels, height, width) format.
1387
+ * @param {[number, number][]} original_sizes The original sizes of each image before it was resized to the model's expected input shape, in (height, width) format.
1388
+ * @param {[number, number][]} reshaped_input_sizes The size of each image as it is fed to the model, in (height, width) format. Used to remove padding.
1389
+ * @param {Object} options Optional parameters for post-processing.
1390
+ * @param {number} [options.mask_threshold] The threshold to use for binarizing the masks.
1391
+ * @param {boolean} [options.binarize] Whether to binarize the masks.
1392
+ * @param {Object} [options.pad_size] The target size the images were padded to before being passed to the model. If `null`, the target size is assumed to be the processor's `pad_size`.
1393
+ * @param {number} [options.pad_size.height] The height the images were padded to.
1394
+ * @param {number} [options.pad_size.width] The width the images were padded to.
1395
+ * @returns {Promise<Tensor[]>} Batched masks in batch_size, num_channels, height, width) format, where (height, width) is given by original_size.
1396
+ */
1397
+ async post_process_masks(masks, original_sizes, reshaped_input_sizes, {
1398
+ mask_threshold = 0.0,
1399
+ binarize = true,
1400
+ pad_size = null,
1401
+ } = {}) {
1402
+ // masks: [1, 1, 3, 256, 256]
1403
+
1404
+ const output_masks = [];
1405
+
1406
+ pad_size = pad_size ?? this.pad_size;
1407
+
1408
+ /** @type {[number, number]} */
1409
+ const target_image_size = [pad_size.height, pad_size.width];
1410
+
1411
+ for (let i = 0; i < original_sizes.length; ++i) {
1412
+ const original_size = original_sizes[i];
1413
+ const reshaped_input_size = reshaped_input_sizes[i];
1414
+
1415
+ // Upscale mask to padded size
1416
+ let interpolated_mask = (await interpolate_4d(
1417
+ masks[i],
1418
+ { mode: 'bilinear', size: target_image_size }
1419
+ ));
1420
+
1421
+ // Crop mask
1422
+ interpolated_mask = interpolated_mask.slice(null, null, [0, reshaped_input_size[0]], [0, reshaped_input_size[1]]);
1423
+
1424
+ // Downscale mask
1425
+ interpolated_mask = (await interpolate_4d(
1426
+ interpolated_mask,
1427
+ { mode: 'bilinear', size: original_size }
1428
+ ));
1429
+
1430
+ if (binarize) {
1431
+ const data = interpolated_mask.data;
1432
+ const binarizedMaskData = new Uint8Array(data.length);
1433
+ for (let i = 0; i < data.length; ++i) {
1434
+ if (data[i] > mask_threshold) {
1435
+ binarizedMaskData[i] = 1;
1436
+ }
1437
+ }
1438
+ interpolated_mask = new Tensor(
1439
+ 'bool',
1440
+ binarizedMaskData,
1441
+ interpolated_mask.dims
1442
+ )
1443
+ }
1444
+
1445
+ output_masks.push(interpolated_mask);
1446
+ }
1447
+
1448
+ return output_masks;
1449
+ }
1450
+
1451
+ /**
1452
+ * Generates a list of crop boxes of different sizes. Each layer has (2**i)**2 boxes for the ith layer.
1453
+ * @param {RawImage} image Input original image
1454
+ * @param {number} target_size Target size of the resized image
1455
+ * @param {Object} options Options for generating crop boxes
1456
+ * @param {number} [options.crop_n_layers] If >0, mask prediction will be run again on crops of the image.
1457
+ * Sets the number of layers to run, where each layer has 2**i_layer number of image crops.
1458
+ * @param {number} [options.overlap_ratio] Sets the degree to which crops overlap. In the first crop layer,
1459
+ * crops will overlap by this fraction of the image length. Later layers with more crops scale down this overlap.
1460
+ * @param {number} [options.points_per_crop] Number of points to sample from each crop.
1461
+ * @param {number} [options.crop_n_points_downscale_factor] The number of points-per-side sampled in layer n is
1462
+ * scaled down by crop_n_points_downscale_factor**n.
1463
+ * @returns {Object} An object containing the crop boxes, number of points per crop, cropped images, and input labels.
1464
+ */
1465
+ generate_crop_boxes(image, target_size, {
1466
+ crop_n_layers = 0,
1467
+ overlap_ratio = 512 / 1500,
1468
+ points_per_crop = 32,
1469
+ crop_n_points_downscale_factor = 1,
1470
+ } = {}) {
1471
+ // TODO: Implement
1472
+ // return { crop_boxes, points_per_crop, cropped_images, input_labels }
1473
+ }
1474
+ }
1475
+
1476
+ export class Swin2SRImageProcessor extends ImageFeatureExtractor {
1477
+ pad_image(pixelData, imgDims, padSize, options = {}) {
1478
+ // NOTE: In this case, `padSize` represents the size of the sliding window for the local attention.
1479
+ // In other words, the image is padded so that its width and height are multiples of `padSize`.
1480
+ const [imageHeight, imageWidth, imageChannels] = imgDims;
1481
+
1482
+ return super.pad_image(pixelData, imgDims, {
1483
+ // NOTE: For Swin2SR models, the original python implementation adds padding even when the image's width/height is already
1484
+ // a multiple of `pad_size`. However, this is most likely a bug (PR: https://github.com/mv-lab/swin2sr/pull/19).
1485
+ // For this reason, we only add padding when the image's width/height is not a multiple of `pad_size`.
1486
+ width: imageWidth + (padSize - imageWidth % padSize) % padSize,
1487
+ height: imageHeight + (padSize - imageHeight % padSize) % padSize,
1488
+ }, {
1489
+ mode: 'symmetric',
1490
+ center: false,
1491
+ constant_values: -1,
1492
+ ...options,
1493
+ })
1494
+ }
1495
+ }
1496
+
1497
+ export class VitMatteImageProcessor extends ImageFeatureExtractor {
1498
+ /**
1499
+ * Calls the feature extraction process on an array of images, preprocesses
1500
+ * each image, and concatenates the resulting features into a single Tensor.
1501
+ * @param {RawImage[]} images The image(s) to extract features from.
1502
+ * @param {RawImage[]} trimaps The trimaps(s) to extract features from.
1503
+ * @returns {Promise<ImageFeatureExtractorResult>} An object containing the concatenated pixel values of the preprocessed images.
1504
+ */
1505
+ async _call(images, trimaps) {
1506
+ if (!Array.isArray(images)) {
1507
+ images = [images];
1508
+ }
1509
+ if (!Array.isArray(trimaps)) {
1510
+ trimaps = [trimaps];
1511
+ }
1512
+
1513
+ const imageData = await Promise.all(images.map(x => this.preprocess(x)));
1514
+ const trimapData = await Promise.all(trimaps.map(x => this.preprocess(x, {
1515
+ do_normalize: false,
1516
+ do_convert_rgb: false,
1517
+ do_convert_grayscale: true,
1518
+ })));
1519
+
1520
+
1521
+ // Stack pixel values
1522
+ const pixel_values = stack(imageData.map(
1523
+ // Concatenate images and trimaps
1524
+ (x, i) => cat([x.pixel_values, trimapData[i].pixel_values], 0)
1525
+ ), 0);
1526
+
1527
+ return {
1528
+ pixel_values,
1529
+
1530
+ // Original sizes of images
1531
+ original_sizes: imageData.map(x => x.original_size),
1532
+
1533
+ // Reshaped sizes of images, before padding or cropping
1534
+ reshaped_input_sizes: imageData.map(x => x.reshaped_input_size),
1535
+ }
1536
+ }
1537
+ }
1538
+
1539
+ export class WhisperFeatureExtractor extends FeatureExtractor {
1540
+
1541
+ constructor(config) {
1542
+ super(config);
1543
+
1544
+ // Prefer given `mel_filters` from preprocessor_config.json, or calculate them if they don't exist.
1545
+ this.config.mel_filters ??= mel_filter_bank(
1546
+ Math.floor(1 + this.config.n_fft / 2), // num_frequency_bins
1547
+ this.config.feature_size, // num_mel_filters
1548
+ 0.0, // min_frequency
1549
+ 8000.0, // max_frequency
1550
+ this.config.sampling_rate, // sampling_rate
1551
+ "slaney", // norm
1552
+ "slaney", // mel_scale
1553
+ );
1554
+
1555
+ this.window = window_function(this.config.n_fft, 'hann');
1556
+ }
1557
+
1558
+ /**
1559
+ * Computes the log-Mel spectrogram of the provided audio waveform.
1560
+ * @param {Float32Array|Float64Array} waveform The audio waveform to process.
1561
+ * @returns {Promise<Tensor>} An object containing the log-Mel spectrogram data as a Float32Array and its dimensions as an array of numbers.
1562
+ */
1563
+ async _extract_fbank_features(waveform) {
1564
+ const features = await spectrogram(
1565
+ waveform,
1566
+ this.window, // window
1567
+ this.config.n_fft, // frame_length
1568
+ this.config.hop_length, // hop_length
1569
+ {
1570
+ power: 2.0,
1571
+ mel_filters: this.config.mel_filters,
1572
+ log_mel: 'log10',
1573
+
1574
+ // Custom
1575
+ max_num_frames: this.config.nb_max_frames, // 3000
1576
+ }
1577
+ )
1578
+
1579
+ const data = features.data;
1580
+ const maxValue = max(data)[0];
1581
+
1582
+ for (let i = 0; i < data.length; ++i) {
1583
+ data[i] = (Math.max(data[i], maxValue - 8.0) + 4.0) / 4.0;
1584
+ }
1585
+
1586
+ return features;
1587
+ }
1588
+
1589
+ /**
1590
+ * Asynchronously extracts features from a given audio using the provided configuration.
1591
+ * @param {Float32Array|Float64Array} audio The audio data as a Float32Array/Float64Array.
1592
+ * @returns {Promise<{ input_features: Tensor }>} A Promise resolving to an object containing the extracted input features as a Tensor.
1593
+ */
1594
+ async _call(audio) {
1595
+ validate_audio_inputs(audio, 'WhisperFeatureExtractor');
1596
+
1597
+ let waveform;
1598
+ if (audio.length > this.config.n_samples) {
1599
+ console.warn(
1600
+ "Attempting to extract features for audio longer than 30 seconds. " +
1601
+ "If using a pipeline to extract transcript from a long audio clip, " +
1602
+ "remember to specify `chunk_length_s` and/or `stride_length_s`."
1603
+ );
1604
+ waveform = audio.slice(0, this.config.n_samples);
1605
+ } else {
1606
+ // pad with zeros
1607
+ waveform = new Float32Array(this.config.n_samples);
1608
+ waveform.set(audio);
1609
+ }
1610
+
1611
+ const features = await this._extract_fbank_features(waveform);
1612
+
1613
+ return {
1614
+ input_features: features.unsqueeze_(0)
1615
+ };
1616
+ }
1617
+ }
1618
+
1619
+ export class Wav2Vec2FeatureExtractor extends FeatureExtractor {
1620
+
1621
+ /**
1622
+ * @param {Float32Array} input_values
1623
+ * @returns {Float32Array}
1624
+ */
1625
+ _zero_mean_unit_var_norm(input_values) {
1626
+ // TODO support batch?
1627
+ const sum = input_values.reduce((a, b) => a + b, 0);
1628
+ const mean = sum / input_values.length;
1629
+ const variance = input_values.reduce((a, b) => a + (b - mean) ** 2, 0) / input_values.length;
1630
+ return input_values.map(x => (x - mean) / Math.sqrt(variance + 1e-7));
1631
+ }
1632
+
1633
+ /**
1634
+ * Asynchronously extracts features from a given audio using the provided configuration.
1635
+ * @param {Float32Array|Float64Array} audio The audio data as a Float32Array/Float64Array.
1636
+ * @returns {Promise<{ input_values: Tensor; attention_mask: Tensor }>} A Promise resolving to an object containing the extracted input features and attention mask as Tensors.
1637
+ */
1638
+ async _call(audio) {
1639
+ validate_audio_inputs(audio, 'Wav2Vec2FeatureExtractor');
1640
+
1641
+ if (audio instanceof Float64Array) {
1642
+ audio = new Float32Array(audio);
1643
+ }
1644
+
1645
+ let input_values = audio;
1646
+
1647
+ // zero-mean and unit-variance normalization
1648
+ if (this.config.do_normalize) {
1649
+ input_values = this._zero_mean_unit_var_norm(input_values);
1650
+ }
1651
+
1652
+ // TODO: allow user to pass in attention mask
1653
+ const shape = [1, input_values.length];
1654
+ return {
1655
+ input_values: new Tensor('float32', input_values, shape),
1656
+ attention_mask: new Tensor('int64', new BigInt64Array(input_values.length).fill(1n), shape)
1657
+ };
1658
+ }
1659
+ }
1660
+
1661
+ export class SeamlessM4TFeatureExtractor extends FeatureExtractor {
1662
+
1663
+ constructor(config) {
1664
+ super(config);
1665
+
1666
+ const sampling_rate = this.config.sampling_rate;
1667
+ const mel_filters = mel_filter_bank(
1668
+ 256, // num_frequency_bins
1669
+ this.config.num_mel_bins, // num_mel_filters
1670
+ 20, // min_frequency
1671
+ Math.floor(sampling_rate / 2), // max_frequency
1672
+ sampling_rate, // sampling_rate
1673
+ null, // norm
1674
+ "kaldi", // mel_scale
1675
+ true, // triangularize_in_mel_space
1676
+ );
1677
+
1678
+ // Do padding:
1679
+ for (let i = 0; i < mel_filters.length; ++i) {
1680
+ mel_filters[i].push(0);
1681
+ }
1682
+ this.mel_filters = mel_filters;
1683
+
1684
+ this.window = window_function(400, 'povey', {
1685
+ periodic: false,
1686
+ })
1687
+ }
1688
+
1689
+ /**
1690
+ * Computes the log-Mel spectrogram of the provided audio waveform.
1691
+ * @param {Float32Array|Float64Array} waveform The audio waveform to process.
1692
+ * @param {number} max_length The maximum number of frames to return.
1693
+ * @returns {Promise<Tensor>} An object containing the log-Mel spectrogram data as a Float32Array and its dimensions as an array of numbers.
1694
+ */
1695
+ async _extract_fbank_features(waveform, max_length) {
1696
+ // NOTE: We don't pad/truncate since that is passed in as `max_num_frames`
1697
+
1698
+ // Kaldi compliance: 16-bit signed integers
1699
+ // 32768 == 2 ** 15
1700
+ waveform = waveform.map((/** @type {number} */ x) => x * 32768)
1701
+
1702
+ return spectrogram(
1703
+ waveform,
1704
+ this.window, // window
1705
+ 400, // frame_length
1706
+ 160, // hop_length
1707
+ {
1708
+ fft_length: 512,
1709
+ power: 2.0,
1710
+ center: false,
1711
+ preemphasis: 0.97,
1712
+ mel_filters: this.mel_filters,
1713
+ log_mel: 'log',
1714
+ mel_floor: 1.192092955078125e-07,
1715
+ remove_dc_offset: true,
1716
+
1717
+ // Custom
1718
+ max_num_frames: max_length,
1719
+ transpose: true,
1720
+ }
1721
+ )
1722
+ }
1723
+
1724
+ /**
1725
+ * Asynchronously extracts features from a given audio using the provided configuration.
1726
+ * @param {Float32Array|Float64Array} audio The audio data as a Float32Array/Float64Array.
1727
+ * @param {Object} options Optional parameters for feature extraction.
1728
+ * @param {boolean} [options.padding=true] Whether to pad the sequence to a multiple of `pad_to_multiple_of`.
1729
+ * @param {number} [options.pad_to_multiple_of=2] The number to pad the sequence to a multiple of.
1730
+ * @param {boolean} [options.do_normalize_per_mel_bins=true] Whether or not to zero-mean unit-variance normalize the input per mel-channel.
1731
+ * @param {boolean} [options.return_attention_mask=true] Whether to return the attention mask.
1732
+ * @returns {Promise<{ input_features: Tensor, attention_mask?: Tensor }>} A Promise resolving to an object containing the extracted input features and attention masks as Tensors.
1733
+ */
1734
+ async _call(audio, {
1735
+ padding = true,
1736
+ pad_to_multiple_of = 2,
1737
+ do_normalize_per_mel_bins = true,
1738
+ return_attention_mask = true,
1739
+ } = {}) {
1740
+ validate_audio_inputs(audio, 'SeamlessM4TFeatureExtractor');
1741
+
1742
+ let features = await this._extract_fbank_features(audio, this.config.max_length);
1743
+
1744
+ if (do_normalize_per_mel_bins) {
1745
+ const [num_features, feature_size] = features.dims;
1746
+ const data = features.data;
1747
+ for (let i = 0; i < feature_size; ++i) {
1748
+ let sum = 0;
1749
+ for (let j = 0; j < num_features; ++j) {
1750
+ sum += data[j * feature_size + i];
1751
+ }
1752
+
1753
+ const mean = sum / num_features;
1754
+
1755
+ let variance = 0;
1756
+ for (let j = 0; j < num_features; ++j) {
1757
+ variance += (data[j * feature_size + i] - mean) ** 2;
1758
+ }
1759
+ variance /= num_features - 1; // NOTE: We use ddof=1
1760
+
1761
+ const std = Math.sqrt(variance + 1e-7);
1762
+ for (let j = 0; j < num_features; ++j) {
1763
+ const index = j * feature_size + i;
1764
+ data[index] = (data[index] - mean) / std;
1765
+ }
1766
+ }
1767
+ }
1768
+
1769
+ let padded_attention_mask;
1770
+ if (padding) {
1771
+ const [num_frames, num_channels] = features.dims;
1772
+ const data = /** @type {Float32Array} */(features.data);
1773
+
1774
+ const pad_size = num_frames % pad_to_multiple_of;
1775
+ if (pad_size > 0) {
1776
+ const padded_data = new Float32Array(num_channels * (num_frames + pad_size));
1777
+ padded_data.set(data)
1778
+ padded_data.fill(this.config.padding_value, data.length)
1779
+
1780
+ const numPaddedFrames = num_frames + pad_size;
1781
+ features = new Tensor(
1782
+ features.type,
1783
+ padded_data,
1784
+ [numPaddedFrames, num_channels],
1785
+ )
1786
+
1787
+ if (return_attention_mask) {
1788
+ padded_attention_mask = new Tensor(
1789
+ 'int64',
1790
+ new BigInt64Array(numPaddedFrames),
1791
+ [1, numPaddedFrames],
1792
+ )
1793
+ padded_attention_mask.data.fill(1n, 0, num_frames);
1794
+ }
1795
+ }
1796
+ }
1797
+
1798
+ const [num_frames, num_channels] = features.dims;
1799
+
1800
+ const stride = this.config.stride;
1801
+ const remainder = num_frames % stride;
1802
+ if (remainder !== 0) {
1803
+ throw new Error(`The number of frames (${num_frames}) must be a multiple of the stride (${stride}).`)
1804
+ }
1805
+
1806
+ const input_features = features.view(
1807
+ 1,
1808
+ Math.floor(num_frames / stride),
1809
+ num_channels * stride,
1810
+ );
1811
+
1812
+ const result = { input_features }
1813
+
1814
+ if (return_attention_mask) {
1815
+ const reshapedNumFrames = input_features.dims[1];
1816
+
1817
+ const attention_mask_data = new BigInt64Array(reshapedNumFrames);
1818
+
1819
+ if (padded_attention_mask) {
1820
+ const padded_attention_mask_data = padded_attention_mask.data;
1821
+ for (let i = 1, j = 0; i < num_frames; i += stride, ++j) {
1822
+ attention_mask_data[j] = padded_attention_mask_data[i];
1823
+ }
1824
+ } else {
1825
+ attention_mask_data.fill(1n);
1826
+ }
1827
+ result.attention_mask = new Tensor(
1828
+ 'int64',
1829
+ attention_mask_data,
1830
+ [1, reshapedNumFrames],
1831
+ );
1832
+ }
1833
+
1834
+ return result;
1835
+ }
1836
+ }
1837
+
1838
+ export class ASTFeatureExtractor extends FeatureExtractor {
1839
+
1840
+
1841
+ constructor(config) {
1842
+ super(config);
1843
+
1844
+ const sampling_rate = this.config.sampling_rate;
1845
+ const mel_filters = mel_filter_bank(
1846
+ 256, // num_frequency_bins
1847
+ this.config.num_mel_bins, // num_mel_filters
1848
+ 20, // min_frequency
1849
+ Math.floor(sampling_rate / 2), // max_frequency
1850
+ sampling_rate, // sampling_rate
1851
+ null, // norm
1852
+ "kaldi", // mel_scale
1853
+ true, // triangularize_in_mel_space
1854
+ );
1855
+
1856
+ // Do padding:
1857
+ for (let i = 0; i < mel_filters.length; ++i) {
1858
+ mel_filters[i].push(0);
1859
+ }
1860
+ this.mel_filters = mel_filters;
1861
+
1862
+ this.window = window_function(400, 'hann', {
1863
+ periodic: false,
1864
+ })
1865
+
1866
+ this.mean = this.config.mean;
1867
+ this.std = this.config.std;
1868
+ }
1869
+
1870
+ /**
1871
+ * Computes the log-Mel spectrogram of the provided audio waveform.
1872
+ * @param {Float32Array|Float64Array} waveform The audio waveform to process.
1873
+ * @param {number} max_length The maximum number of frames to return.
1874
+ * @returns {Promise<Tensor>} An object containing the log-Mel spectrogram data as a Float32Array and its dimensions as an array of numbers.
1875
+ */
1876
+ async _extract_fbank_features(waveform, max_length) {
1877
+ // NOTE: We don't pad/truncate since that is passed in as `max_num_frames`
1878
+ return spectrogram(
1879
+ waveform,
1880
+ this.window, // window
1881
+ 400, // frame_length
1882
+ 160, // hop_length
1883
+ {
1884
+ fft_length: 512,
1885
+ power: 2.0,
1886
+ center: false,
1887
+ preemphasis: 0.97,
1888
+ mel_filters: this.mel_filters,
1889
+ log_mel: 'log',
1890
+ mel_floor: 1.192092955078125e-07,
1891
+ remove_dc_offset: true,
1892
+
1893
+ // Custom
1894
+ max_num_frames: max_length,
1895
+ transpose: true,
1896
+ }
1897
+ )
1898
+ }
1899
+
1900
+
1901
+ /**
1902
+ * Asynchronously extracts features from a given audio using the provided configuration.
1903
+ * @param {Float32Array|Float64Array} audio The audio data as a Float32Array/Float64Array.
1904
+ * @returns {Promise<{ input_values: Tensor }>} A Promise resolving to an object containing the extracted input features as a Tensor.
1905
+ */
1906
+ async _call(audio) {
1907
+ validate_audio_inputs(audio, 'ASTFeatureExtractor');
1908
+
1909
+ const features = await this._extract_fbank_features(audio, this.config.max_length);
1910
+ if (this.config.do_normalize) {
1911
+ // Normalize the input audio spectrogram to have mean=0, std=0.5
1912
+ const denom = this.std * 2;
1913
+ const features_data = features.data;
1914
+ for (let i = 0; i < features_data.length; ++i) {
1915
+ features_data[i] = (features_data[i] - this.mean) / denom;
1916
+ }
1917
+ }
1918
+
1919
+ return {
1920
+ input_values: features.unsqueeze_(0)
1921
+ };
1922
+ }
1923
+ }
1924
+
1925
+ export class ClapFeatureExtractor extends FeatureExtractor {
1926
+
1927
+ constructor(config) {
1928
+ super(config);
1929
+
1930
+ this.mel_filters = mel_filter_bank(
1931
+ this.config.nb_frequency_bins, // num_frequency_bins
1932
+ this.config.feature_size, // num_mel_filters
1933
+ this.config.frequency_min, // min_frequency
1934
+ this.config.frequency_max, // max_frequency
1935
+ this.config.sampling_rate, // sampling_rate
1936
+ null, // norm
1937
+ "htk", // mel_scale
1938
+ );
1939
+
1940
+ this.mel_filters_slaney = mel_filter_bank(
1941
+ this.config.nb_frequency_bins, // num_frequency_bins
1942
+ this.config.feature_size, // num_mel_filters
1943
+ this.config.frequency_min, // min_frequency
1944
+ this.config.frequency_max, // max_frequency
1945
+ this.config.sampling_rate, // sampling_rate
1946
+ "slaney", // norm
1947
+ "slaney", // mel_scale
1948
+ );
1949
+
1950
+ this.window = window_function(this.config.fft_window_size, 'hann')
1951
+
1952
+ }
1953
+
1954
+
1955
+ /**
1956
+ * Extracts the mel spectrogram and prepares it for the mode based on the `truncation` and `padding` arguments.
1957
+ *
1958
+ * Four different path are possible:
1959
+ * - `truncation="fusion"` and the length of the waveform is greater than the max length: the mel spectrogram
1960
+ * will be computed on the entire audio. 3 random crops and a dowsampled version of the full mel spectrogram
1961
+ * are then stacked together. They will later be used for `feature_fusion`.
1962
+ * - `truncation="rand_trunc"` and the length of the waveform is smaller than the max length: the audio is
1963
+ * padded based on `padding`.
1964
+ * - `truncation="fusion"` and the length of the waveform is smaller than the max length: the audio is padded
1965
+ * based on `padding`, and is repeated `4` times.
1966
+ * - `truncation="rand_trunc"` and the length of the waveform is greater than the max length: the mel
1967
+ * spectrogram will be computed on a random crop of the waveform.
1968
+ *
1969
+ * @param {Float32Array|Float64Array} waveform The input waveform.
1970
+ * @param {number} max_length The maximum length of the waveform.
1971
+ * @param {string} truncation The truncation strategy to use.
1972
+ * @param {string} padding The padding strategy to use.
1973
+ * @returns {Promise<Tensor>} An object containing the mel spectrogram data as a Float32Array, its dimensions as an array of numbers, and a boolean indicating whether the waveform was longer than the max length.
1974
+ * @private
1975
+ */
1976
+ async _get_input_mel(waveform, max_length, truncation, padding) {
1977
+
1978
+ /** @type {Tensor} */
1979
+ let input_mel;
1980
+ let longer = false;
1981
+ const diff = waveform.length - max_length;
1982
+ if (diff > 0) {
1983
+ if (truncation === 'rand_trunc') {
1984
+ longer = true;
1985
+ const idx = Math.floor(Math.random() * (diff + 1));
1986
+ waveform = waveform.subarray(idx, idx + max_length);
1987
+
1988
+ input_mel = await this._extract_fbank_features(waveform, this.mel_filters_slaney, this.config.nb_max_samples);
1989
+ } else {
1990
+ // TODO implement fusion strategy
1991
+ throw new Error(`Truncation strategy "${truncation}" not implemented`)
1992
+ }
1993
+ } else {
1994
+ if (diff < 0) {
1995
+ let padded = new Float64Array(max_length); // already padded with zeros
1996
+ padded.set(waveform);
1997
+
1998
+ if (padding === 'repeat') {
1999
+ for (let i = waveform.length; i < max_length; i += waveform.length) {
2000
+ padded.set(waveform.subarray(0, Math.min(waveform.length, max_length - i)), i);
2001
+ }
2002
+ } else if (padding === 'repeatpad') {
2003
+ for (let i = waveform.length; i < -diff; i += waveform.length) {
2004
+ padded.set(waveform, i);
2005
+ }
2006
+ }
2007
+ waveform = padded;
2008
+ }
2009
+
2010
+ if (truncation === 'fusion') {
2011
+ throw new Error(`Truncation strategy "${truncation}" not implemented`)
2012
+ }
2013
+
2014
+ input_mel = await this._extract_fbank_features(waveform, this.mel_filters_slaney, this.config.nb_max_samples);
2015
+ }
2016
+
2017
+ return input_mel.unsqueeze_(0);
2018
+ }
2019
+
2020
+ /**
2021
+ * Compute the log-mel spectrogram of the provided `waveform` using the Hann window.
2022
+ * In CLAP, two different filter banks are used depending on the truncation pattern:
2023
+ * - `self.mel_filters`: they correspond to the default parameters of `torchaudio` which can be obtained from
2024
+ * calling `torchaudio.transforms.MelSpectrogram().mel_scale.fb`. These filters are used when `truncation`
2025
+ * is set to `"fusion"`.
2026
+ * - `self.mel_filteres_slaney` : they correspond to the default parameters of `librosa` which used
2027
+ * `librosa.filters.mel` when computing the mel spectrogram. These filters were only used in the original
2028
+ * implementation when the truncation mode is not `"fusion"`.
2029
+ *
2030
+ * @param {Float32Array|Float64Array} waveform The audio waveform to process.
2031
+ * @param {number[][]} mel_filters The mel filters to use.
2032
+ * @param {number} [max_length=null] The maximum number of frames to return.
2033
+ * @returns {Promise<Tensor>} An object containing the log-Mel spectrogram data as a Float32Array and its dimensions as an array of numbers.
2034
+ */
2035
+ async _extract_fbank_features(waveform, mel_filters, max_length = null) {
2036
+ // NOTE: We don't pad/truncate since that is passed in as `max_num_frames`
2037
+ return spectrogram(
2038
+ waveform,
2039
+ this.window, // window
2040
+ this.config.fft_window_size, // frame_length
2041
+ this.config.hop_length, // hop_length
2042
+ {
2043
+ power: 2.0,
2044
+ mel_filters,
2045
+ log_mel: 'dB',
2046
+
2047
+ // Custom
2048
+ max_num_frames: max_length,
2049
+ do_pad: false,
2050
+ transpose: true,
2051
+ }
2052
+ )
2053
+ }
2054
+
2055
+
2056
+ /**
2057
+ * Asynchronously extracts features from a given audio using the provided configuration.
2058
+ * @param {Float32Array|Float64Array} audio The audio data as a Float32Array/Float64Array.
2059
+ * @returns {Promise<{ input_features: Tensor }>} A Promise resolving to an object containing the extracted input features as a Tensor.
2060
+ */
2061
+ async _call(audio, {
2062
+ max_length = null,
2063
+ } = {}) {
2064
+ validate_audio_inputs(audio, 'ClapFeatureExtractor');
2065
+
2066
+ // convert to mel spectrogram, truncate and pad if needed.
2067
+ const padded_inputs = await this._get_input_mel(
2068
+ audio,
2069
+ max_length ?? this.config.nb_max_samples,
2070
+ this.config.truncation,
2071
+ this.config.padding,
2072
+ );
2073
+
2074
+ return {
2075
+ input_features: padded_inputs.unsqueeze_(0),
2076
+ }
2077
+ }
2078
+ }
2079
+
2080
+
2081
+ export class PyAnnoteFeatureExtractor extends FeatureExtractor {
2082
+ /**
2083
+ * Asynchronously extracts features from a given audio using the provided configuration.
2084
+ * @param {Float32Array|Float64Array} audio The audio data as a Float32Array/Float64Array.
2085
+ * @returns {Promise<{ input_values: Tensor; }>} The extracted input features.
2086
+ */
2087
+ async _call(audio) {
2088
+ validate_audio_inputs(audio, 'PyAnnoteFeatureExtractor');
2089
+
2090
+ if (audio instanceof Float64Array) {
2091
+ audio = new Float32Array(audio);
2092
+ }
2093
+
2094
+ const shape = [
2095
+ 1, /* batch_size */
2096
+ 1, /* num_channels */
2097
+ audio.length, /* num_samples */
2098
+ ];
2099
+ return {
2100
+ input_values: new Tensor('float32', audio, shape),
2101
+ };
2102
+ }
2103
+
2104
+ /**
2105
+ * NOTE: Can return fractional values. `Math.ceil` will ensure correct value.
2106
+ * @param {number} samples The number of frames in the audio.
2107
+ * @returns {number} The number of frames in the audio.
2108
+ */
2109
+ samples_to_frames(samples) {
2110
+ return ((samples - this.config.offset) / this.config.step);
2111
+ }
2112
+
2113
+ /**
2114
+ * Post-processes the speaker diarization logits output by the model.
2115
+ * @param {Tensor} logits The speaker diarization logits output by the model.
2116
+ * @param {number} num_samples Number of samples in the input audio.
2117
+ * @returns {Array<Array<{ id: number, start: number, end: number, confidence: number }>>} The post-processed speaker diarization results.
2118
+ */
2119
+ post_process_speaker_diarization(logits, num_samples) {
2120
+ const ratio = (
2121
+ num_samples / this.samples_to_frames(num_samples)
2122
+ ) / this.config.sampling_rate;
2123
+
2124
+ const results = [];
2125
+ for (const scores of logits.tolist()) {
2126
+ const accumulated_segments = [];
2127
+
2128
+ let current_speaker = -1;
2129
+ for (let i = 0; i < scores.length; ++i) {
2130
+ const probabilities = softmax(scores[i]);
2131
+ const [score, id] = max(probabilities);
2132
+ const [start, end] = [i, i + 1];
2133
+
2134
+ if (id !== current_speaker) {
2135
+ // Speaker has changed
2136
+ current_speaker = id;
2137
+ accumulated_segments.push({ id, start, end, score });
2138
+ } else {
2139
+ // Continue the current segment
2140
+ accumulated_segments.at(-1).end = end;
2141
+ accumulated_segments.at(-1).score += score;
2142
+ }
2143
+ }
2144
+
2145
+ results.push(accumulated_segments.map(
2146
+ // Convert frame-space to time-space
2147
+ // and compute the confidence
2148
+ ({ id, start, end, score }) => ({
2149
+ id,
2150
+ start: start * ratio,
2151
+ end: end * ratio,
2152
+ confidence: score / (end - start),
2153
+ })
2154
+ ));
2155
+ }
2156
+ return results;
2157
+ }
2158
+
2159
+ }
2160
+
2161
+ export class WeSpeakerFeatureExtractor extends FeatureExtractor {
2162
+
2163
+ constructor(config) {
2164
+ super(config);
2165
+
2166
+ const sampling_rate = this.config.sampling_rate;
2167
+ const mel_filters = mel_filter_bank(
2168
+ 256, // num_frequency_bins
2169
+ this.config.num_mel_bins, // num_mel_filters
2170
+ 20, // min_frequency
2171
+ Math.floor(sampling_rate / 2), // max_frequency
2172
+ sampling_rate, // sampling_rate
2173
+ null, // norm
2174
+ "kaldi", // mel_scale
2175
+ true, // triangularize_in_mel_space
2176
+ );
2177
+
2178
+ // Do padding:
2179
+ for (let i = 0; i < mel_filters.length; ++i) {
2180
+ mel_filters[i].push(0);
2181
+ }
2182
+ this.mel_filters = mel_filters;
2183
+
2184
+ this.window = window_function(400, 'hamming', {
2185
+ periodic: false,
2186
+ })
2187
+ this.min_num_frames = this.config.min_num_frames;
2188
+ }
2189
+
2190
+ /**
2191
+ * Computes the log-Mel spectrogram of the provided audio waveform.
2192
+ * @param {Float32Array|Float64Array} waveform The audio waveform to process.
2193
+ * @returns {Promise<Tensor>} An object containing the log-Mel spectrogram data as a Float32Array and its dimensions as an array of numbers.
2194
+ */
2195
+ async _extract_fbank_features(waveform) {
2196
+ // Kaldi compliance: 16-bit signed integers
2197
+ // 32768 == 2 ** 15
2198
+ waveform = waveform.map((/** @type {number} */ x) => x * 32768)
2199
+
2200
+ return spectrogram(
2201
+ waveform,
2202
+ this.window, // window
2203
+ 400, // frame_length
2204
+ 160, // hop_length
2205
+ {
2206
+ fft_length: 512,
2207
+ power: 2.0,
2208
+ center: false,
2209
+ preemphasis: 0.97,
2210
+ mel_filters: this.mel_filters,
2211
+ log_mel: 'log',
2212
+ mel_floor: 1.192092955078125e-07,
2213
+ remove_dc_offset: true,
2214
+
2215
+ // Custom
2216
+ transpose: true,
2217
+ min_num_frames: this.min_num_frames,
2218
+ }
2219
+ )
2220
+ }
2221
+
2222
+
2223
+ /**
2224
+ * Asynchronously extracts features from a given audio using the provided configuration.
2225
+ * @param {Float32Array|Float64Array} audio The audio data as a Float32Array/Float64Array.
2226
+ * @returns {Promise<{ input_features: Tensor }>} A Promise resolving to an object containing the extracted input features as a Tensor.
2227
+ */
2228
+ async _call(audio) {
2229
+ validate_audio_inputs(audio, 'WeSpeakerFeatureExtractor');
2230
+
2231
+ const features = (await this._extract_fbank_features(audio)).unsqueeze_(0);
2232
+
2233
+ if (this.config.fbank_centering_span === null) {
2234
+ // center features with global average
2235
+ const meanData = /** @type {Float32Array} */ (features.mean(1).data);
2236
+ const featuresData = /** @type {Float32Array} */(features.data);
2237
+ const [batch_size, num_frames, feature_size] = features.dims;
2238
+
2239
+ for (let i = 0; i < batch_size; ++i) {
2240
+ const offset1 = i * num_frames * feature_size;
2241
+ const offset2 = i * feature_size;
2242
+ for (let j = 0; j < num_frames; ++j) {
2243
+ const offset3 = offset1 + j * feature_size;
2244
+ for (let k = 0; k < feature_size; ++k) {
2245
+ featuresData[offset3 + k] -= meanData[offset2 + k];
2246
+ }
2247
+ }
2248
+ }
2249
+ }
2250
+
2251
+ return {
2252
+ input_features: features
2253
+ };
2254
+ }
2255
+ }
2256
+
2257
+ export class SpeechT5FeatureExtractor extends FeatureExtractor { }
2258
+
2259
+ /**
2260
+ * Represents a Processor that extracts features from an input.
2261
+ * @extends Callable
2262
+ */
2263
+ export class Processor extends Callable {
2264
+ /**
2265
+ * Creates a new Processor with the given feature extractor.
2266
+ * @param {FeatureExtractor} feature_extractor The function used to extract features from the input.
2267
+ */
2268
+ constructor(feature_extractor) {
2269
+ super();
2270
+ this.feature_extractor = feature_extractor;
2271
+ // TODO use tokenizer here?
2272
+ }
2273
+
2274
+ /**
2275
+ * Calls the feature_extractor function with the given input.
2276
+ * @param {any} input The input to extract features from.
2277
+ * @param {...any} args Additional arguments.
2278
+ * @returns {Promise<any>} A Promise that resolves with the extracted features.
2279
+ */
2280
+ async _call(input, ...args) {
2281
+ return await this.feature_extractor(input, ...args);
2282
+ }
2283
+ }
2284
+
2285
+ export class SamProcessor extends Processor {
2286
+ /**
2287
+ * @borrows SamImageProcessor#_call as _call
2288
+ */
2289
+ async _call(...args) {
2290
+ return await this.feature_extractor(...args);
2291
+ }
2292
+
2293
+ /**
2294
+ * @borrows SamImageProcessor#post_process_masks as post_process_masks
2295
+ */
2296
+ post_process_masks(...args) {
2297
+ // @ts-ignore
2298
+ return this.feature_extractor.post_process_masks(...args);
2299
+ }
2300
+ /**
2301
+ * @borrows SamImageProcessor#reshape_input_points as reshape_input_points
2302
+ */
2303
+ reshape_input_points(...args) {
2304
+ // @ts-ignore
2305
+ return this.feature_extractor.reshape_input_points(...args);
2306
+ }
2307
+ }
2308
+
2309
+ /**
2310
+ * Represents a WhisperProcessor that extracts features from an audio input.
2311
+ * @extends Processor
2312
+ */
2313
+ export class WhisperProcessor extends Processor {
2314
+ /**
2315
+ * Calls the feature_extractor function with the given audio input.
2316
+ * @param {any} audio The audio input to extract features from.
2317
+ * @returns {Promise<any>} A Promise that resolves with the extracted features.
2318
+ */
2319
+ async _call(audio) {
2320
+ return await this.feature_extractor(audio)
2321
+ }
2322
+ }
2323
+
2324
+
2325
+ export class Wav2Vec2ProcessorWithLM extends Processor {
2326
+ /**
2327
+ * Calls the feature_extractor function with the given audio input.
2328
+ * @param {any} audio The audio input to extract features from.
2329
+ * @returns {Promise<any>} A Promise that resolves with the extracted features.
2330
+ */
2331
+ async _call(audio) {
2332
+ return await this.feature_extractor(audio)
2333
+ }
2334
+ }
2335
+
2336
+ export class PyAnnoteProcessor extends Processor {
2337
+ /**
2338
+ * Calls the feature_extractor function with the given audio input.
2339
+ * @param {any} audio The audio input to extract features from.
2340
+ * @returns {Promise<any>} A Promise that resolves with the extracted features.
2341
+ */
2342
+ async _call(audio) {
2343
+ return await this.feature_extractor(audio)
2344
+ }
2345
+
2346
+ post_process_speaker_diarization(...args) {
2347
+ // @ts-ignore
2348
+ return this.feature_extractor.post_process_speaker_diarization(...args);
2349
+ }
2350
+
2351
+ }
2352
+
2353
+ export class SpeechT5Processor extends Processor {
2354
+ /**
2355
+ * Calls the feature_extractor function with the given input.
2356
+ * @param {any} input The input to extract features from.
2357
+ * @returns {Promise<any>} A Promise that resolves with the extracted features.
2358
+ */
2359
+ async _call(input) {
2360
+ return await this.feature_extractor(input)
2361
+ }
2362
+ }
2363
+
2364
+ export class OwlViTProcessor extends Processor { }
2365
+
2366
+ export class Florence2Processor extends Processor {
2367
+ constructor(feature_extractor) {
2368
+ super(feature_extractor);
2369
+
2370
+ const {
2371
+ tasks_answer_post_processing_type,
2372
+ task_prompts_without_inputs,
2373
+ task_prompts_with_input,
2374
+ } = feature_extractor.config;
2375
+
2376
+ /** @type {Map<string, string>} */
2377
+ this.tasks_answer_post_processing_type = new Map(Object.entries(tasks_answer_post_processing_type ?? {}));
2378
+
2379
+ /** @type {Map<string, string>} */
2380
+ this.task_prompts_without_inputs = new Map(Object.entries(task_prompts_without_inputs ?? {}));
2381
+
2382
+ /** @type {Map<string, string>} */
2383
+ this.task_prompts_with_input = new Map(Object.entries(task_prompts_with_input ?? {}));
2384
+
2385
+ this.regexes = {
2386
+ quad_boxes: /(.+?)<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>/gm,
2387
+ bboxes: /([^<]+)?<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>/gm,
2388
+ }
2389
+ this.size_per_bin = 1000;
2390
+ }
2391
+
2392
+ /**
2393
+ * Helper function to construct prompts from input texts
2394
+ * @param {string|string[]} text
2395
+ * @returns {string[]}
2396
+ */
2397
+ construct_prompts(text) {
2398
+ if (typeof text === 'string') {
2399
+ text = [text];
2400
+ }
2401
+
2402
+ const prompts = [];
2403
+ for (const t of text) {
2404
+ // 1. fixed task prompts without additional inputs
2405
+ if (this.task_prompts_without_inputs.has(t)) {
2406
+ prompts.push(this.task_prompts_without_inputs.get(t));
2407
+ }
2408
+ // 2. task prompts with additional inputs
2409
+ else {
2410
+ for (const [task, prompt] of this.task_prompts_with_input) {
2411
+ if (t.includes(task)) {
2412
+ prompts.push(prompt.replaceAll('{input}', t).replaceAll(task, ''));
2413
+ break;
2414
+ }
2415
+ }
2416
+
2417
+ // 3. default prompt
2418
+ if (prompts.length !== text.length) {
2419
+ prompts.push(t);
2420
+ }
2421
+ }
2422
+ }
2423
+ return prompts;
2424
+ }
2425
+
2426
+ /**
2427
+ * Post-process the output of the model to each of the task outputs.
2428
+ * @param {string} text The text to post-process.
2429
+ * @param {string} task The task to post-process the text for.
2430
+ * @param {[number, number]} image_size The size of the image. height x width.
2431
+ */
2432
+ post_process_generation(text, task, image_size) {
2433
+ const task_answer_post_processing_type = this.tasks_answer_post_processing_type.get(task) ?? 'pure_text';
2434
+
2435
+ // remove the special tokens
2436
+ text = text.replaceAll('<s>', '').replaceAll('</s>', '');
2437
+
2438
+ let final_answer;
2439
+ switch (task_answer_post_processing_type) {
2440
+ case 'pure_text':
2441
+ final_answer = text;
2442
+ break;
2443
+
2444
+ case 'description_with_bboxes':
2445
+ case 'bboxes':
2446
+ case 'phrase_grounding':
2447
+ case 'ocr':
2448
+ const key = task_answer_post_processing_type === 'ocr' ? 'quad_boxes' : 'bboxes';
2449
+ const matches = text.matchAll(this.regexes[key]);
2450
+ const labels = [];
2451
+ const items = [];
2452
+ for (const [_, label, ...locations] of matches) {
2453
+ // Push new label, or duplicate the last label
2454
+ labels.push(label ? label.trim() : labels.at(-1) ?? '');
2455
+ items.push(locations.map((x, i) =>
2456
+ // NOTE: Add 0.5 to use the center position of the bin as the coordinate.
2457
+ (Number(x) + 0.5) / this.size_per_bin * image_size[i % 2])
2458
+ );
2459
+ }
2460
+ final_answer = { labels, [key]: items };
2461
+ break;
2462
+
2463
+ default:
2464
+ throw new Error(`Task "${task}" (of type "${task_answer_post_processing_type}") not yet implemented.`);
2465
+ }
2466
+
2467
+ return { [task]: final_answer }
2468
+ }
2469
+ }
2470
+
2471
+ //////////////////////////////////////////////////
2472
+ /**
2473
+ * Helper class which is used to instantiate pretrained processors with the `from_pretrained` function.
2474
+ * The chosen processor class is determined by the type specified in the processor config.
2475
+ *
2476
+ * **Example:** Load a processor using `from_pretrained`.
2477
+ * ```javascript
2478
+ * let processor = await AutoProcessor.from_pretrained('openai/whisper-tiny.en');
2479
+ * ```
2480
+ *
2481
+ * **Example:** Run an image through a processor.
2482
+ * ```javascript
2483
+ * let processor = await AutoProcessor.from_pretrained('Xenova/clip-vit-base-patch16');
2484
+ * let image = await RawImage.read('https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/football-match.jpg');
2485
+ * let image_inputs = await processor(image);
2486
+ * // {
2487
+ * // "pixel_values": {
2488
+ * // "dims": [ 1, 3, 224, 224 ],
2489
+ * // "type": "float32",
2490
+ * // "data": Float32Array [ -1.558687686920166, -1.558687686920166, -1.5440893173217773, ... ],
2491
+ * // "size": 150528
2492
+ * // },
2493
+ * // "original_sizes": [
2494
+ * // [ 533, 800 ]
2495
+ * // ],
2496
+ * // "reshaped_input_sizes": [
2497
+ * // [ 224, 224 ]
2498
+ * // ]
2499
+ * // }
2500
+ * ```
2501
+ */
2502
+ export class AutoProcessor {
2503
+ static FEATURE_EXTRACTOR_CLASS_MAPPING = {
2504
+ ImageFeatureExtractor,
2505
+ WhisperFeatureExtractor,
2506
+ ViTFeatureExtractor,
2507
+ MobileViTFeatureExtractor,
2508
+ MobileViTImageProcessor,
2509
+ MobileNetV1FeatureExtractor,
2510
+ MobileNetV2FeatureExtractor,
2511
+ MobileNetV3FeatureExtractor,
2512
+ MobileNetV4FeatureExtractor,
2513
+ OwlViTFeatureExtractor,
2514
+ Owlv2ImageProcessor,
2515
+ CLIPFeatureExtractor,
2516
+ CLIPImageProcessor,
2517
+ Florence2Processor,
2518
+ ChineseCLIPFeatureExtractor,
2519
+ SiglipImageProcessor,
2520
+ ConvNextFeatureExtractor,
2521
+ ConvNextImageProcessor,
2522
+ SegformerFeatureExtractor,
2523
+ BitImageProcessor,
2524
+ DPTImageProcessor,
2525
+ DPTFeatureExtractor,
2526
+ GLPNFeatureExtractor,
2527
+ BeitFeatureExtractor,
2528
+ DeiTFeatureExtractor,
2529
+ DetrFeatureExtractor,
2530
+ RTDetrImageProcessor,
2531
+ YolosFeatureExtractor,
2532
+ DonutFeatureExtractor,
2533
+ NougatImageProcessor,
2534
+ EfficientNetImageProcessor,
2535
+
2536
+ ViTImageProcessor,
2537
+ VitMatteImageProcessor,
2538
+ SamImageProcessor,
2539
+ Swin2SRImageProcessor,
2540
+ Wav2Vec2FeatureExtractor,
2541
+ SeamlessM4TFeatureExtractor,
2542
+ SpeechT5FeatureExtractor,
2543
+ ASTFeatureExtractor,
2544
+ ClapFeatureExtractor,
2545
+ PyAnnoteFeatureExtractor,
2546
+ WeSpeakerFeatureExtractor,
2547
+ }
2548
+
2549
+ static PROCESSOR_CLASS_MAPPING = {
2550
+ WhisperProcessor,
2551
+ Wav2Vec2ProcessorWithLM,
2552
+ PyAnnoteProcessor,
2553
+ SamProcessor,
2554
+ SpeechT5Processor,
2555
+ OwlViTProcessor,
2556
+ Florence2Processor,
2557
+ }
2558
+
2559
+ /**
2560
+ * Instantiate one of the processor classes of the library from a pretrained model.
2561
+ *
2562
+ * The processor class to instantiate is selected based on the `feature_extractor_type` property of the config object
2563
+ * (either passed as an argument or loaded from `pretrained_model_name_or_path` if possible)
2564
+ *
2565
+ * @param {string} pretrained_model_name_or_path The name or path of the pretrained model. Can be either:
2566
+ * - A string, the *model id* of a pretrained processor hosted inside a model repo on huggingface.co.
2567
+ * Valid model ids can be located at the root-level, like `bert-base-uncased`, or namespaced under a
2568
+ * user or organization name, like `dbmdz/bert-base-german-cased`.
2569
+ * - A path to a *directory* containing processor files, e.g., `./my_model_directory/`.
2570
+ * @param {import('./utils/hub.js').PretrainedOptions} options Additional options for loading the processor.
2571
+ *
2572
+ * @returns {Promise<Processor>} A new instance of the Processor class.
2573
+ */
2574
+ static async from_pretrained(pretrained_model_name_or_path, {
2575
+ progress_callback = null,
2576
+ config = null,
2577
+ cache_dir = null,
2578
+ local_files_only = false,
2579
+ revision = 'main',
2580
+ } = {}) {
2581
+
2582
+ let preprocessorConfig = config ?? await getModelJSON(pretrained_model_name_or_path, 'preprocessor_config.json', true, {
2583
+ progress_callback,
2584
+ config,
2585
+ cache_dir,
2586
+ local_files_only,
2587
+ revision,
2588
+ })
2589
+
2590
+ // Determine feature extractor class
2591
+ // TODO: Ensure backwards compatibility with old configs
2592
+ let key = preprocessorConfig.feature_extractor_type ?? preprocessorConfig.image_processor_type;
2593
+ let feature_extractor_class = this.FEATURE_EXTRACTOR_CLASS_MAPPING[key];
2594
+
2595
+ if (!feature_extractor_class) {
2596
+ if (preprocessorConfig.size !== undefined) {
2597
+ // Assume ImageFeatureExtractor
2598
+ console.warn(`Feature extractor type "${key}" not found, assuming ImageFeatureExtractor due to size parameter in config.`);
2599
+ feature_extractor_class = ImageFeatureExtractor;
2600
+ } else {
2601
+ throw new Error(`Unknown Feature Extractor type: ${key}`);
2602
+ }
2603
+ }
2604
+
2605
+ // If no associated processor class, use default
2606
+ let processor_class = this.PROCESSOR_CLASS_MAPPING[preprocessorConfig.processor_class] ?? Processor;
2607
+
2608
+ // Instantiate processor and feature extractor
2609
+ let feature_extractor = new feature_extractor_class(preprocessorConfig);
2610
+ return new processor_class(feature_extractor);
2611
+ }
2612
+ }
2613
+ //////////////////////////////////////////////////
2614
+