diff --git a/src/utils/image.js b/src/utils/image.js index 33bdf11d8..ded253f42 100644 --- a/src/utils/image.js +++ b/src/utils/image.js @@ -1,10 +1,13 @@ +/** + * @typedef {import('./tensor.js').DataArray} DataArray + */ /** - * @file Helper module for image processing. - * - * These functions and classes are only used internally, + * @file Helper module for image processing. + * + * These functions and classes are only used internally, * meaning an end-user shouldn't need to access anything here. - * + * * @module utils/image */ @@ -91,7 +94,7 @@ export class RawImage { this.channels = channels; } - /** + /** * Returns the size of the image (width, height). * @returns {[number, number]} The size of the image (width, height). */ @@ -101,9 +104,9 @@ export class RawImage { /** * Helper method for reading an image from a variety of input types. - * @param {RawImage|string|URL} input + * @param {RawImage|string|URL} input * @returns The image object. - * + * * **Example:** Read image from a URL. * ```javascript * let image = await RawImage.read('https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/football-match.jpg'); @@ -181,7 +184,7 @@ export class RawImage { /** * Helper method to create a new Image from a tensor - * @param {Tensor} tensor + * @param {Tensor} tensor */ static fromTensor(tensor, channel_format = 'CHW') { if (tensor.dims.length !== 3) { @@ -304,6 +307,46 @@ export class RawImage { return this._update(newData, this.width, this.height, 4); } + /** + * Apply an alpha mask to the image. + * @param {RawImage} mask The mask to apply. Values should be between 0 and 255, and be a single channel. + * @returns {Promise} The masked image. + * @throws {Error} If the mask is not the same size as the image. + * @throws {Error} If the image does not have 4 channels. + * @throws {Error} If the mask is not a single channel. + */ + async putAlpha(mask) { + if (mask.width !== this.width || mask.height !== this.height) { + throw new Error('Mask must be the same size as the image'); + } + + // We want the current image to have an alpha channel, but the mask will + // just be a single channel. + if (this.channels !== 4) { + throw new Error('Image must have 4 channels'); + } + if (mask.channels !== 1) { + throw new Error('Mask must have 1 channel'); + } + + const numPixels = this.width * this.height; + for (let i = 0; i < numPixels; i++) { + const maskPixel = mask.data[i]; + if (typeof maskPixel === 'undefined') { + throw new Error('Invalid mask'); + } + + // Ensure that the alpha is a range from 0 to 255, and not a range + // from 0 to 1. + const alpha = maskPixel < 1 ? maskPixel * 255 : maskPixel; + // Calculate the offset, multiplying by 4 because of the number of + // channels, then offset by 3 to get the alpha channel. + this.data[(i * 4) + 3] = alpha; + } + + return this; + } + /** * Resize the image to the given dimensions. This method uses the canvas API to perform the resizing. * @param {number} width The width of the new image. @@ -355,7 +398,7 @@ export class RawImage { case 'nearest': case 'bilinear': case 'bicubic': - // Perform resizing using affine transform. + // Perform resizing using affine transform. // This matches how the python Pillow library does it. img = img.affine([width / this.width, 0, 0, height / this.height], { interpolator: resampleMethod @@ -368,7 +411,7 @@ export class RawImage { img = img.resize({ width, height, fit: 'fill', - kernel: 'lanczos3', // PIL Lanczos uses a kernel size of 3 + kernel: 'lanczos3', // PIL Lanczos uses a kernel size of 3 }); break; @@ -447,7 +490,7 @@ export class RawImage { // Create canvas object for this image const canvas = this.toCanvas(); - // Create a new canvas of the desired size. This is needed since if the + // Create a new canvas of the desired size. This is needed since if the // image is too small, we need to pad it with black pixels. const ctx = createCanvasFunction(crop_width, crop_height).getContext('2d'); @@ -495,7 +538,7 @@ export class RawImage { // Create canvas object for this image const canvas = this.toCanvas(); - // Create a new canvas of the desired size. This is needed since if the + // Create a new canvas of the desired size. This is needed since if the // image is too small, we need to pad it with black pixels. const ctx = createCanvasFunction(crop_width, crop_height).getContext('2d'); @@ -742,4 +785,4 @@ export class RawImage { } }); } -} \ No newline at end of file +}