将跟踪模型 (.pt) 与 C++ 和 OpenCV 结合使用时输出图像的差异

Question

我已经根据EnlightenGAN重新训练了模型。此外，我还跟踪了模型，以便使用 libTorch v1.6 在 C++ 应用程序中执行它。但是，与 python（执行跟踪模型）版本相比，我得到的结果略有不同。

该模型需要输入 RGB 张量和注意力图图像张量作为输入。注意图基本上是告知模型需要对比度增强的图像区域。

以下代码用于推断 python 中 PT 模型的输出。


def getTransform():
    transform_list = []
    transform_list += [transforms.ToTensor(),
                       transforms.Normalize((0.5, 0.5, 0.5),
                                            (0.5, 0.5, 0.5))]
    return transforms.Compose(transform_list)

def convertToCV(tensor):
    
    tensor = torch.squeeze(tensor)
    tensor = tensor.cpu().float().detach()
    tensor = torch.unsqueeze(tensor, 0)
    tensor = tensor.permute(1, 2, 0)
    tensor = ((tensor +1)/2.0) * 255.0
    tensor = tensor.numpy()
    return tensor

def proprocess(image):
    
    transform = getTransform()
    trgbImage = transform(image)
    r,g,b = trgbImage[0]+1, trgbImage[1]+1, trgbImage[2]+1
    tattentionImage = 1. - ((0.299*r+0.587*g+0.114*b)/2.)
    tattentionImage = torch.unsqueeze(tattentionImage, 0)
    trgbImage = torch.unsqueeze(trgbImage, 0)
    tattentionImage = torch.unsqueeze(tattentionImage, 0)

    return trgbImage, tattentionImage

def run(inputPath, OutputPath):

    modelToLoad = torch.jit.load("./EGAN.pt")
    print("OK")
    count =0 
    for filename in os.listdir(inputPath):
        
        print("Processing Image : ", filename)
        inputImage = cv2.imread(os.path.join(inputPath,filename))
        
        rgbImage, attentionImage = proprocess(inputImage)
        
        fake, real = modelToLoad.forward(rgbImage,attentionImage )
        
        fake_B = convertToCV(fake)
        fake_B1 = cv2.cvtColor(fake_B, cv2.COLOR_RGB2BGR)
        cv2.imwrite(OutputPath + "pic1.png" , fake_B )

推理代码的 C++ 版本如下

#define A 0.299
#define B 0.5870
#define C 0.114

cv::Mat torchTensortoCVMat1C(torch::Tensor& tensor)
{
    tensor = tensor.squeeze(0);
    tensor = tensor.to(torch::kCPU).to(torch::kFloat32).detach();
    tensor = tensor.permute({ 1, 2, 0 }).contiguous();
    tensor = tensor.mul(0.5).add(0.5).mul(255.0);
    tensor = tensor.to(torch::kU8);

    int64_t height = tensor.size(0);
    int64_t width  = tensor.size(1);
    cv::Mat mat    = cv::Mat(cv::Size(width, height), CV_8UC3, tensor.data_ptr<uchar>());
    return mat.clone();
}

std::vector<torch::jit::IValue> CV2Tensor(const cv::Mat& cv_Image)
{
    torch::Tensor tInputImage = (torch::from_blob(cv_Image.data, { cv_Image.rows, cv_Image.cols, cv_Image.channels() }, torch::kByte));
    tInputImage = tInputImage.to(torch::kFloat).div(255);
    tInputImage = tInputImage.sub(0.5).div(0.5).permute({ 2, 0, 1 });

    torch::Tensor red   = tInputImage[0] + 1 ;
    torch::Tensor green = tInputImage[1] + 1 ;
    torch::Tensor blue  = tInputImage[2] + 1;

    red   = red.mul(A);
    green = green.mul(B);
    blue  = blue.mul(C);

    torch::Tensor channelSum = red.add(green).add(blue);
    channelSum = channelSum.div(2.);
    torch::Tensor tGrayImage = 1. - channelSum;

    tGrayImage.unsqueeze_(0);
    tGrayImage.unsqueeze_(0);
    tInputImage.unsqueeze_(0);

    std::vector<torch::jit::IValue> input;
    input.push_back(tInputImage);
    input.push_back(tGrayImage);

    return input;
}

void enhanceImage(const std::string& Img, torch::jit::script::Module& network,const std::string& outputPath, std::string& fileName)
{
    cv::Mat rgbImage;
    cv::Mat inputImage = cv::imread(Img);
    std::vector<torch::jit::IValue> input = CV2Tensor(inputImage);

    try
    {
        auto outputs = network.forward(input).toTuple();
        torch::Tensor resultFake = outputs->elements()[0].toTensor();
    
        cv::Mat output1 = torchTensortoCVMat(resultFake);
        cv::imshow("out1.png", output1);
        cv::waitKey(0);
    }
    catch (std::exception& e)
    {
        std::cout << e.what() << std::endl;
    }
}

我也检查了所有步骤的张量输出，它们是一样的。然而，在转换之后，输出图像的颜色从输入图像的较亮区域流出，如下所示。

Python 版本 C++ 版本

我尝试了很多次，但我完全不知道该如何解决这个问题。欢迎任何帮助。

谢谢。 PS : 如果需要更多信息，请告诉我。

Answer 1

我在上面的实现中发现了问题。在 CPP 版本中，我在去规范化后没有限制值。我已经设置了一个钳位功能，现在它按预期工作。如果有人遇到同样的问题，编辑部分如下：

tensor = tensor.mul(0.5).add(0.5).mul(255.0); -- > tensor = tensor.mul(0.5).add(0.5).mul(255.0).clamp(0, 255);

没有限制它，因为它会导致图像较亮区域溢出。

将跟踪模型 (.pt) 与 C++ 和 OpenCV 结合使用时输出图像的差异

Difference In the output image when using traced model(.pt) with C++ and OpenCV

c++

python

pytorch

libtorch