Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add functionality to split RawImage into channels; Add functionality to return a vertical slice of a Tensor #978

Open
wants to merge 5 commits into
base: v3
Choose a base branch
from

Conversation

BritishWerewolf
Copy link
Contributor

@BritishWerewolf BritishWerewolf commented Oct 16, 2024

Some models, like U2Net, require images to be processed with all the reds, greens and blue pixels grouped together.

As a result, this PR provides the ability to do something like this.

// Get our image, then split into individual channels.
const image = await RawImage.fromURL("https://picsum.photos/600/400");

// ... snip ...
// Preprocess the image by resizing and padding to a square of 320x320
// See these PRs for functionality to do this: #971, #972

// Image pixel values are an array of [R, G, B, R, G, B,…]
// Use this to separate out each individual channel.
const [R, G, B] = image.toChannels();
// Convert from an array of Uint8Array into regular array, then flatten.
// We now have an array of [R, R, …, G, G, …, B, B, …]
const image_data = [R, G, B].map(d => Array.from(d)).flat();

// Create the tensor for this image
const tensor = new Tensor('float32', image_data, [1, image.channels, image.width, image.height]);
const model = await AutoModel.from_pretrained('u2netp');
const outputs = await model({ "input": tensor });

In addition to the toChannels method, I created another for Tensor's too.

With .slice we could return specific rows, this PR adds the ability to return vertical slices of a Tensor.
For some examples, check out the tests I've added.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant