@huggingface/transformers 3.0.1 → 3.1.0

package/src/models/sam/image_processing_sam.js

@@ -0,0 +1,242 @@
+import { 
+    ImageProcessor,
+} from "../../base/image_processors_utils.js";
+import { calculateDimensions } from "../../utils/core.js";
+
+import {
+    interpolate_4d,
+    Tensor,
+} from "../../utils/tensor.js";
+
+
+/**
+ * @typedef {object} SamImageProcessorResult
+ * @property {Tensor} pixel_values
+ * @property {import("../../base/image_processors_utils.js").HeightWidth[]} original_sizes
+ * @property {import("../../base/image_processors_utils.js").HeightWidth[]} reshaped_input_sizes
+ * @property {Tensor} [input_points]
+ * @property {Tensor} [input_labels]
+ * @property {Tensor} [input_boxes]
+ */
+
+export class SamImageProcessor extends ImageProcessor {
+
+    /**
+     * 
+     * @param {any} input_points 
+     * @param {import("../../base/image_processors_utils.js").HeightWidth[]} original_sizes 
+     * @param {import("../../base/image_processors_utils.js").HeightWidth[]} reshaped_input_sizes 
+     * @returns {Tensor}
+     */
+    reshape_input_points(input_points, original_sizes, reshaped_input_sizes, is_bounding_box = false) {
+
+        // Make deep copy to avoid altering user's input
+        input_points = structuredClone(input_points);
+        let shape = calculateDimensions(input_points);
+
+        // TODO: add support for 2D input_points
+        if (shape.length === 3) {
+            // Correct user's input
+            if (!is_bounding_box) {
+                shape = [1, ...shape];
+            }
+            input_points = [input_points];
+        } else if (shape.length !== 4) {
+            throw Error("The input_points must be a 4D tensor of shape `batch_size`, `point_batch_size`, `nb_points_per_image`, `2`.")
+        }
+
+        // Reshape input points
+        for (let i = 0; i < input_points.length; ++i) { // batch_size
+            let originalImageSize = original_sizes[i];
+            let reshapedImageSize = reshaped_input_sizes[i];
+
+            let resizeFactors = [
+                reshapedImageSize[0] / originalImageSize[0],
+                reshapedImageSize[1] / originalImageSize[1]
+            ]
+
+            for (let j = 0; j < input_points[i].length; ++j) { // point_batch_size
+                for (let k = 0; k < input_points[i][j].length; ++k) { // nb_points_per_image
+                    for (let w = 0; w < input_points[i][j][k].length; ++w) { // 2 or 4
+                        input_points[i][j][k][w] *= resizeFactors[w % 2];
+                    }
+                }
+            }
+        }
+
+        return new Tensor(
+            'float32',
+            Float32Array.from(input_points.flat(Infinity)),
+            shape
+        )
+
+    }
+
+    /**
+     * 
+     * @param {any} input_labels 
+     * @param {Tensor} input_points 
+     * @returns {Tensor}
+     */
+    add_input_labels(input_labels, input_points) {
+        let shape = calculateDimensions(input_labels);
+        if (shape.length === 2) {
+            // Correct user's input
+            shape = [1, ...shape];
+            input_labels = [input_labels];
+        } else if (shape.length !== 3) {
+            throw Error("The input_points must be a 4D tensor of shape `batch_size`, `point_batch_size`, `nb_points_per_image`, `2`.")
+        }
+
+        if (shape.some((x, i) => x !== input_points.dims[i])) {
+            throw Error(`The first ${shape.length} dimensions of 'input_points' and 'input_labels' must be the same.`)
+        }
+        return new Tensor(
+            'int64',
+            input_labels.flat(Infinity).map(BigInt),
+            shape,
+        )
+    }
+    /**
+     * @param {any[]} images The URL(s) of the image(s) to extract features from.
+     * @param {Object} [options] Additional options for the processor.
+     * @param {any} [options.input_points=null] A 3D or 4D array, representing the input points provided by the user.
+     * - 3D: `[point_batch_size, nb_points_per_image, 2]`. In this case, `batch_size` is assumed to be 1.
+     * - 4D: `[batch_size, point_batch_size, nb_points_per_image, 2]`.
+     * @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.
+     * - 2D: `[point_batch_size, nb_points_per_image]`. In this case, `batch_size` is assumed to be 1.
+     * - 3D: `[batch_size, point_batch_size, nb_points_per_image]`.
+     * @param {number[][][]} [options.input_boxes=null] A 3D array of shape `(batch_size, num_boxes, 4)`, representing the input boxes provided by the user.
+     * This is used by the prompt encoder to encode the prompt. Generally yields to much better generated masks.
+     * The processor will generate a tensor, with each dimension corresponding respectively to the image batch size,
+     * the number of boxes per image and the coordinates of the top left and botton right point of the box.
+     * In the order (`x1`, `y1`, `x2`, `y2`):
+     * - `x1`: the x coordinate of the top left point of the input box
+     * - `y1`: the y coordinate of the top left point of the input box
+     * - `x2`: the x coordinate of the bottom right point of the input box
+     * - `y2`: the y coordinate of the bottom right point of the input box
+     * @returns {Promise<SamImageProcessorResult>}
+     */
+    async _call(images, {
+        input_points = null,
+        input_labels = null,
+        input_boxes = null
+    } = {}) {
+        // TODO allow user to use preprocessed images
+        /** @type {SamImageProcessorResult} */
+        const processed = await super._call(images);
+
+        if (input_points) {
+            processed.input_points = this.reshape_input_points(
+                input_points, processed.original_sizes, processed.reshaped_input_sizes
+            );
+        }
+
+        if (input_labels) {
+            if (!processed.input_points) {
+                throw Error("`input_points` must be provided if `input_labels` are provided.")
+            }
+            processed.input_labels = this.add_input_labels(input_labels, processed.input_points);
+        }
+
+        if (input_boxes) {
+            processed.input_boxes = this.reshape_input_points(
+                input_boxes, processed.original_sizes, processed.reshaped_input_sizes, true,
+            );
+        }
+
+        return processed;
+    }
+
+    /**
+     * Remove padding and upscale masks to the original image size.
+     * @param {Tensor} masks Batched masks from the mask_decoder in (batch_size, num_channels, height, width) format.
+     * @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.
+     * @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.
+     * @param {Object} options Optional parameters for post-processing.
+     * @param {number} [options.mask_threshold] The threshold to use for binarizing the masks.
+     * @param {boolean} [options.binarize] Whether to binarize the masks.
+     * @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`.
+     * @param {number} [options.pad_size.height] The height the images were padded to.
+     * @param {number} [options.pad_size.width] The width the images were padded to.
+     * @returns {Promise<Tensor[]>} Batched masks in batch_size, num_channels, height, width) format, where (height, width) is given by original_size.
+     */
+    async post_process_masks(masks, original_sizes, reshaped_input_sizes, {
+        mask_threshold = 0.0,
+        binarize = true,
+        pad_size = null,
+    } = {}) {
+        // masks: [1, 1, 3, 256, 256]
+
+        const output_masks = [];
+
+        pad_size = pad_size ?? this.pad_size;
+
+        /** @type {[number, number]} */
+        const target_image_size = [pad_size.height, pad_size.width];
+
+        for (let i = 0; i < original_sizes.length; ++i) {
+            const original_size = original_sizes[i];
+            const reshaped_input_size = reshaped_input_sizes[i];
+
+            // Upscale mask to padded size
+            let interpolated_mask = (await interpolate_4d(
+                masks[i],
+                { mode: 'bilinear', size: target_image_size }
+            ));
+
+            // Crop mask
+            interpolated_mask = interpolated_mask.slice(null, null, [0, reshaped_input_size[0]], [0, reshaped_input_size[1]]);
+
+            // Downscale mask
+            interpolated_mask = (await interpolate_4d(
+                interpolated_mask,
+                { mode: 'bilinear', size: original_size }
+            ));
+
+            if (binarize) {
+                const data = interpolated_mask.data;
+                const binarizedMaskData = new Uint8Array(data.length);
+                for (let i = 0; i < data.length; ++i) {
+                    if (data[i] > mask_threshold) {
+                        binarizedMaskData[i] = 1;
+                    }
+                }
+                interpolated_mask = new Tensor(
+                    'bool',
+                    binarizedMaskData,
+                    interpolated_mask.dims
+                )
+            }
+
+            output_masks.push(interpolated_mask);
+        }
+
+        return output_masks;
+    }
+
+    /**
+     * Generates a list of crop boxes of different sizes. Each layer has (2**i)**2 boxes for the ith layer.
+     * @param {import("../../utils/image.js").RawImage} image Input original image
+     * @param {number} target_size Target size of the resized image
+     * @param {Object} options Options for generating crop boxes 
+     * @param {number} [options.crop_n_layers] If >0, mask prediction will be run again on crops of the image.
+     * Sets the number of layers to run, where each layer has 2**i_layer number of image crops.
+     * @param {number} [options.overlap_ratio] Sets the degree to which crops overlap. In the first crop layer,
+     * crops will overlap by this fraction of the image length. Later layers with more crops scale down this overlap.
+     * @param {number} [options.points_per_crop] Number of points to sample from each crop.
+     * @param {number} [options.crop_n_points_downscale_factor] The number of points-per-side sampled in layer n is
+     * scaled down by crop_n_points_downscale_factor**n.
+     * @returns {Object} An object containing the crop boxes, number of points per crop, cropped images, and input labels.
+     */
+    generate_crop_boxes(image, target_size, {
+        crop_n_layers = 0,
+        overlap_ratio = 512 / 1500,
+        points_per_crop = 32,
+        crop_n_points_downscale_factor = 1,
+    } = {}) {
+        // TODO: Implement
+        // return { crop_boxes, points_per_crop, cropped_images, input_labels }
+    }
+}
+

package/src/models/sam/processing_sam.js

@@ -0,0 +1,20 @@
+import { Processor } from "../../base/processing_utils.js";
+import { AutoImageProcessor } from "../auto/image_processing_auto.js";
+
+export class SamProcessor extends Processor {
+    static image_processor_class = AutoImageProcessor
+
+    async _call(...args) {
+        return await this.image_processor(...args);
+    }
+
+    post_process_masks(...args) {
+        // @ts-ignore
+        return this.image_processor.post_process_masks(...args);
+    }
+
+    reshape_input_points(...args) {
+        // @ts-ignore
+        return this.image_processor.reshape_input_points(...args);
+    }
+}

package/src/models/sapiens/image_processing_sapiens.js

@@ -0,0 +1,13 @@
+import { 
+    ImageProcessor,
+    post_process_semantic_segmentation,
+} from "../../base/image_processors_utils.js";
+
+
+export class SapiensImageProcessor extends ImageProcessor {
+    /** @type {typeof post_process_semantic_segmentation} */
+    post_process_semantic_segmentation(...args) {
+        return post_process_semantic_segmentation(...args);
+    }
+}
+export class SapiensFeatureExtractor extends SapiensImageProcessor { }

package/src/models/seamless_m4t/feature_extraction_seamless_m4t.js

@@ -0,0 +1,180 @@
+import { FeatureExtractor, validate_audio_inputs } from '../../base/feature_extraction_utils.js';
+import { Tensor } from '../../utils/tensor.js';
+import { mel_filter_bank, spectrogram, window_function } from '../../utils/audio.js';
+
+export class SeamlessM4TFeatureExtractor extends FeatureExtractor {
+
+    constructor(config) {
+        super(config);
+
+        const sampling_rate = this.config.sampling_rate;
+        const mel_filters = mel_filter_bank(
+            256, // num_frequency_bins
+            this.config.num_mel_bins, // num_mel_filters
+            20, // min_frequency
+            Math.floor(sampling_rate / 2), // max_frequency
+            sampling_rate, // sampling_rate
+            null, // norm
+            "kaldi", // mel_scale
+            true, // triangularize_in_mel_space
+        );
+
+        // Do padding:
+        for (let i = 0; i < mel_filters.length; ++i) {
+            mel_filters[i].push(0);
+        }
+        this.mel_filters = mel_filters;
+
+        this.window = window_function(400, 'povey', {
+            periodic: false,
+        })
+    }
+
+    /**
+     * Computes the log-Mel spectrogram of the provided audio waveform.
+     * @param {Float32Array|Float64Array} waveform The audio waveform to process.
+     * @param {number} max_length The maximum number of frames to return.
+     * @returns {Promise<Tensor>} An object containing the log-Mel spectrogram data as a Float32Array and its dimensions as an array of numbers.
+     */
+    async _extract_fbank_features(waveform, max_length) {
+        // NOTE: We don't pad/truncate since that is passed in as `max_num_frames`
+
+        // Kaldi compliance: 16-bit signed integers
+        // 32768 == 2 ** 15
+        waveform = waveform.map((/** @type {number} */ x) => x * 32768)
+
+        return spectrogram(
+            waveform,
+            this.window, // window
+            400, // frame_length
+            160, // hop_length
+            {
+                fft_length: 512,
+                power: 2.0,
+                center: false,
+                preemphasis: 0.97,
+                mel_filters: this.mel_filters,
+                log_mel: 'log',
+                mel_floor: 1.192092955078125e-07,
+                remove_dc_offset: true,
+
+                // Custom
+                max_num_frames: max_length,
+                transpose: true,
+            }
+        )
+    }
+
+    /**
+     * Asynchronously extracts features from a given audio using the provided configuration.
+     * @param {Float32Array|Float64Array} audio The audio data as a Float32Array/Float64Array.
+     * @param {Object} options Optional parameters for feature extraction.
+     * @param {boolean} [options.padding=true] Whether to pad the sequence to a multiple of `pad_to_multiple_of`.
+     * @param {number} [options.pad_to_multiple_of=2] The number to pad the sequence to a multiple of.
+     * @param {boolean} [options.do_normalize_per_mel_bins=true] Whether or not to zero-mean unit-variance normalize the input per mel-channel.
+     * @param {boolean} [options.return_attention_mask=true] Whether to return the attention mask.
+     * @returns {Promise<{ input_features: Tensor, attention_mask?: Tensor }>} A Promise resolving to an object containing the extracted input features and attention masks as Tensors.
+     */
+    async _call(audio, {
+        padding = true,
+        pad_to_multiple_of = 2,
+        do_normalize_per_mel_bins = true,
+        return_attention_mask = true,
+    } = {}) {
+        validate_audio_inputs(audio, 'SeamlessM4TFeatureExtractor');
+
+        let features = await this._extract_fbank_features(audio, this.config.max_length);
+
+        if (do_normalize_per_mel_bins) {
+            const [num_features, feature_size] = features.dims;
+            const data = features.data;
+            for (let i = 0; i < feature_size; ++i) {
+                let sum = 0;
+                for (let j = 0; j < num_features; ++j) {
+                    sum += data[j * feature_size + i];
+                }
+
+                const mean = sum / num_features;
+
+                let variance = 0;
+                for (let j = 0; j < num_features; ++j) {
+                    variance += (data[j * feature_size + i] - mean) ** 2;
+                }
+                variance /= num_features - 1; // NOTE: We use ddof=1
+
+                const std = Math.sqrt(variance + 1e-7);
+                for (let j = 0; j < num_features; ++j) {
+                    const index = j * feature_size + i;
+                    data[index] = (data[index] - mean) / std;
+                }
+            }
+        }
+
+        let padded_attention_mask;
+        if (padding) {
+            const [num_frames, num_channels] = features.dims;
+            const data = /** @type {Float32Array} */(features.data);
+
+            const pad_size = num_frames % pad_to_multiple_of;
+            if (pad_size > 0) {
+                const padded_data = new Float32Array(num_channels * (num_frames + pad_size));
+                padded_data.set(data)
+                padded_data.fill(this.config.padding_value, data.length)
+
+                const numPaddedFrames = num_frames + pad_size;
+                features = new Tensor(
+                    features.type,
+                    padded_data,
+                    [numPaddedFrames, num_channels],
+                )
+
+                if (return_attention_mask) {
+                    padded_attention_mask = new Tensor(
+                        'int64',
+                        new BigInt64Array(numPaddedFrames),
+                        [1, numPaddedFrames],
+                    )
+                    padded_attention_mask.data.fill(1n, 0, num_frames);
+                }
+            }
+        }
+
+        const [num_frames, num_channels] = features.dims;
+
+        const stride = this.config.stride;
+        const remainder = num_frames % stride;
+        if (remainder !== 0) {
+            throw new Error(`The number of frames (${num_frames}) must be a multiple of the stride (${stride}).`)
+        }
+
+        const input_features = features.view(
+            1,
+            Math.floor(num_frames / stride),
+            num_channels * stride,
+        );
+
+        const result = { input_features }
+
+        if (return_attention_mask) {
+            const reshapedNumFrames = input_features.dims[1];
+
+            const attention_mask_data = new BigInt64Array(reshapedNumFrames);
+
+            if (padded_attention_mask) {
+                const padded_attention_mask_data = padded_attention_mask.data;
+                for (let i = 1, j = 0; i < num_frames; i += stride, ++j) {
+                    attention_mask_data[j] = padded_attention_mask_data[i];
+                }
+            } else {
+                attention_mask_data.fill(1n);
+            }
+            result.attention_mask = new Tensor(
+                'int64',
+                attention_mask_data,
+                [1, reshapedNumFrames],
+            );
+        }
+
+        return result;
+    }
+}

package/src/models/segformer/image_processing_segformer.js

@@ -0,0 +1,13 @@
+import { 
+    ImageProcessor,
+    post_process_semantic_segmentation,
+} from "../../base/image_processors_utils.js";
+
+
+export class SegformerImageProcessor extends ImageProcessor {
+    /** @type {typeof post_process_semantic_segmentation} */
+    post_process_semantic_segmentation(...args) {
+        return post_process_semantic_segmentation(...args);
+    }
+}
+export class SegformerFeatureExtractor extends SegformerImageProcessor { }

package/src/models/siglip/image_processing_siglip.js

@@ -0,0 +1,5 @@
+import { 
+    ImageProcessor,
+} from "../../base/image_processors_utils.js";
+
+export class SiglipImageProcessor extends ImageProcessor { }

package/src/models/speecht5/feature_extraction_speecht5.js

@@ -0,0 +1,4 @@
+
+import { FeatureExtractor } from "../../base/feature_extraction_utils.js";
+
+export class SpeechT5FeatureExtractor extends FeatureExtractor { }

package/src/models/speecht5/processing_speecht5.js

@@ -0,0 +1,17 @@
+import { Processor } from "../../base/processing_utils.js";
+import { AutoTokenizer } from "../../tokenizers.js";
+import { AutoFeatureExtractor } from "../auto/feature_extraction_auto.js";
+
+export class SpeechT5Processor extends Processor {
+    static tokenizer_class = AutoTokenizer
+    static feature_extractor_class = AutoFeatureExtractor
+
+    /**
+     * Calls the feature_extractor function with the given input.
+     * @param {any} input The input to extract features from.
+     * @returns {Promise<any>} A Promise that resolves with the extracted features.
+     */
+    async _call(input) {
+        return await this.feature_extractor(input)
+    }
+}

package/src/models/swin2sr/image_processing_swin2sr.js

@@ -0,0 +1,24 @@
+import { 
+    ImageProcessor,
+} from "../../base/image_processors_utils.js";
+
+export class Swin2SRImageProcessor extends ImageProcessor {
+    pad_image(pixelData, imgDims, padSize, options = {}) {
+        // NOTE: In this case, `padSize` represents the size of the sliding window for the local attention.
+        // In other words, the image is padded so that its width and height are multiples of `padSize`.
+        const [imageHeight, imageWidth, imageChannels] = imgDims;
+
+        return super.pad_image(pixelData, imgDims, {
+            // NOTE: For Swin2SR models, the original python implementation adds padding even when the image's width/height is already
+            // a multiple of `pad_size`. However, this is most likely a bug (PR: https://github.com/mv-lab/swin2sr/pull/19).
+            // For this reason, we only add padding when the image's width/height is not a multiple of `pad_size`.
+            width: imageWidth + (padSize - imageWidth % padSize) % padSize,
+            height: imageHeight + (padSize - imageHeight % padSize) % padSize,
+        }, {
+            mode: 'symmetric',
+            center: false,
+            constant_values: -1,
+            ...options,
+        })
+    }
+}

package/src/models/vit/image_processing_vit.js

@@ -0,0 +1,7 @@
+import { 
+    ImageProcessor,
+} from "../../base/image_processors_utils.js";
+
+export class ViTImageProcessor extends ImageProcessor { }
+export class ViTFeatureExtractor extends ViTImageProcessor { }
+

package/src/models/vitmatte/image_processing_vitmatte.js

@@ -0,0 +1,50 @@
+import { 
+    ImageProcessor,
+} from "../../base/image_processors_utils.js";
+
+import {
+    stack,
+    cat,
+} from "../../utils/tensor.js";
+
+export class VitMatteImageProcessor extends ImageProcessor {
+    /**
+     * Calls the feature extraction process on an array of images, preprocesses
+     * each image, and concatenates the resulting features into a single Tensor.
+     * @param {import("../../utils/image.js").RawImage[]} images The image(s) to extract features from.
+     * @param {import("../../utils/image.js").RawImage[]} trimaps The trimaps(s) to extract features from.
+     * @returns {Promise<import("../../base/image_processors_utils.js").ImageProcessorResult>} An object containing the concatenated pixel values of the preprocessed images.
+     */
+    async _call(images, trimaps) {
+        if (!Array.isArray(images)) {
+            images = [images];
+        }
+        if (!Array.isArray(trimaps)) {
+            trimaps = [trimaps];
+        }
+
+        const imageData = await Promise.all(images.map(x => this.preprocess(x)));
+        const trimapData = await Promise.all(trimaps.map(x => this.preprocess(x, {
+            do_normalize: false,
+            do_convert_rgb: false,
+            do_convert_grayscale: true,
+        })));
+
+
+        // Stack pixel values
+        const pixel_values = stack(imageData.map(
+            // Concatenate images and trimaps
+            (x, i) => cat([x.pixel_values, trimapData[i].pixel_values], 0)
+        ), 0);
+
+        return {
+            pixel_values,
+
+            // Original sizes of images
+            original_sizes: imageData.map(x => x.original_size),
+
+            // Reshaped sizes of images, before padding or cropping
+            reshaped_input_sizes: imageData.map(x => x.reshaped_input_size),
+        }
+    }
+}

package/src/models/vitpose/image_processing_vitpose.js

@@ -0,0 +1,89 @@
+import {
+    ImageProcessor,
+} from "../../base/image_processors_utils.js";
+
+export class VitPoseImageProcessor extends ImageProcessor {
+
+    /**
+     * Transform the heatmaps into keypoint predictions and transform them back to the image.
+     * NOTE: This is a naive implementation and does not include advanced post-processing techniques,
+     * so the results may not be as accurate as the original implementation.
+     * @param {import('../../utils/tensor.js').Tensor} outputs The model outputs.
+     * @param {[number, number, number, number][][]} boxes List or array of bounding boxes for each image.
+     * Each box should be a list of 4 floats representing the bounding box coordinates in COCO format (top_left_x, top_left_y, width, height).
+     * @returns {{
+     *   bbox: [number, number, number, number],
+     *   scores: number[],
+     *   labels: number[],
+     *   keypoints: [number, number][]
+     * }[][]} List of keypoints predictions for each image.
+     */
+    post_process_pose_estimation(outputs, boxes, {
+        threshold = null,
+        // TODO:
+        // kernel_size = 11,
+        // target_sizes = null,
+    } = {}) {
+        // NOTE: boxes are 3D (batch_size, num_boxes, 4)
+        const heatmaps = outputs.tolist();
+        const [batch_size, num_classes, height, width] = outputs.dims;
+
+        const results = [];
+        for (let b = 0; b < batch_size; ++b) {
+            const heatmap = heatmaps[b];
+            const bboxes = boxes[b];
+
+            const batch_results = [];
+            for (let n = 0; n < bboxes.length; ++n) {
+                const bbox = bboxes[n];
+
+                const keypoints = [];
+                const scores = [];
+                const labels = [];
+
+                const xScale = bbox.at(-2) / width;
+                const yScale = bbox.at(-1) / height;
+                for (let c = 0; c < heatmap.length; ++c) {
+                    let [xWeightedSum, yWeightedSum] = [0, 0];
+                    let sum = 0;
+                    let score = -Infinity;
+                    const row = heatmap[c];
+                    for (let y = 0; y < row.length; ++y) {
+                        const col = row[y];
+                        for (let x = 0; x < col.length; ++x) {
+                            const value = col[x];
+                            sum += value;
+
+                            score = Math.max(score, value);
+
+                            // Get weighted sum of positions
+                            // TODO: Determine best offsets
+                            xWeightedSum += (x + 0.5) * value;
+                            yWeightedSum += (y) * value;
+                        }
+                    }
+
+                    // Ignore low scores, if threshold is set
+                    if (threshold != null && score < threshold) continue;
+
+                    /** @type {[number, number]} */
+                    const keypoint = [
+                        xScale * xWeightedSum / sum,
+                        yScale * yWeightedSum / sum,
+                    ]
+                    keypoints.push(keypoint);
+                    labels.push(c);
+                    scores.push(score);
+                }
+                batch_results.push({
+                    bbox,
+                    scores,
+                    labels,
+                    keypoints,
+                });
+            }
+            results.push(batch_results);
+        }
+        return results;
+    }
+}