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;
}
原文:https://www.cnblogs.com/cheungxiongwei/p/11836869.html