QImage 如何和 Tensor 相互转换?

时间:2019-11-11
本文章向大家介绍QImage 如何和 Tensor 相互转换?,主要包括QImage 如何和 Tensor 相互转换?使用实例、应用技巧、基本知识点总结和需要注意事项,具有一定的参考价值,需要的朋友可以参考一下。
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