def test__encode_decode_quantize_true(self, mock__random_data):
"""Test apply quantize and then call decode"""
# setup
random_data_fixture = torch.load(os.path.join(self.fixtures_path, 'tensors/random.ts'))
mock__random_data.return_value = random_data_fixture
cover_fixture = torch.load(os.path.join(self.fixtures_path, 'tensors/img.ts'))
steganogan = self.VoidSteganoGAN()
steganogan.encoder = MagicMock(return_value=cover_fixture)
steganogan.decoder = MagicMock(return_value=MagicMock())
# run
generated, payload, decoded = steganogan._encode_decode(cover_fixture, quantize=True)
# assert
expected_quantize_cover = (255 * (cover_fixture + 1.0) / 2.0).long()
expected_quantize_cover = 2.0 * expected_quantize_cover.float() / 255.0 - 1.0
mock__random_data.called_once_with(cover_fixture)
steganogan.encoder.assert_called_once_with(cover_fixture, random_data_fixture)
assert_called_with_tensors(steganogan.decoder, [call(expected_quantize_cover)])
assert (generated == expected_quantize_cover).all()
assert (payload == random_data_fixture).all()
assert decoded == steganogan.decoder.return_value
def test_forward_1_layer(self):
"""If there is only one layer it must be called with image as the only argument."""
# setup
layer1 = Mock(return_value=torch.Tensor([[5, 6], [7, 8]]))
self.test_encoder._models = [layer1]
# run
image = torch.Tensor([[1, 2], [3, 4]])
result = self.test_encoder.forward(image, 'some_data')
# assert
assert (result == torch.Tensor([[5, 6], [7, 8]])).all()
call_1 = call(torch.Tensor([[1, 2], [3, 4]]))
assert_called_with_tensors(layer1, [call_1])
def test_forward_more_than_2_layers(self):
"""If there are more than 2 layers, they must be called adding data to each result"""
# setup
layer1 = Mock(return_value=torch.Tensor([[5, 6], [7, 8]]))
layer2 = Mock(return_value=torch.Tensor([[9, 10], [11, 12]]))
layer3 = Mock(return_value=torch.Tensor([[13, 14], [15, 16]]))
self.test_encoder._models = [layer1, layer2, layer3]
# run
image = torch.Tensor([[1, 2], [3, 4]])
data = torch.Tensor([[-1, -2], [-3, -4]])
result = self.test_encoder.forward(image, data)
# asserts
call_layer_1 = call(torch.Tensor([[1, 2], [3, 4]]))
call_layer_2 = call(torch.Tensor([[5, 6, -1, -2], [7, 8, -3, -4]]))
call_layer_3 = call(
torch.Tensor([[5, 6, 9, 10, -1, -2], [7, 8, 11, 12, -3, -4]]))
assert_called_with_tensors(layer1, [call_layer_1])
assert_called_with_tensors(layer2, [call_layer_2])
assert_called_with_tensors(layer3, [call_layer_3])
assert (result == torch.Tensor([[13, 14], [15, 16]])).all()
def test_forward_add_image(self):
"""If add_image is true, image must be added to the result."""
# setup
self.test_encoder.add_image = True
layer1 = Mock(return_value=torch.Tensor([[5, 6], [7, 8]]))
layer2 = Mock(return_value=torch.Tensor([[9, 10], [11, 12]]))
layer3 = Mock(return_value=torch.Tensor([[13, 14], [15, 16]]))
self.test_encoder._models = [layer1, layer2, layer3]
# run
image = torch.Tensor([[1, 2], [3, 4]])
data = torch.Tensor([[-1, -2], [-3, -4]])
result = self.test_encoder.forward(image, data)
# asserts
call_layer_1 = call(torch.Tensor([[1, 2], [3, 4]]))
call_layer_2 = call(torch.Tensor([[5, 6, -1, -2], [7, 8, -3, -4]]))
call_layer_3 = call(
torch.Tensor([[5, 6, 9, 10, -1, -2], [7, 8, 11, 12, -3, -4]]))
assert_called_with_tensors(layer1, [call_layer_1])
assert_called_with_tensors(layer2, [call_layer_2])
assert_called_with_tensors(layer3, [call_layer_3])
assert (result == torch.Tensor([[14, 16], [18, 20]])).all()