diff --git a/src/utils/image.js b/src/utils/image.js index 33bdf11d8..754af1d35 100644 --- a/src/utils/image.js +++ b/src/utils/image.js @@ -1,10 +1,10 @@ /** - * @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 +91,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 +101,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 +181,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 +304,46 @@ export class RawImage { return this._update(newData, this.width, this.height, 4); } + /** + * Apply an alpha mask to the image. Operates in place. + * @param {RawImage} mask The mask to apply. It should have a single channel. + * @returns {RawImage} 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. + */ + putAlpha(mask) { + if (mask.width !== this.width || mask.height !== this.height) { + throw new Error(`Expected mask size to be ${this.width}x${this.height}, but got ${mask.width}x${mask.height}`); + } + if (mask.channels !== 1) { + throw new Error(`Expected mask to have 1 channel, but got ${mask.channels}`); + } + + const this_data = this.data; + const mask_data = mask.data; + const num_pixels = this.width * this.height; + if (this.channels === 3) { + // Convert to RGBA and simultaneously apply mask to alpha channel + const newData = new Uint8ClampedArray(num_pixels * 4); + for (let i = 0, in_offset = 0, out_offset = 0; i < num_pixels; ++i) { + newData[out_offset++] = this_data[in_offset++]; + newData[out_offset++] = this_data[in_offset++]; + newData[out_offset++] = this_data[in_offset++]; + newData[out_offset++] = mask_data[i]; + } + return this._update(newData, this.width, this.height, 4); + + } else if (this.channels === 4) { + // Apply mask to alpha channel in place + for (let i = 0; i < num_pixels; ++i) { + this_data[4 * i + 3] = mask_data[i]; + } + return this; + } + throw new Error(`Expected image to have 3 or 4 channels, but got ${this.channels}`); + } + /** * 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 +395,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 +408,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 +487,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 +535,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');