Tasha Upchurch
e18961e639
- This test validates the functionality of the `TemporalSpatialSmoothing` handler, which smooths image stacks in both temporal and spatial dimensions. - The test uses random frames of noise and verifies the correctness of the handler. - Made minor changes in `image.py` - Updated code formatting for improved readability and consistency. - Added error handling for cases when creating a PIL image from a numpy array. - No significant changes were made in `video.py`. These changes primarily focus on adding test coverage and maintaining code quality.
139 lines
5.0 KiB
Python
139 lines
5.0 KiB
Python
from unittest import TestCase
|
|
import os
|
|
import torch
|
|
os.environ["ETERNAL_KERNEL_LITEGRAPH_NODES_TEST"] = "True"
|
|
|
|
class TestImageStackToVideoFile(TestCase):
|
|
def test_handler(self):
|
|
# test with random frames of noise
|
|
import numpy as np
|
|
from custom_nodes.EternalKernelLiteGraphNodes.video import ImageStackToVideoFile
|
|
|
|
# create a random stack of images
|
|
stack = np.random.randint(0, 255, (100, 100, 100, 3), dtype=np.uint8)
|
|
# it expects a tensor
|
|
stack = torch.from_numpy(stack)
|
|
# create the node
|
|
node = ImageStackToVideoFile()
|
|
# run the node
|
|
node.handler(image_stack=stack, video_out="test.mp4")
|
|
# check that the file exists
|
|
# find the actual filepath for the generated file, will be ../../../video/test.mp4
|
|
test_video_file_fullpath = os.path.join(
|
|
os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))), "video",
|
|
"test.mp4")
|
|
|
|
self.assertTrue(os.path.exists(test_video_file_fullpath))
|
|
# remove the file
|
|
os.remove(test_video_file_fullpath)
|
|
|
|
|
|
class TestTemporalSpatialSmoothing(TestCase):
|
|
def test_handler(self):
|
|
# test with random frames of noise
|
|
|
|
from custom_nodes.EternalKernelLiteGraphNodes.video import TemporalSpatialSmoothing, ImageStackToVideoFile
|
|
from custom_nodes.EternalKernelLiteGraphNodes.image import torch_image_show
|
|
|
|
# use only torch
|
|
# create a random stack of images
|
|
# use floats
|
|
stack = torch.rand((48, 512, 512, 3))
|
|
# _,path= ImageStackToVideoFile().handler(image_stack=stack, video_out="test_start.mp4",fps=24)
|
|
# launch the default video player
|
|
# os.system(f"start {path}")
|
|
|
|
# create the node
|
|
node = TemporalSpatialSmoothing()
|
|
# run the node
|
|
result = node.handler(image_stack=stack, spatial_kernel_size=3, temporal_kernel_size=3, spatial_sigma=.5,
|
|
temporal_sigma=.5)
|
|
# check that the file exists
|
|
_, path = ImageStackToVideoFile().handler(image_stack=result[0], video_out="test_result.mp4", fps=24)
|
|
# launch the default video player
|
|
os.system(f"start {path}")
|
|
|
|
# torch_image_show(result[0])
|
|
|
|
|
|
class TestTemporalOpticalFlowSmoothing(TestCase):
|
|
def test_handler(self):
|
|
from custom_nodes.EternalKernelLiteGraphNodes.video import TemporalOpticalFlowSmoothing, ImageStackToVideoFile
|
|
import time
|
|
import numpy as np
|
|
# create a random stack of images
|
|
#stack = torch.rand((24*4, 512, 512, 3))
|
|
h=256
|
|
w=512
|
|
fps = 24
|
|
seconds = 2
|
|
stack = torch.rand((fps * seconds, h, w, 3))
|
|
# now add a box that moves from left to right with a constant speed
|
|
# the box will be 20% of the width and 20% of the height
|
|
# the box will move from left to right
|
|
|
|
# create the box
|
|
box = torch.zeros((fps * seconds, int(h), int(w), 3))
|
|
# fill the box with ones
|
|
box = box + 1
|
|
# the box width
|
|
bw = min(int(w * .2),int(h*.2))
|
|
# the box height
|
|
bh = int(h * .2)
|
|
# the box x position
|
|
|
|
t = stack.shape[0]
|
|
for i in range(t):
|
|
pt = i / t
|
|
|
|
# the box x position
|
|
bx = int((w - bw) * pt)
|
|
# now add a jitter
|
|
bx += int(np.random.randint(-10, 10, 1))
|
|
# the box y position
|
|
by = int(h * .333)
|
|
# fill the box with ones
|
|
box[i,:,:,:] = 0
|
|
box[i, by:by + bh, bx:bx + bw, :] = 1
|
|
|
|
# create a mask for the box region
|
|
#box_mask = torch.zeros_like(stack)
|
|
#box_mask[:, by:by + bh, bx:bx + bw, :] = 1
|
|
|
|
# set the color of the stack to blue where the box mask is 1
|
|
stack = torch.where(box == 1, torch.tensor([.24, .5, .7], dtype=torch.float32), stack)
|
|
|
|
nm = "box1"
|
|
_, path = ImageStackToVideoFile().handler(image_stack=stack, video_out=nm, fps=fps)
|
|
assert os.path.exists(path), f"Output video file doesn't exist: {path}"
|
|
os.system(f"start {path}")
|
|
time.sleep(1)
|
|
|
|
# create the node
|
|
node = TemporalOpticalFlowSmoothing()
|
|
# run the node
|
|
result = node.handler(image_stack=stack, device="cuda")
|
|
result = result[0]
|
|
|
|
|
|
#result = torch.cat(result[0], dim=0)
|
|
#print(result.shape)
|
|
# now we can permute the dimensions
|
|
#result = result.permute(0, 2,3, 1)
|
|
|
|
|
|
_, path = ImageStackToVideoFile().handler(image_stack=result, video_out="test_result_optical_flow.mp4", fps=fps)
|
|
assert os.path.exists(path), f"Output video file doesn't exist: {path}"
|
|
os.system(f"start {path}")
|
|
|
|
## Check the shape of the output
|
|
#assert len(result) == stack.shape[0] - 1
|
|
|
|
# Each element in the result should have the shape of [1, 2, H, W]
|
|
#for flow in result:
|
|
# assert flow.shape == torch.Size([1, 2, 512, 512])
|
|
|
|
# check the file generation and video play
|
|
|
|
|