Загруженная с помощью opencv модель циклогана Torch не выдает нужного изображения - PullRequest
0 голосов
/ 17 января 2019

Я пытаюсь загрузить предварительно обученную модель резака (.t7) в OpenCV. Модель CycleGAN, которая преобразует изображения лошадей в изображения зебры. Модель можно найти здесь: https://people.eecs.berkeley.edu/~taesung_park/CycleGAN/models/

Я использовал обе модели horse2zebra.t7 и horse2zebra_cpu.t7, но они оба возвращают мозаичное черно-белое изображение вместо изображения зебры.

Это пример входного изображения:

enter image description here

И это вывод:

enter image description here

Код:

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)

1 Ответ

0 голосов
/ 20 января 2019

Есть две вещи, пропущенные. Во-первых, model.forward()[0,:,:,:] возвращает 3x256x256 3D BLOB-объект (плоское изображение с каналами во втором измерении). Однако OpenCV работает с чередующимися изображениями, поэтому вам нужно переставить размеры на np.transpose: 256x256x3.

Диапазон выходных значений также важен. Для изображения из вопроса я получил значения от [-0.832621, 0.891473]. Нам нужно нормализовать его до [0, 255] типа Uint8. все это одно сделано 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. Это изображение, которое я получил из сценария выше:

enter image description here

...