@huggingface/transformers/src/models/vitmatte/image_processing_vitmatte.js

State-of-the-art Machine Learning for the web. Run 🤗 Transformers directly in your browser, with no need for a server!

import { ImageProcessor } from '../../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("../../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),
        };
    }
}