QImage 如何和 Tensor 相互转换?
torch::Tensor fromQImage(QImage image)
{
int width = image.width();
int height = image.height();
int depth = image.depth();
int channels = depth / 8;
const torch::TensorOptions option(torch::kUInt8);
torch::Tensor tensor = torch::from_blob(image.bits(),{width * height,channels},option);//R G B A
auto result = torch::zeros({1,channels,height*width});//N C H C
if(channels == 4){
/*!
R G B A
R G B A
R G B A
=>
R R R
G G G
B B B
A A A
*/
tensor = tensor.transpose(0,1);
auto R = tensor[0];
auto G = tensor[1];
auto B = tensor[2];
auto A = tensor[3];
result[0][0] = B;
result[0][1] = G;
result[0][2] = R;
result[0][3] = A;
result = result.view({1,channels,height,width}).div(255.0);//N C H C
//std::cout << result << std::endl;
return result;
}
if(channels == 3){
/*!
R G B
R G B
R G B
=>
R R R
G G G
B B B
*/
tensor = tensor.transpose(0,1);
auto R = tensor[0];
auto G = tensor[1];
auto B = tensor[2];
result[0][0] = B;
result[0][1] = G;
result[0][2] = R;
result = result.view({1,channels,height,width}).div(255.0);//N C H C
//std::cout << result << std::endl;
return result;
}
if(channels == 1){
return result;
}
return result;
}
QImage TensorToQImage(const torch::Tensor &tensor)
{
QImage image;
int dim = tensor.dim();
if(dim != 4){
qFatal("dim must be 4.");
}
//std::cout << tensor.size(0) << tensor.size(1) << tensor.size(2) << tensor.size(3) << std::endl;
int channels = tensor.size(1);
int width = tensor.size(3);
int height = tensor.size(2);
// fill QImage
if(channels == 1){
#pragma omp simd
image = QImage(width,height,QImage::Format_Grayscale8);
for(int w = 0;w < width;++w){
for(int h = 0;h < height;++h){
QRgb gray = tensor[0][0][h][w].item<float>() * 255.0;
image.setPixel(w,h,gray);
}
}
}
// fill QImage
if(channels == 3){
image = QImage(width,height,QImage::Format_RGB888);
#pragma omp simd
for(int w = 0;w < width;++w){
for(int h = 0;h < height;++h){
int r = tensor[0][0][h][w].item<float>() * 255.0;
int g = tensor[0][1][h][w].item<float>() * 255.0;
int b = tensor[0][2][h][w].item<float>() * 255.0;
QRgb rgb = qRgb(r,g,b);
image.setPixel(w,h,rgb);
}
}
}
// fill QImage
if(channels == 4){
image = QImage(width,height,QImage::Format_RGB32);
#pragma omp simd
for(int w = 0;w < width;++w){
for(int h = 0;h < height;++h){
int r = tensor[0][0][h][w].item<float>() * 255.0;
int g = tensor[0][1][h][w].item<float>() * 255.0;
int b = tensor[0][2][h][w].item<float>() * 255.0;
int a = tensor[0][3][h][w].item<float>() * 255.0;
QRgb rgb = qRgba(r,g,b,a);
image.setPixel(w,h,rgb);
}
}
}
return QImage();
}
/*!
QImage to torch::Tensor N x C x H x W
*/
torch::Tensor QImageToTensor(const QImage &image)
{
int width = image.width();
int height = image.height();
int depth = image.depth();
int channels = depth / 8;
// create tensor
torch::TensorOptions option(torch::kFloat32);
torch::Tensor tensor = torch::zeros({1,channels,height,width},option);//N C H W
bool isOk = false;
// fill tensor
if(channels == 1){
#pragma omp simd
for(int w = 0;w < width;++w){
for(int h = 0;h < height;++h){
QRgb rgb = image.pixel(w,h);
tensor[0][0][h][w] = qGray(rgb)/255.0;//GRAY
}
}
isOk = true;
}
// fill tensor
if(channels == 3){
#pragma omp simd
for(int w = 0;w < width;++w){
for(int h = 0;h < height;++h){
QRgb rgb = image.pixel(w,h);
tensor[0][0][h][w] = qRed(rgb)/255.0;//R
tensor[0][1][h][w] = qGreen(rgb)/255.0;//G
tensor[0][2][h][w] = qBlue(rgb)/255.0;//B
}
}
isOk = true;
}
// fill tensor
if(channels == 4){
#pragma omp simd
for(int w = 0;w < width;++w){
for(int h = 0;h < height;++h){
QRgb rgb = image.pixel(w,h);
tensor[0][0][h][w] = qRed(rgb)/255.0;//R
tensor[0][1][h][w] = qGreen(rgb)/255.0;//G
tensor[0][2][h][w] = qBlue(rgb)/255.0;//B
tensor[0][3][h][w] = qAlpha(rgb)/255.0;//A
}
}
isOk = true;
}
if(!isOk){
qFatal("channels must be 1, 3, or 4.");
}
//std::cout << tensor << std::endl;
return tensor;
}
转载请注明出处并保持作品的完整性,谢谢