168 lines
6.8 KiB
Python
168 lines
6.8 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
|
|
import av
|
|
import numpy as np
|
|
from PIL import Image
|
|
import torch
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(ROOT))
|
|
if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"):
|
|
del sys.modules["nodes"]
|
|
|
|
from nodes._media_helpers import decode_video_frame, video_metadata
|
|
from nodes._transform_engine import (
|
|
TransformSpec,
|
|
stable_file_fingerprint,
|
|
transform_pil,
|
|
transform_tensor_batch,
|
|
)
|
|
|
|
|
|
def solid(width=12, height=8, color=(30, 80, 140, 255)):
|
|
return Image.new("RGBA", (width, height), color)
|
|
|
|
|
|
class TransformEngineTests(unittest.TestCase):
|
|
def test_identity_preserves_dimensions_and_pixels(self):
|
|
source = solid()
|
|
output, mask, geometry = transform_pil(source, TransformSpec())
|
|
self.assertEqual(output.size, source.size)
|
|
self.assertEqual(geometry.output_width, 12)
|
|
self.assertEqual(np.asarray(mask).max(), 0)
|
|
self.assertTrue(np.all(np.asarray(output) == np.array([30, 80, 140])))
|
|
|
|
def test_source_alpha_is_part_of_mask(self):
|
|
source = solid(5, 5)
|
|
source.putpixel((2, 2), (255, 0, 0, 0))
|
|
_, mask, _ = transform_pil(source, TransformSpec())
|
|
self.assertEqual(mask.getpixel((2, 2)), 255)
|
|
self.assertEqual(mask.getpixel((0, 0)), 0)
|
|
|
|
def test_positive_and_negative_rotation_boundaries(self):
|
|
for angle in (-90, 90, 179.9, -179.9):
|
|
with self.subTest(angle=angle):
|
|
output, mask, geometry = transform_pil(solid(11, 7), TransformSpec(rotation_degrees=angle))
|
|
self.assertEqual(output.size, mask.size)
|
|
self.assertGreaterEqual(geometry.rotated_width, 7)
|
|
self.assertGreaterEqual(geometry.rotated_height, 7)
|
|
if abs(angle) % 90 > 0.01:
|
|
self.assertGreater(np.asarray(mask).max(), 0)
|
|
|
|
def test_crop_clamps_to_rotated_source(self):
|
|
output, _, geometry = transform_pil(
|
|
solid(20, 10), TransformSpec(crop_x=18, crop_y=9, crop_width=999, crop_height=999)
|
|
)
|
|
self.assertEqual(output.size, (2, 1))
|
|
self.assertEqual((geometry.crop_x, geometry.crop_y), (18, 9))
|
|
|
|
def test_ratio_is_enforced_in_backend(self):
|
|
output, _, _ = transform_pil(
|
|
solid(100, 100), TransformSpec(crop_aspect_ratio="16:9")
|
|
)
|
|
self.assertEqual(output.size, (100, 56))
|
|
|
|
def test_every_padding_side_changes_output_and_mask(self):
|
|
cases = {
|
|
"pad_left": ((7, 4), (0, 1)),
|
|
"pad_top": ((5, 6), (1, 0)),
|
|
"pad_right": ((7, 4), (6, 1)),
|
|
"pad_bottom": ((5, 6), (1, 5)),
|
|
}
|
|
for field, (expected_size, generated_pixel) in cases.items():
|
|
with self.subTest(field=field):
|
|
output, mask, _ = transform_pil(solid(5, 4), TransformSpec(**{field: 2}))
|
|
self.assertEqual(output.size, expected_size)
|
|
self.assertEqual(mask.getpixel(generated_pixel), 255)
|
|
|
|
def test_fill_color_is_used_for_padding_and_rotation_voids(self):
|
|
output, _, _ = transform_pil(
|
|
solid(4, 3, (255, 0, 0, 255)),
|
|
TransformSpec(rotation_degrees=45, pad_left=2, fill_color="#123456"),
|
|
)
|
|
self.assertEqual(output.getpixel((0, 0)), (0x12, 0x34, 0x56))
|
|
|
|
def test_feather_keeps_generated_pixels_and_softens_boundary(self):
|
|
_, mask, _ = transform_pil(solid(20, 20), TransformSpec(pad_left=10, feather=3))
|
|
values = np.asarray(mask)
|
|
self.assertEqual(values[:, :5].min(), 255)
|
|
self.assertTrue(np.any((values[:, 10:15] > 0) & (values[:, 10:15] < 255)))
|
|
|
|
def test_canvas_multiple_rounds_only_right_and_bottom(self):
|
|
output, mask, geometry = transform_pil(
|
|
solid(101, 99), TransformSpec(canvas_multiple=8)
|
|
)
|
|
self.assertEqual(output.size, (104, 104))
|
|
self.assertEqual((geometry.pad_left, geometry.pad_top), (0, 0))
|
|
self.assertEqual((geometry.pad_right, geometry.pad_bottom), (3, 5))
|
|
self.assertEqual(mask.getpixel((103, 103)), 255)
|
|
|
|
def test_tensor_batch_preserves_bhwc_and_bhw(self):
|
|
source = torch.zeros((3, 9, 7, 3), dtype=torch.float32)
|
|
output, mask, _ = transform_tensor_batch(source, TransformSpec(pad_bottom=1))
|
|
self.assertEqual(tuple(output.shape), (3, 10, 7, 3))
|
|
self.assertEqual(tuple(mask.shape), (3, 10, 7))
|
|
self.assertEqual(output.dtype, source.dtype)
|
|
|
|
def test_invalid_shape_and_oversized_padding_are_actionable(self):
|
|
with self.assertRaisesRegex(ValueError, "expected BHWC"):
|
|
transform_tensor_batch(torch.zeros((4, 4, 3)), TransformSpec())
|
|
with self.assertRaisesRegex(ValueError, "pad_left"):
|
|
transform_pil(solid(), TransformSpec(pad_left=40000))
|
|
|
|
def test_fingerprint_changes_with_inputs_and_file_state(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path = Path(directory, "source.bin")
|
|
path.write_bytes(b"one")
|
|
first = stable_file_fingerprint(path, {"crop": 0})
|
|
second = stable_file_fingerprint(path, {"crop": 1})
|
|
self.assertNotEqual(first, second)
|
|
time.sleep(0.01)
|
|
path.write_bytes(b"two-two")
|
|
os.utime(path, None)
|
|
self.assertNotEqual(first, stable_file_fingerprint(path, {"crop": 0}))
|
|
|
|
|
|
class VideoDecodeTests(unittest.TestCase):
|
|
def _write_video(self, path: Path):
|
|
container = av.open(str(path), mode="w")
|
|
stream = container.add_stream("mpeg4", rate=5)
|
|
stream.width = 16
|
|
stream.height = 16
|
|
stream.pix_fmt = "yuv420p"
|
|
for index in range(5):
|
|
array = np.zeros((16, 16, 3), dtype=np.uint8)
|
|
array[..., 0] = index * 45
|
|
frame = av.VideoFrame.from_ndarray(array, format="rgb24")
|
|
for packet in stream.encode(frame):
|
|
container.mux(packet)
|
|
for packet in stream.encode():
|
|
container.mux(packet)
|
|
container.close()
|
|
|
|
def test_first_middle_last_frame_decoding(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path = Path(directory, "frames.mp4")
|
|
self._write_video(path)
|
|
metadata = video_metadata(path)
|
|
self.assertEqual(metadata["width"], 16)
|
|
self.assertGreaterEqual(metadata["frame_count"], 5)
|
|
reds = []
|
|
for index in (0, 2, 4):
|
|
frame, actual_index, _ = decode_video_frame(path, "frame index", index, 0.0)
|
|
self.assertEqual(actual_index, index)
|
|
reds.append(float(np.asarray(frame)[..., 0].mean()))
|
|
self.assertLess(reds[0], reds[1])
|
|
self.assertLess(reds[1], reds[2])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|