python - 使用 opencv 加载的 Torch cyclegan 模型不输出所需的图像

Question

我正在尝试在 OpenCV 中加载预训练的 torch (.t7) 模型。该模型是一个 CycleGAN，可将马图像转换为斑马图像。该模型可以在这里找到:https://people.eecs.berkeley.edu/~taesung_park/CycleGAN/models/

我使用过 horse2zebra.t7 和 horse2zebra_cpu.t7 模型，但它们都返回平铺的黑白图像而不是斑马图像。

这是一个示例输入图像:

这是输出:

代码:

import cv2
import numpy as np

model = cv2.dnn.readNetFromTorch('./cyclegan_horse2zebra_cpu.t7')

image = cv2.imread('./images/1.jpg')

blob = cv2.dnn.blobFromImage(image, 1, (256, 256))
model.setInput(blob)

out = model.forward()[0,:,:,:]
out = np.reshape(out, (256, 256, 3))

cv2.imshow('image', out)
cv2.waitKey(0)

cv2.imwrite('out.png', out)

Answer 1

遗漏了两件事。第一个是 model.forward()[0,:,:,:] 返回 3x256x256 3D blob(平面图像， channel 在第二维)。但是，OpenCV 处理交错图像，因此您需要通过 np.transpose 排列尺寸:256x256x3。

输出值范围也很重要。对于问题中的图像，我从 [-0.832621, 0.891473] 中获取了值。我们需要将其规范化为 Uint8 类型的 [0, 255]。所有这一切都是由 cv::normalize 制造的.

import cv2
import numpy as np

model = cv2.dnn.readNetFromTorch('./horse2zebra_cpu.t7')

image = cv2.imread('./images/1.jpg')

blob = cv2.dnn.blobFromImage(image, 1, (256, 256))
model.setInput(blob)

out = model.forward()[0,:,:,:]
out = cv2.normalize(out, dst=None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)

out = np.transpose(out, (1, 2, 0))

cv2.imshow('image', out)
cv2.waitKey(0)

cv2.imwrite('out.png', out)

请注意，原始框架可能有不同的后处理过程。如果您可以比较 OpenCV 和 Torch 的输出，那就太好了。这是我从上面的脚本中得到的图像: