397 lines
20 KiB
Python
397 lines
20 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import inspect
|
|
import io
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
import unittest.mock
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(ROOT))
|
|
sys.path.insert(0, str(ROOT / "tests"))
|
|
if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"):
|
|
del sys.modules["nodes"]
|
|
|
|
from nodes._transform_engine import (
|
|
TransformSpec,
|
|
transform_tensor_batch,
|
|
transform_tensor_batch_chunked,
|
|
)
|
|
from nodes._inpaint_crop_helpers import STITCHER_KIND, apply_stitch
|
|
from nodes._transform_inputs import resize_inputs, transform_inputs
|
|
from nodes.node_inpaint_crop_stitch import AusBossStitchInpaint
|
|
from nodes.node_video_crop_rotate_pad_clip import AusBossVideoCropRotatePadClip, snap_frame_count
|
|
from test_video_load_helpers import FPS, FRAMES, HEIGHT, WIDTH, ensure_core_video_api, write_test_video
|
|
|
|
|
|
def transform_defaults(**overrides) -> dict:
|
|
values = {name: spec[1]["default"] for name, spec in transform_inputs().items()}
|
|
values.update({name: spec[1]["default"] for name, spec in resize_inputs().items()})
|
|
values.update(overrides)
|
|
return values
|
|
|
|
|
|
class ClipDefaultsTests(unittest.TestCase):
|
|
def test_a_fresh_clip_node_is_outpaint_ready(self):
|
|
# The in-context video models the clip feeds paint pure black behind
|
|
# a hard edge; a grey or feathered band comes back untouched.
|
|
required = AusBossVideoCropRotatePadClip.INPUT_TYPES()["required"]
|
|
self.assertEqual(required["feather"][1]["default"], 0)
|
|
self.assertEqual(required["fill_color"][1]["default"], "#000000")
|
|
self.assertEqual(transform_inputs()["feather"][1]["default"], 24)
|
|
|
|
|
|
class ChunkedTransformTests(unittest.TestCase):
|
|
def test_chunks_match_the_single_pass_transform(self):
|
|
torch.manual_seed(7)
|
|
batch = torch.rand((5, 40, 56, 3))
|
|
spec = TransformSpec(rotation_degrees=12.5, pad_left=9, pad_top=4, feather=6, canvas_multiple=8)
|
|
whole_out, whole_mask, whole_geometry = transform_tensor_batch(batch, spec)
|
|
out, mask, geometry = transform_tensor_batch_chunked(batch, spec, chunk_size=2)
|
|
self.assertEqual(geometry, whole_geometry)
|
|
self.assertTrue(torch.equal(out, whole_out))
|
|
self.assertTrue(torch.equal(mask, whole_mask))
|
|
|
|
def test_empty_batch_is_reported(self):
|
|
with self.assertRaisesRegex(ValueError, "no decodable frames"):
|
|
transform_tensor_batch_chunked(torch.zeros((0, 8, 8, 3)), TransformSpec())
|
|
|
|
|
|
class VideoClipNodeTests(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls._tmp = tempfile.TemporaryDirectory()
|
|
cls.video = Path(cls._tmp.name) / "clip.mp4"
|
|
write_test_video(cls.video, with_audio=True)
|
|
# The second source reads a video by the name ComfyUI gives it below
|
|
# its input, output or temp folder; the fixture's temporary folder
|
|
# stands in for temp.
|
|
from nodes import _media_helpers
|
|
|
|
root = Path(cls._tmp.name).resolve()
|
|
cls.server_file = f"{cls.video.name} [temp]"
|
|
cls._env = unittest.mock.patch.object(_media_helpers, "_server_folders", lambda: {"temp": root})
|
|
cls._env.start()
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
cls._env.stop()
|
|
cls._tmp.cleanup()
|
|
|
|
def run_node(self, **overrides):
|
|
values = transform_defaults()
|
|
args = {
|
|
"video": "",
|
|
"source_mode": "local path",
|
|
"local_path": self.server_file,
|
|
"start_seconds": 0.0,
|
|
"end_seconds": 0.0,
|
|
**values,
|
|
}
|
|
args.update(overrides)
|
|
return asyncio.run(AusBossVideoCropRotatePadClip().load_transform(**args))
|
|
|
|
def test_linked_frame_bounds_override_stale_timeline_and_cap(self):
|
|
out = self.run_node(start_seconds=9, end_seconds=2, start_frame=6, end_frame=18,
|
|
max_frames=1, frame_load_cap=4)
|
|
self.assertEqual(out[3], 4)
|
|
self.assertAlmostEqual(out[7], 4 / FPS)
|
|
reference = self.run_node(start_seconds=6 / FPS, end_seconds=18 / FPS, max_frames=4)
|
|
self.assertTrue(torch.equal(out[0], reference[0]))
|
|
|
|
def test_force_rate_preserves_time_drops_or_repeats_frames(self):
|
|
original = self.run_node()
|
|
for rate in (FPS / 2, FPS * 2):
|
|
out = self.run_node(force_rate=rate)
|
|
self.assertEqual(out[3], int(FRAMES * rate / FPS))
|
|
self.assertEqual(out[4], rate)
|
|
self.assertAlmostEqual(out[7], original[7])
|
|
if rate > FPS:
|
|
self.assertTrue(torch.equal(out[0][0], out[0][1]))
|
|
else:
|
|
self.assertTrue(torch.equal(out[0][1], original[0][2]))
|
|
|
|
def test_force_rate_cap_and_thinning_apply_before_snap(self):
|
|
out = self.run_node(start_frame=3, end_frame=21, force_rate=24,
|
|
every_nth=2, frame_load_cap=10, frame_snap="4n+1")
|
|
self.assertEqual(out[3], 9)
|
|
self.assertEqual(out[4], 12)
|
|
self.assertAlmostEqual(out[7], 0.75)
|
|
self.assertAlmostEqual(out[2]["waveform"].shape[-1] / out[2]["sample_rate"], 0.75, delta=0.01)
|
|
|
|
def test_validation_reads_only_the_source(self):
|
|
# ComfyUI files a failed check once per input VALIDATE_INPUTS names
|
|
# and skips its own range, list and link checks for them, so only the
|
|
# source is named. A stale widget window beside a linked bound can
|
|
# no longer block the run; the window is checked when the node runs.
|
|
validate = AusBossVideoCropRotatePadClip.VALIDATE_INPUTS
|
|
self.assertEqual(list(inspect.signature(validate).parameters), ["video", "source_mode", "local_path"])
|
|
self.assertIs(validate("", "local path", self.server_file), True)
|
|
self.assertEqual(
|
|
validate("", "local path", ""),
|
|
"Video Crop + Rotate + Pad: Pick a video from the node's Outputs list first.",
|
|
)
|
|
|
|
def test_decode_errors_name_this_node_and_only_its_inputs(self):
|
|
# The shared decode used to report "Load Video" and advise
|
|
# custom_width/custom_height, which this node does not have.
|
|
with self.assertRaisesRegex(ValueError, r"^Video Crop \+ Rotate \+ Pad -> Clip starts at 99\.00s"):
|
|
self.run_node(start_seconds=99.0)
|
|
from nodes import _video_load_helpers
|
|
|
|
with unittest.mock.patch.object(_video_load_helpers, "_available_memory_bytes", return_value=1000):
|
|
with self.assertRaises(ValueError) as caught:
|
|
self.run_node()
|
|
message = str(caught.exception)
|
|
self.assertTrue(message.startswith("Video Crop + Rotate + Pad -> Clip would need"), message)
|
|
self.assertIn("every_nth", message)
|
|
self.assertNotIn("custom_width", message)
|
|
message.encode("ascii")
|
|
|
|
def test_invalid_force_rate_and_backwards_frame_bounds_fail(self):
|
|
for rate in (-1, float("nan"), 1001):
|
|
with self.assertRaises(ValueError):
|
|
self.run_node(force_rate=rate)
|
|
with self.assertRaises(ValueError):
|
|
self.run_node(start_frame=18, end_frame=6)
|
|
|
|
def test_fixed_frames_move_as_one_window_and_override_limits(self):
|
|
out = self.run_node(start_seconds=99, end_seconds=0.1, end_frame=1,
|
|
fixed_frames=12, frame_load_cap=2, max_frames=1, frame_snap="8n+1")
|
|
reference = self.run_node(start_seconds=1, end_seconds=2)
|
|
self.assertEqual(out[3], 12)
|
|
self.assertTrue(torch.equal(out[0], reference[0]))
|
|
self.assertTrue(torch.equal(out[9], reference[9]))
|
|
self.assertEqual(out[8]["canvas"].shape[0], 12)
|
|
self.assertAlmostEqual(out[7], 1)
|
|
self.assertAlmostEqual(out[2]["waveform"].shape[-1] / out[2]["sample_rate"], 1, delta=0.01)
|
|
|
|
def test_fixed_frames_uses_output_rate_after_thinning(self):
|
|
for rate, nth, count in [(24, 2, 12), (24, 1, 24), (0, 2, 6), (10, 1, 9)]:
|
|
out = self.run_node(start_frame=23, fixed_frames=count, force_rate=rate, every_nth=nth)
|
|
self.assertEqual(out[3], count)
|
|
self.assertEqual(out[4], (rate or FPS) / nth)
|
|
self.assertAlmostEqual(out[7], count / out[4])
|
|
self.assertAlmostEqual(out[2]["waveform"].shape[-1] / out[2]["sample_rate"], out[7], delta=0.01)
|
|
|
|
def test_last_seconds_loads_only_the_end(self):
|
|
whole = self.run_node()
|
|
tail = self.run_node(last_seconds=6 / FPS)
|
|
self.assertEqual(tail[3], 6)
|
|
self.assertTrue(torch.equal(tail[0], whole[0][-6:]))
|
|
self.assertEqual(self.run_node(last_seconds=0)[3], FRAMES)
|
|
|
|
def test_last_seconds_counts_back_from_out_and_stops_at_in(self):
|
|
whole = self.run_node()
|
|
before_out = self.run_node(end_seconds=12 / FPS, last_seconds=4 / FPS)
|
|
self.assertTrue(torch.equal(before_out[0], whole[0][8:12]))
|
|
self.assertEqual(self.run_node(start_seconds=FRAMES / FPS - 3 / FPS, last_seconds=9 / FPS)[3], 3)
|
|
self.assertEqual(self.run_node(last_seconds=9999)[3], FRAMES)
|
|
|
|
def test_fixed_frames_too_long_is_explicit_and_zero_preserves_free_trim(self):
|
|
with self.assertRaisesRegex(ValueError, "source is only"):
|
|
self.run_node(fixed_frames=FRAMES + 1)
|
|
free = self.run_node(start_seconds=0.5, end_seconds=1, fixed_frames=0)
|
|
self.assertEqual(free[3], 6)
|
|
# A backwards free window is ignored once Fixed frames places its own.
|
|
self.assertEqual(self.run_node(start_seconds=9, end_seconds=2, fixed_frames=12)[3], 12)
|
|
|
|
def test_padding_applies_to_every_frame_and_lands_in_the_mask(self):
|
|
frames, mask, audio, count, fps, width, height, duration = self.run_node(pad_left=16, feather=0)[:8]
|
|
self.assertEqual(tuple(frames.shape), (FRAMES, HEIGHT, WIDTH + 16, 3))
|
|
self.assertEqual(tuple(mask.shape), (FRAMES, HEIGHT, WIDTH + 16))
|
|
self.assertTrue(torch.all(mask[:, :, :16] == 1.0))
|
|
self.assertTrue(torch.all(mask[:, :, 16:] == 0.0))
|
|
# Fill color #808080 in the new band, on every frame.
|
|
self.assertTrue(torch.allclose(frames[:, :, :16], torch.full((FRAMES, HEIGHT, 16, 3), 128 / 255)))
|
|
self.assertEqual((count, width, height), (FRAMES, WIDTH + 16, HEIGHT))
|
|
self.assertAlmostEqual(fps, float(FPS))
|
|
self.assertAlmostEqual(duration, FRAMES / FPS)
|
|
self.assertGreater(audio["waveform"].shape[-1], 0)
|
|
|
|
def test_divisible_by_grows_no_strip_on_an_unpadded_edge(self):
|
|
# Padded on the left only: the leftover joins that band, and the
|
|
# height, which nobody padded, is trimmed to the multiple instead.
|
|
multiple = 24
|
|
frames, mask, _, count, _, width, height, _ = self.run_node(pad_left=3, feather=0, canvas_multiple=multiple)[:8]
|
|
band = -(-(WIDTH + 3) // multiple) * multiple - WIDTH
|
|
self.assertEqual((width, height), (WIDTH + band, HEIGHT // multiple * multiple))
|
|
self.assertEqual(tuple(frames.shape), (count, height, width, 3))
|
|
self.assertTrue(torch.all(mask[:, :, :band] == 1.0))
|
|
self.assertTrue(torch.all(mask[:, :, band:] == 0.0))
|
|
|
|
def test_trim_thinning_and_crop_compose(self):
|
|
frames, mask, _, count, fps, width, height, duration = self.run_node(
|
|
start_seconds=0.5, end_seconds=1.5, every_nth=2, crop_x=8, crop_width=32, feather=0
|
|
)[:8]
|
|
self.assertEqual(count, 6)
|
|
self.assertEqual(tuple(frames.shape), (6, HEIGHT, 32, 3))
|
|
self.assertEqual((width, height), (32, HEIGHT))
|
|
self.assertAlmostEqual(fps, FPS / 2)
|
|
self.assertAlmostEqual(duration, 1.0)
|
|
self.assertEqual(float(mask.max()), 0.0)
|
|
|
|
def test_preview_position_never_changes_the_fingerprint(self):
|
|
common = dict(video="", source_mode="local path", local_path=self.server_file, start_seconds=0.0, end_seconds=0.0)
|
|
values = transform_defaults(pad_right=8)
|
|
one = AusBossVideoCropRotatePadClip.IS_CHANGED(**common, frame_index=0, frame_time=0.0, **values)
|
|
two = AusBossVideoCropRotatePadClip.IS_CHANGED(**common, frame_index=17, frame_time=1.4, seek_mode="time seconds", **values)
|
|
self.assertEqual(one, two)
|
|
trimmed = AusBossVideoCropRotatePadClip.IS_CHANGED(**{**common, "end_seconds": 1.0}, **values)
|
|
self.assertNotEqual(one, trimmed)
|
|
padded = AusBossVideoCropRotatePadClip.IS_CHANGED(**common, **transform_defaults(pad_right=9))
|
|
self.assertNotEqual(one, padded)
|
|
|
|
def test_a_backwards_window_stops_the_node_when_it_runs(self):
|
|
with self.assertRaisesRegex(ValueError, "needs start_seconds smaller than end_seconds"):
|
|
self.run_node(start_seconds=2.0, end_seconds=1.0)
|
|
|
|
def test_pad_fit_preserves_source_pixels_and_mask(self):
|
|
original = self.run_node(feather=0)[0]
|
|
frames, mask, _, _, _, width, height, _ = self.run_node(
|
|
crop_aspect_ratio="free", crop_width=WIDTH, crop_height=HEIGHT,
|
|
pad_top=17, pad_bottom=18, feather=0,
|
|
)[:8]
|
|
self.assertEqual((width, height), (WIDTH, HEIGHT + 35))
|
|
self.assertTrue(torch.equal(frames[:, 17:17 + HEIGHT], original))
|
|
self.assertTrue(torch.all(mask[:, 17:17 + HEIGHT] == 0))
|
|
self.assertTrue(torch.all(mask[:, :17] == 1))
|
|
self.assertTrue(torch.all(mask[:, -18:] == 1))
|
|
|
|
def test_limit_shortens_audio_and_duration_together(self):
|
|
_, _, audio, count, fps, _, _, duration = self.run_node(
|
|
start_seconds=0.5, end_seconds=1.5, every_nth=2, max_frames=3,
|
|
)[:8]
|
|
self.assertEqual(count, 3)
|
|
self.assertAlmostEqual(duration, count / fps)
|
|
self.assertAlmostEqual(audio["waveform"].shape[-1] / audio["sample_rate"], duration, delta=0.01)
|
|
|
|
def test_frame_snap_keeps_a_video_model_count_and_the_audio_follows(self):
|
|
free = self.run_node(start_seconds=0.0, end_seconds=0.0)
|
|
self.assertEqual(free[3], FRAMES)
|
|
snapped = self.run_node(start_seconds=0.0, end_seconds=0.0, frame_snap="4n+1")
|
|
frames, mask, audio, count, fps, _, _, duration, stitcher, original = snapped
|
|
self.assertEqual(original.shape[0], count)
|
|
self.assertEqual(tuple(original.shape[1:3]), (HEIGHT, WIDTH))
|
|
self.assertEqual(count, ((FRAMES - 1) // 4) * 4 + 1)
|
|
self.assertEqual(frames.shape[0], count)
|
|
self.assertEqual(mask.shape[0], count)
|
|
self.assertEqual(stitcher["canvas"].shape[0], count)
|
|
self.assertAlmostEqual(duration, count / fps)
|
|
self.assertAlmostEqual(audio["waveform"].shape[-1] / audio["sample_rate"], duration, delta=0.01)
|
|
|
|
def test_frame_snap_rules(self):
|
|
self.assertEqual(snap_frame_count(78, "8n+1"), 73)
|
|
self.assertEqual(snap_frame_count(97, "8n+1"), 97)
|
|
self.assertEqual(snap_frame_count(78, "4n+1"), 77)
|
|
self.assertEqual(snap_frame_count(78, "free"), 78)
|
|
self.assertEqual(snap_frame_count(5, "8n+1"), 1)
|
|
self.assertEqual(snap_frame_count(0, "8n+1"), 0)
|
|
self.assertEqual(snap_frame_count(30, "nonsense"), 30)
|
|
|
|
def test_frame_snap_changes_the_fingerprint(self):
|
|
common = dict(video="", source_mode="local path", local_path=self.server_file, start_seconds=0.0, end_seconds=0.0)
|
|
values = transform_defaults()
|
|
one = AusBossVideoCropRotatePadClip.IS_CHANGED(**common, **values)
|
|
two = AusBossVideoCropRotatePadClip.IS_CHANGED(**common, frame_snap="8n+1", **values)
|
|
self.assertNotEqual(one, two)
|
|
|
|
def test_stitcher_pastes_the_source_back_over_a_generated_clip(self):
|
|
out = self.run_node(pad_left=16, pad_right=8, feather=0, stitch_blend=4)
|
|
frames, mask, stitcher = out[0], out[1], out[8]
|
|
self.assertEqual(stitcher["kind"], STITCHER_KIND)
|
|
self.assertTrue(torch.equal(stitcher["canvas"], frames))
|
|
self.assertEqual(tuple(stitcher["blend"].shape), tuple(mask.shape))
|
|
self.assertEqual(stitcher["source_bbox"], (16, 0, 16 + WIDTH, HEIGHT))
|
|
# The paste covers the bands (to within the blur kernel's float sum)
|
|
# and ramps a few pixels into the source; beyond the ramp it is zero.
|
|
self.assertTrue(torch.all(stitcher["blend"][:, :, :16] >= 0.999))
|
|
self.assertTrue(torch.all(stitcher["blend"][:, :, 16 + 12 : -(8 + 12)] == 0.0))
|
|
generated = torch.rand_like(frames)
|
|
stitched = apply_stitch(stitcher, generated)
|
|
self.assertEqual(tuple(stitched.shape), tuple(frames.shape))
|
|
self.assertTrue(torch.equal(stitched[:, :, 32:-24], frames[:, :, 32:-24]))
|
|
self.assertTrue(torch.allclose(stitched[:, :, :16], generated[:, :, :16], atol=1e-5))
|
|
self.assertTrue(torch.allclose(stitched[:, :, -8:], generated[:, :, -8:], atol=1e-5))
|
|
|
|
def test_stitcher_without_padding_returns_the_frames_untouched(self):
|
|
out = self.run_node(feather=0)
|
|
frames, stitcher = out[0], out[8]
|
|
self.assertEqual(float(stitcher["blend"].max()), 0.0)
|
|
self.assertTrue(torch.equal(apply_stitch(stitcher, torch.rand_like(frames)), frames))
|
|
|
|
def test_a_shorter_generated_clip_stitches_with_a_matching_mask(self):
|
|
# frame_snap "free" keeps every frame, but LTX hands back 8n+1: the
|
|
# stitch keeps the leading frames and blend_mask has to match them.
|
|
out = self.run_node(pad_left=16, feather=0, stitch_blend=4)
|
|
frames, stitcher = out[0], out[8]
|
|
self.assertEqual(stitcher["blend"].shape[0], FRAMES) # one mask per frame
|
|
kept = snap_frame_count(FRAMES, "8n+1")
|
|
self.assertLess(kept, FRAMES)
|
|
generated = torch.rand_like(frames[:kept])
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
stitched, blend_mask = AusBossStitchInpaint().stitch(stitcher, generated)
|
|
self.assertEqual(tuple(stitched.shape), (kept, *frames.shape[1:]))
|
|
self.assertTrue(torch.equal(blend_mask, stitcher["blend"][:kept]))
|
|
self.assertTrue(torch.equal(stitched[:, :, 32:], frames[:kept, :, 32:]))
|
|
|
|
def test_stitch_grow_moves_the_paste_boundary(self):
|
|
base = self.run_node(pad_left=16, feather=0, stitch_blend=0)[8]["blend"]
|
|
grown = self.run_node(pad_left=16, feather=0, stitch_blend=0, stitch_grow=6)[8]["blend"]
|
|
shrunk = self.run_node(pad_left=16, feather=0, stitch_blend=0, stitch_grow=-6)[8]["blend"]
|
|
self.assertTrue(torch.all(base[:, :, :16] == 1.0) and torch.all(base[:, :, 16:] == 0.0))
|
|
self.assertTrue(torch.all(grown[:, :, :22] == 1.0) and torch.all(grown[:, :, 22:] == 0.0))
|
|
self.assertTrue(torch.all(shrunk[:, :, :10] == 1.0) and torch.all(shrunk[:, :, 10:] == 0.0))
|
|
|
|
def test_stitch_settings_change_the_fingerprint(self):
|
|
common = dict(video="", source_mode="local path", local_path=self.server_file, start_seconds=0.0, end_seconds=0.0)
|
|
values = transform_defaults(pad_left=8)
|
|
one = AusBossVideoCropRotatePadClip.IS_CHANGED(**common, **values)
|
|
two = AusBossVideoCropRotatePadClip.IS_CHANGED(**common, stitch_blend=48, **values)
|
|
self.assertNotEqual(one, two)
|
|
|
|
def test_resize_keeps_frames_and_masks_aligned(self):
|
|
if not ensure_core_video_api(self):
|
|
self.skipTest("Set AUSBOSS_COMFY_ROOT to test ComfyUI's resize implementation.")
|
|
out = self.run_node(
|
|
max_frames=2, pad_left=8, resize_to_megapixels=True,
|
|
megapixels=0.01, resolution_steps=8,
|
|
)
|
|
frames, mask, _, count, _, width, height, _ = out[:8]
|
|
self.assertEqual(count, 2)
|
|
self.assertEqual(tuple(frames.shape[1:3]), (height, width))
|
|
self.assertEqual(tuple(mask.shape), (count, height, width))
|
|
self.assertEqual(width % 8, 0)
|
|
self.assertEqual(height % 8, 0)
|
|
stitcher = out[8]
|
|
self.assertEqual(tuple(stitcher["canvas"].shape), tuple(frames.shape))
|
|
self.assertEqual(tuple(stitcher["blend"].shape), (count, height, width))
|
|
x0, _, x1, _ = stitcher["source_bbox"]
|
|
self.assertTrue(0 < x0 < x1 <= width)
|
|
|
|
def test_resize_leaves_a_clean_step_between_the_clip_and_its_band(self):
|
|
if not ensure_core_video_api(self):
|
|
self.skipTest("Set AUSBOSS_COMFY_ROOT to test ComfyUI's resize implementation.")
|
|
# The frames are resized by themselves and padded after. Resizing the
|
|
# padded canvas as one picture left a column part fill, part mask.
|
|
out = self.run_node(
|
|
max_frames=2, pad_left=8, feather=0, resize_to_megapixels=True,
|
|
megapixels=0.01, resolution_steps=8,
|
|
)
|
|
frames, mask, stitcher = out[0], out[1], out[8]
|
|
x0 = stitcher["source_bbox"][0]
|
|
self.assertGreater(x0, 0)
|
|
self.assertTrue(torch.all(frames[:, :, :x0] == 128 / 255))
|
|
self.assertTrue(torch.all(mask[:, :, :x0] == 1.0))
|
|
self.assertTrue(torch.all(mask[:, :, x0:] == 0.0))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|