[feature] Add support for juce::image to torch::tensor and back about pytorch-cpp

[mfl28]

I got this to work with the following code, without using OpenCV:

torch::Tensor image_to_tensor(const juce::Image& image) {
    juce::Image::BitmapData image_data(image, juce::Image::BitmapData::readOnly);

    auto num_channels = format_to_num_channels(image.getFormat());

    // output tensor format: [num_channels, height, width], values in interval [0, 1]
    return torch::from_blob(image_data.data, {image.getHeight(), image.getWidth(), num_channels}, torch::kUInt8)
           .clone()
           .to(torch::kFloat32)
           .permute({2, 0, 1})
           .div_(255);
}

juce::Image tensor_to_image(const torch::Tensor& tensor) {
    // expected tensor format: [num_channels, height, width], values in interval [0, 1]
    auto image_format = num_channels_to_format(tensor.size(0));
    auto image_width = tensor.size(2);
    auto image_height = tensor.size(1);

    auto image_tensor = tensor.detach()
                              .permute({1, 2, 0})
                              .mul_(255)
                              .to(torch::kUInt8)
                              .cpu()
                              .contiguous();

    juce::Image image(image_format, image_width, image_height, true);
    juce::Image::BitmapData image_data(image, juce::Image::BitmapData::writeOnly);

    std::memcpy(image_data.data, image_tensor.data_ptr(), image_tensor.numel() * sizeof(torch::kUInt8));
    return image;
}

with the format <--> channels conversion functions:

juce::Image::PixelFormat num_channels_to_format(int num_channels) {
    switch(num_channels) {
        case 1:
            return juce::Image::SingleChannel;
        case 3:
            return juce::Image::RGB;
        case 4:
            return juce::Image::ARGB;
        default:
            throw std::runtime_error("Invalid number of channels");
    }
}

int format_to_num_channels(juce::Image::PixelFormat format) {
    switch(format) {
        case juce::Image::SingleChannel:
            return 1;
        case juce::Image::RGB:
            return 3;
        case juce::Image::ARGB:
            return 4;
        default:
            throw std::runtime_error("Invalid image format");
    }
}

Note that I just did a quick test with an RGB jpg file and did not check every possible format.

[mfl28]

Hi!
Can you tell us where you are getting errors? In your torch::tensor --> cv::Mat function or in the cv::Mat --> juce::Image function? Which are these errors?

Is make crash in middle of cv:mat to juce image. I dont want OpenCV at all.

Amazing will make test and tell if compile no crash.

[prabhuomkar]

@paklau99988 any update?

[mfl28]

Closing this for now, will be reopened if there are further questions.