ausboss
b62805685b
Load Video and the Clip node divide their fps when they skip frames, but Select Every Nth changed only the batch, so a thinned clip played too fast unless the rate was recomputed by hand. An optional fps socket and an appended fps output now do that division.
168 lines
6.7 KiB
Python
168 lines
6.7 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
import sys
|
|
import unittest
|
|
|
|
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._batch_helpers import (
|
|
merge_batches,
|
|
select_every_nth,
|
|
select_one_based_frame,
|
|
split_batch,
|
|
)
|
|
|
|
|
|
def numbered_batch(count: int, height: int = 2, width: int = 2) -> torch.Tensor:
|
|
"""A BHWC batch where every frame is filled with its one-based number."""
|
|
batch = torch.zeros((count, height, width, 3), dtype=torch.float32)
|
|
for index in range(count):
|
|
batch[index] = float(index + 1)
|
|
return batch
|
|
|
|
|
|
def frame_numbers(batch: torch.Tensor) -> list[float]:
|
|
return [float(frame[0, 0, 0]) for frame in batch]
|
|
|
|
|
|
class SelectFrameTests(unittest.TestCase):
|
|
def test_one_based_selection_returns_the_right_frame(self):
|
|
frames = numbered_batch(5)
|
|
self.assertEqual(float(select_one_based_frame(frames, 1)[0, 0, 0, 0]), 1.0)
|
|
self.assertEqual(float(select_one_based_frame(frames, 5)[0, 0, 0, 0]), 5.0)
|
|
|
|
def test_negative_numbers_count_from_the_end(self):
|
|
frames = numbered_batch(5)
|
|
self.assertEqual(float(select_one_based_frame(frames, -1)[0, 0, 0, 0]), 5.0)
|
|
self.assertEqual(float(select_one_based_frame(frames, -2)[0, 0, 0, 0]), 4.0)
|
|
self.assertEqual(float(select_one_based_frame(frames, -5)[0, 0, 0, 0]), 1.0)
|
|
|
|
def test_selection_stays_a_one_image_batch_either_way(self):
|
|
frames = numbered_batch(3)
|
|
self.assertEqual(tuple(select_one_based_frame(frames, 2).shape), (1, 2, 2, 3))
|
|
self.assertEqual(tuple(select_one_based_frame(frames, -1).shape), (1, 2, 2, 3))
|
|
|
|
def test_out_of_range_reports_the_available_range(self):
|
|
frames = numbered_batch(3)
|
|
for number in (0, 4):
|
|
with self.assertRaisesRegex(ValueError, "frames 1 through 3"):
|
|
select_one_based_frame(frames, number)
|
|
|
|
def test_zero_and_past_the_start_negatives_stay_invalid(self):
|
|
frames = numbered_batch(3)
|
|
for number in (0, -4):
|
|
with self.assertRaisesRegex(ValueError, "-1 through -3"):
|
|
select_one_based_frame(frames, number)
|
|
|
|
|
|
class SelectEveryNthTests(unittest.TestCase):
|
|
def test_keeps_every_other_frame_from_the_first(self):
|
|
frames = numbered_batch(6)
|
|
self.assertEqual(frame_numbers(select_every_nth(frames, 2, 0)), [1.0, 3.0, 5.0])
|
|
|
|
def test_offset_skips_frames_before_the_first_kept_one(self):
|
|
frames = numbered_batch(6)
|
|
self.assertEqual(frame_numbers(select_every_nth(frames, 2, 1)), [2.0, 4.0, 6.0])
|
|
self.assertEqual(frame_numbers(select_every_nth(frames, 3, 2)), [3.0, 6.0])
|
|
|
|
def test_nth_of_one_keeps_everything_after_the_offset(self):
|
|
frames = numbered_batch(4)
|
|
self.assertEqual(frame_numbers(select_every_nth(frames, 1, 2)), [3.0, 4.0])
|
|
|
|
def test_offset_beyond_the_batch_reports_the_count(self):
|
|
frames = numbered_batch(3)
|
|
with self.assertRaisesRegex(ValueError, "only\\s+contains 3"):
|
|
select_every_nth(frames, 2, 3)
|
|
|
|
def test_invalid_nth_and_offset_are_refused(self):
|
|
frames = numbered_batch(3)
|
|
with self.assertRaisesRegex(ValueError, "at least 1"):
|
|
select_every_nth(frames, 0, 0)
|
|
with self.assertRaisesRegex(ValueError, "0 or more"):
|
|
select_every_nth(frames, 2, -1)
|
|
|
|
|
|
class SelectEveryNthNodeTests(unittest.TestCase):
|
|
def test_fps_output_divides_the_wired_rate_by_nth(self):
|
|
from nodes.node_batch_utils import AusBossSelectEveryNth
|
|
|
|
node = AusBossSelectEveryNth()
|
|
kept, rate = node.select(numbered_batch(6), 2, 0, fps=24.0)
|
|
self.assertEqual(frame_numbers(kept), [1.0, 3.0, 5.0])
|
|
self.assertEqual(rate, 12.0)
|
|
# Nothing wired: the frames still come out, and the rate reads 0.
|
|
kept, rate = node.select(numbered_batch(6), 3, 0)
|
|
self.assertEqual(frame_numbers(kept), [1.0, 4.0])
|
|
self.assertEqual(rate, 0.0)
|
|
self.assertEqual(AusBossSelectEveryNth.RETURN_NAMES, ("images", "fps"))
|
|
|
|
|
|
class SplitBatchTests(unittest.TestCase):
|
|
def test_splits_after_the_one_based_index(self):
|
|
frames = numbered_batch(5)
|
|
first, second = split_batch(frames, 2)
|
|
self.assertEqual(frame_numbers(first), [1.0, 2.0])
|
|
self.assertEqual(frame_numbers(second), [3.0, 4.0, 5.0])
|
|
|
|
def test_neither_side_may_be_empty(self):
|
|
frames = numbered_batch(3)
|
|
for index in (0, 3, 4):
|
|
with self.assertRaisesRegex(ValueError, "between 1 and 2"):
|
|
split_batch(frames, index)
|
|
|
|
def test_a_single_frame_cannot_be_split(self):
|
|
with self.assertRaisesRegex(ValueError, "at least 2 frames"):
|
|
split_batch(numbered_batch(1), 1)
|
|
|
|
|
|
class MergeBatchesTests(unittest.TestCase):
|
|
def test_concatenates_a_first(self):
|
|
merged = merge_batches(numbered_batch(2), numbered_batch(3))
|
|
self.assertEqual(frame_numbers(merged), [1.0, 2.0, 1.0, 2.0, 3.0])
|
|
|
|
def test_resize_to_a_brings_b_to_a_size(self):
|
|
batch_a = numbered_batch(2, height=4, width=6)
|
|
batch_b = numbered_batch(1, height=2, width=2)
|
|
merged = merge_batches(batch_a, batch_b, "resize to a")
|
|
self.assertEqual(tuple(merged.shape), (3, 4, 6, 3))
|
|
# A constant frame survives bilinear resizing exactly.
|
|
self.assertEqual(float(merged[2, 0, 0, 0]), 1.0)
|
|
|
|
def test_resize_to_b_brings_a_to_b_size(self):
|
|
batch_a = numbered_batch(2, height=4, width=6)
|
|
batch_b = numbered_batch(1, height=2, width=2)
|
|
merged = merge_batches(batch_a, batch_b, "resize to b")
|
|
self.assertEqual(tuple(merged.shape), (3, 2, 2, 3))
|
|
|
|
def test_error_policy_names_both_sizes(self):
|
|
batch_a = numbered_batch(1, height=4, width=6)
|
|
batch_b = numbered_batch(1, height=2, width=2)
|
|
with self.assertRaisesRegex(ValueError, "6x4.*2x2"):
|
|
merge_batches(batch_a, batch_b, "error")
|
|
|
|
def test_matching_sizes_never_resize(self):
|
|
batch_a = numbered_batch(2)
|
|
batch_b = numbered_batch(2)
|
|
merged = merge_batches(batch_a, batch_b, "error")
|
|
self.assertEqual(tuple(merged.shape), (4, 2, 2, 3))
|
|
|
|
def test_channel_mismatch_is_refused(self):
|
|
batch_a = numbered_batch(1)
|
|
batch_b = torch.zeros((1, 2, 2, 4), dtype=torch.float32)
|
|
with self.assertRaisesRegex(ValueError, "3-channel and 4-channel"):
|
|
merge_batches(batch_a, batch_b)
|
|
|
|
def test_unknown_policy_is_refused(self):
|
|
with self.assertRaisesRegex(ValueError, "mismatch policy"):
|
|
merge_batches(numbered_batch(1), numbered_batch(1), "stretch")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|