62 lines
2.7 KiB
Python
62 lines
2.7 KiB
Python
from io import BytesIO
|
|
import unittest
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
import torch
|
|
|
|
from plugin_loader import plugin
|
|
from chicktools.chick.codec import CarrierCodecError, decode_bytes, encode_bytes, image_to_png_bytes
|
|
|
|
|
|
class CarrierCodecTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.source = Image.new("RGB", (8, 8), (12, 34, 56))
|
|
output = BytesIO()
|
|
self.source.save(output, format="PNG")
|
|
self.raw = output.getvalue()
|
|
|
|
def test_round_trip_all_lsb_widths(self):
|
|
for compress in (2, 6, 8):
|
|
with self.subTest(compress=compress):
|
|
carrier = encode_bytes(self.raw, "png", compress=compress, title="sample")
|
|
self.assertGreaterEqual(carrier.width, 640)
|
|
recovered, extension = decode_bytes(carrier)
|
|
self.assertEqual((recovered, extension), (self.raw, "png"))
|
|
|
|
def test_password_round_trip_and_rejection(self):
|
|
carrier = encode_bytes(self.raw, "png", password="local-test", compress=6)
|
|
self.assertEqual(decode_bytes(carrier, "local-test"), (self.raw, "png"))
|
|
with self.assertRaises(CarrierCodecError):
|
|
decode_bytes(carrier)
|
|
with self.assertRaises(CarrierCodecError):
|
|
decode_bytes(carrier, "wrong")
|
|
|
|
def test_large_payload_across_bit_chunk_boundaries_and_png_save(self):
|
|
raw = bytes(range(256)) * 1200 + b"\x00\x00"
|
|
for width in (2, 6, 8):
|
|
with self.subTest(width=width):
|
|
carrier = encode_bytes(raw, "bin", password="chunks", compress=width)
|
|
output = BytesIO()
|
|
carrier.save(output, format="PNG")
|
|
with Image.open(BytesIO(output.getvalue())) as saved:
|
|
self.assertEqual(decode_bytes(saved, "chunks"), (raw, "bin"))
|
|
|
|
def test_image_conversion_rejects_batches(self):
|
|
with self.assertRaises(CarrierCodecError):
|
|
image_to_png_bytes(np.zeros((2, 8, 8, 3), dtype=np.float32))
|
|
|
|
def test_node_encodes_a_real_image_batch(self):
|
|
source = torch.zeros((1, 8, 8, 3), dtype=torch.float32)
|
|
source[:, :, :, 0] = 0.25
|
|
carrier = plugin.NODE_CLASS_MAPPINGS["ChickHideNode"]().encode(
|
|
password="node-test", title="", fps=24, compress=2, combine_video=False, images=source
|
|
)[0]
|
|
self.assertEqual(tuple(carrier.shape[:1]), (1,))
|
|
carrier_image = Image.fromarray(np.rint(carrier[0].numpy() * 255).astype(np.uint8), mode="RGB")
|
|
recovered, extension = decode_bytes(carrier_image, "node-test")
|
|
self.assertEqual(extension, "png")
|
|
with Image.open(BytesIO(recovered)) as restored:
|
|
self.assertEqual(restored.size, (8, 8))
|
|
self.assertEqual(restored.getpixel((0, 0)), (64, 0, 0))
|