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.
637 lines
23 KiB
Python
637 lines
23 KiB
Python
### info for code completion AI ###
|
|
"""
|
|
all of these classes are plugins for comfyui and follow the same pattern
|
|
all of the images are torch tensors and it is unknown and unimportant if they are on the cpu or gpu
|
|
all image inputs are (B,W,H,C)
|
|
|
|
avoid numpy and PIL as much as possible
|
|
"""
|
|
import os
|
|
from abc import ABC, ABCMeta, abstractmethod
|
|
from typing import List, Tuple, Union
|
|
import cv2
|
|
import torch
|
|
from cv2 import VideoWriter, VideoWriter_fourcc, VideoCapture
|
|
from torch import dtype
|
|
from custom_nodes.EternalKernelLiteGraphNodes.image import torch_image_show
|
|
|
|
NODE_CLASS_MAPPINGS = {} # this is the dictionary that will be used to register the nodes
|
|
|
|
# setup directories
|
|
root_p = os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
|
|
# input_dir = os.path.join(root_p, 'input', 'video')
|
|
# output_dir = os.path.join(root_p, 'output', 'video')
|
|
|
|
video_dir = os.path.join(root_p, 'video')
|
|
|
|
if not os.path.exists(video_dir):
|
|
os.makedirs(video_dir)
|
|
|
|
|
|
# video_folders = lambda s: sorted(
|
|
# [name for name in os.listdir(s.input_dir) if os.path.isdir(os.path.join(s.input_dir, name))])
|
|
# turn the above in the a full function
|
|
def video_folders(s):
|
|
ret = sorted(
|
|
[name for name in os.listdir(s.input_dir) if os.path.isdir(os.path.join(s.input_dir, name))])
|
|
return ret
|
|
|
|
|
|
def video_files(s):
|
|
# any files in in_video_folders ending in .mp4
|
|
ret = sorted([name for name in os.listdir(s.input_dir) if
|
|
os.path.isfile(os.path.join(s.input_dir, name)) and name.endswith('.mp4')])
|
|
return ret
|
|
|
|
|
|
# start with an ABC to define the common widget interface
|
|
|
|
class WidgetMetaclass(ABCMeta):
|
|
"""A metaclass that automatically registers classes."""
|
|
|
|
def __init__(cls, name, bases, attrs):
|
|
if (ABC in bases) or (VideoWidget in bases):
|
|
pass
|
|
else:
|
|
NODE_CLASS_MAPPINGS[name] = cls
|
|
|
|
super().__init__(name, bases, attrs)
|
|
|
|
|
|
class VideoWidget(ABC, metaclass=WidgetMetaclass):
|
|
"""Abstract base class for simple construction"""
|
|
|
|
input_dir = video_dir
|
|
output_dir = video_dir
|
|
|
|
CATEGORY = "video"
|
|
RETURN_TYPES = ("VIDEO",)
|
|
FUNCTION = "handler"
|
|
|
|
@classmethod
|
|
@abstractmethod
|
|
def INPUT_TYPES(cls):
|
|
"""All subclasses must provide a INPUT_TYPES classmethod"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def handler(self, one: object, two: object):
|
|
"""All subclasses must provide a handler method"""
|
|
pass
|
|
|
|
|
|
class VideoFileToVideoFolder(VideoWidget, metaclass=WidgetMetaclass):
|
|
"""Abstract base class for simple construction"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required":
|
|
{
|
|
"video_in": (video_files(cls),),
|
|
"folder_out": (video_folders(cls),),
|
|
},
|
|
|
|
"optional":
|
|
{"text":
|
|
("STRING", {"multiline": False}),
|
|
}
|
|
}
|
|
|
|
@abstractmethod
|
|
def handler(self, one, two, text):
|
|
return (one, two,)
|
|
|
|
|
|
class VideoFolderToImage(VideoWidget, metaclass=WidgetMetaclass):
|
|
"""Abstract base class for simple construction"""
|
|
RETURN_TYPES = ("IMAGE",)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required":
|
|
{
|
|
"folder_in": (video_folders(cls),),
|
|
"idx_slice_start": ("INT", {"min": 0, "max": 10000, "step": 1, "default": 0}),
|
|
"idx_slice_stop": ("INT", {"min": 0, "max": 10000, "step": 1, "default": 1}),
|
|
"slice_idx_step": ("INT", {"min": 0, "max": 1000, "step": 1, "default": 1}),
|
|
"idx": ("INT", {"min": 0, "max": 10000, "step": 1, "default": 0}),
|
|
},
|
|
|
|
"optional":
|
|
{"text":
|
|
("STRING", {"multiline": False}),
|
|
}
|
|
}
|
|
|
|
@abstractmethod
|
|
def handler(self, one, two, text):
|
|
return (one,)
|
|
|
|
|
|
class VideoFileToImage(VideoWidget, metaclass=WidgetMetaclass):
|
|
"""Abstract base class for simple construction"""
|
|
RETURN_TYPES = ("IMAGE",)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required":
|
|
{
|
|
"video_in": (video_files(cls),),
|
|
"idx_slice_start": ("INT", {"min": 0, "max": 10000, "step": 1, "default": 0}),
|
|
"idx_slice_stop": ("INT", {"min": 0, "max": 10000, "step": 1, "default": 1}),
|
|
"slice_idx_step": ("INT", {"min": 0, "max": 1000, "step": 1, "default": 1}),
|
|
"idx": ("INT", {"min": 0, "max": 10000, "step": 1, "default": 0}),
|
|
},
|
|
|
|
"optional":
|
|
{"text":
|
|
("STRING", {"multiline": False}),
|
|
}
|
|
}
|
|
|
|
@abstractmethod
|
|
def handler(self, one, two, text):
|
|
return (one,)
|
|
|
|
|
|
## end of abstract base classes ##
|
|
|
|
class VideoToFrames(VideoFileToVideoFolder, metaclass=WidgetMetaclass):
|
|
"""use cv2 to open the video and use it as a source for image stacks"""
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.old = self.INPUT_TYPES
|
|
# self.INPUT_TYPES = lambda s: self.old().update({})
|
|
|
|
def handler(self, video_in, folder_out, text):
|
|
video_in_path = os.path.join(self.input_dir, video_in)
|
|
video_out_path = os.path.join(self.output_dir, folder_out)
|
|
|
|
vidcap = cv2.VideoCapture(video_in_path)
|
|
success, image = vidcap.read()
|
|
count = 0
|
|
while success:
|
|
cv2.imwrite(os.path.join(video_out_path, "frame%d.jpg" % count), image)
|
|
success, image = vidcap.read()
|
|
count += 1
|
|
|
|
return (folder_out,)
|
|
|
|
|
|
class VideoFramesToImageStack(VideoFolderToImage, metaclass=WidgetMetaclass):
|
|
"""Use cv2 to open the video and save the videos individual frames"""
|
|
|
|
def handler(self, folder_in, text, idx_slice_start, idx_slice_stop, slice_idx_step, idx):
|
|
"""
|
|
Return slices of time from the video, sequences of images, defined by the start, stop and step
|
|
- torch tensors (T,H,W,C)
|
|
"""
|
|
import glob
|
|
|
|
folder_in_path = os.path.join(self.input_dir, folder_in)
|
|
image_list = []
|
|
|
|
# List and sort all PNG or JPG files in the folder
|
|
file_fps = sorted(
|
|
glob.glob(os.path.join(folder_in_path, '*.png')) + glob.glob(os.path.join(folder_in_path, '*.jpg')))
|
|
|
|
# Ensure the slice doesn't go beyond the end of the file list
|
|
if idx_slice_stop + idx > len(file_fps):
|
|
idx_slice_stop = len(file_fps) - idx
|
|
|
|
for i in range(idx_slice_start + idx, idx_slice_stop + idx, slice_idx_step):
|
|
for fn in file_fps[i:i + slice_idx_step]:
|
|
image = cv2.imread(fn)
|
|
if image is not None:
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # converting to RGB format
|
|
image_list.append(image)
|
|
|
|
# Convert list of images to tensor (T,H,W,C)
|
|
tensor = torch.stack([torch.tensor(img, dtype=torch.float32) for img in image_list])
|
|
tensor = tensor / 255.0
|
|
return (tensor,)
|
|
|
|
|
|
class VideoFileToImageStack(VideoFileToImage, metaclass=WidgetMetaclass):
|
|
"""Use cv2 to open a video and save the videos individual frames"""
|
|
|
|
def handler(self, video_in, text, idx_slice_start, idx_slice_stop, slice_idx_step, idx):
|
|
"""
|
|
Return slices of time from the video, sequences of images, defined by the start, stop and step
|
|
- torch tensors (T,H,W,C)
|
|
"""
|
|
image_list = []
|
|
video_in_fp = os.path.join(self.input_dir, video_in)
|
|
vidcap = cv2.VideoCapture(video_in_fp)
|
|
if not vidcap.isOpened():
|
|
print("Error: Unable to open video file")
|
|
return
|
|
|
|
start = (idx * slice_idx_step) + idx_slice_start
|
|
stop = (idx * slice_idx_step) + idx_slice_stop
|
|
|
|
for frame_count in range(start, stop, slice_idx_step):
|
|
# Seek to frame_count frame
|
|
vidcap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)
|
|
|
|
success, image = vidcap.read()
|
|
if not success:
|
|
print("Error: Unable to read frame from video file")
|
|
break
|
|
# image_list.append(image)
|
|
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
image_list.append(image)
|
|
|
|
# Convert list of images to tensor (T,H,W,C)
|
|
tensor = torch.stack([torch.tensor(img, dtype=torch.float32) for img in image_list])
|
|
tensor = tensor / 255.0
|
|
|
|
return (tensor,)
|
|
|
|
|
|
class ImageStackToVideoFrames(metaclass=WidgetMetaclass):
|
|
"""Really just saves the image stack to a folder"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
return {"required":
|
|
{
|
|
"image_stack": ("IMAGE",),
|
|
"folder_out": ("STRING", {"multiline": False, "default": "video01"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
CATEGORY = "video"
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "handler"
|
|
|
|
def handler(self, **kwargs):
|
|
image_stack = kwargs["image_stack"]
|
|
folder_out = kwargs["folder_out"]
|
|
folder_out_path = video_dir
|
|
|
|
# full path to the output folder
|
|
folder_out_path = os.path.join(video_dir, folder_out)
|
|
|
|
# Create the output folder if it doesn't exist
|
|
if not os.path.exists(folder_out_path):
|
|
os.makedirs(folder_out_path)
|
|
# move any current files in the output folder to a random folder inside the video folder
|
|
import uuid
|
|
import shutil
|
|
random_folder = os.path.join(folder_out_path, str(uuid.uuid4()))
|
|
os.makedirs(random_folder)
|
|
for file in os.listdir(folder_out_path):
|
|
file_path = os.path.join(folder_out_path, file)
|
|
if os.path.isfile(file_path):
|
|
shutil.move(file_path, random_folder)
|
|
|
|
# Save each image in the image stack to the output folder
|
|
for i in range(image_stack.shape[0]):
|
|
image = image_stack[i].numpy()
|
|
# make sure the image is in the correct color format
|
|
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
|
# convert to 0-255
|
|
image = image * 255
|
|
# use zfill to get a good ordering of the frames
|
|
file_name = os.path.join(folder_out_path, "frame" + str(i).zfill(6) + ".jpg")
|
|
# use the best compression nearly lossless
|
|
cv2.imwrite(file_name, image, [cv2.IMWRITE_JPEG_QUALITY, 99])
|
|
|
|
image_stack = None
|
|
del image_stack
|
|
return ()
|
|
|
|
|
|
class ImageStackToVideoFile(metaclass=WidgetMetaclass):
|
|
"""Really just saves the image stack to a video file"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
return {"required":
|
|
{
|
|
"image_stack": ("IMAGE",),
|
|
"video_out": ("STRING", {"multiline": False, "default": "video01.mp4"}),
|
|
},
|
|
"optional":
|
|
{"fps": ("FLOAT", {"default": 30, "min": 1, "max": 120}),
|
|
"func": ("FUNC", {"default": "NONE"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("FUNC", "STRING",)
|
|
RETURN_NAMES = ("FUNC", "video_out_path",)
|
|
CATEGORY = "video"
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "handler"
|
|
|
|
def handler(self, **kwargs):
|
|
"""frames are in the format (T,H,W,C)"""
|
|
import cv2
|
|
import os
|
|
import numpy as np
|
|
import uuid
|
|
import shutil
|
|
import inspect
|
|
|
|
func = kwargs.get("func", None)
|
|
|
|
if func:
|
|
kwargs = func(**kwargs)
|
|
|
|
self.ARGS = kwargs
|
|
self.IN_FUNC = func
|
|
# get this functions source code
|
|
try:
|
|
self.CODE = inspect.getsource(ImageStackToVideoFile.handler)
|
|
except OSError:
|
|
self.CODE = "Unable to get source code"
|
|
|
|
|
|
self.FUNC = self.handler
|
|
|
|
try:
|
|
image_stack = kwargs["image_stack"]
|
|
except KeyError:
|
|
# this means we are just passing our function through
|
|
return (self,)
|
|
|
|
video_out = kwargs["video_out"]
|
|
fps = kwargs["fps"]
|
|
video_out_path = os.path.join(video_dir, video_out)
|
|
|
|
# check to see if the video file already exists
|
|
if os.path.exists(video_out_path):
|
|
# move the existing file to a random folder inside the video folder
|
|
random_folder = os.path.join(video_dir, str(uuid.uuid4()))
|
|
os.makedirs(random_folder)
|
|
shutil.move(video_out_path, random_folder)
|
|
|
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
|
video = cv2.VideoWriter(video_out_path + ".mp4", fourcc, float(fps),
|
|
(image_stack.shape[2], image_stack.shape[1]))
|
|
|
|
for i in range(image_stack.shape[0]):
|
|
image = image_stack[i].numpy()
|
|
|
|
# Check that the image is in the correct format
|
|
if image.dtype != np.uint8:
|
|
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
|
image = (image * 255).astype(np.uint8) # convert to 0-255 and cast to uint8
|
|
|
|
video.write(image)
|
|
|
|
video.release()
|
|
image_stack = None
|
|
del image_stack
|
|
return (self, video_out_path + ".mp4",)
|
|
|
|
|
|
class TemporalSpatialSmoothing(metaclass=WidgetMetaclass):
|
|
"""Smooths the image stack in time and space"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
return {"required":
|
|
{
|
|
"image_stack": ("IMAGE",),
|
|
"spatial_kernel_size": ("INT", {"default": 3, "min": 1, "max": 50}),
|
|
"temporal_kernel_size": ("INT", {"default": 3, "min": 1, "max": 50}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "video"
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "handler"
|
|
|
|
@staticmethod
|
|
def gaussian_kernel_1d(kernel_size=3, sigma=1.0):
|
|
import numpy as np
|
|
if kernel_size % 2 == 0:
|
|
kernel_size += 1
|
|
x_coord = torch.arange(-kernel_size // 2 + 1., kernel_size // 2 + 1.).float()
|
|
gaussian_kernel = (1. / (sigma * torch.sqrt(torch.tensor(2. * np.pi)))) * torch.exp(
|
|
-x_coord ** 2. / (2 * sigma ** 2))
|
|
gaussian_kernel = gaussian_kernel / torch.sum(gaussian_kernel)
|
|
return gaussian_kernel
|
|
|
|
def handler(self, **kwargs):
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
image_stack = kwargs["image_stack"]
|
|
spatial_kernel_size = kwargs["spatial_kernel_size"]
|
|
temporal_kernel_size = kwargs["temporal_kernel_size"]
|
|
spatial_sigma = kwargs.get("spatial_sigma", 1.0)
|
|
temporal_sigma = kwargs.get("temporal_sigma", 1.0)
|
|
|
|
# make sure image_stack is a torch tensor and has the right dimensions
|
|
if not torch.is_tensor(image_stack):
|
|
image_stack = torch.tensor(image_stack, dtype=torch.float32)
|
|
|
|
if len(image_stack.shape) < 5: # if it doesn't have a channel dimension
|
|
image_stack = image_stack.unsqueeze(0) # add a batch dimension
|
|
image_stack = image_stack.permute(0, 4, 1, 2,
|
|
3) # rearrange dimensions to: (batch, channel, depth, height, width)
|
|
|
|
# create 1D Gaussian kernels for each dimension
|
|
spatial_kernel = self.gaussian_kernel_1d(spatial_kernel_size, spatial_sigma)
|
|
temporal_kernel = self.gaussian_kernel_1d(temporal_kernel_size, temporal_sigma)
|
|
|
|
# apply the convolution along each dimension separately
|
|
for dim, kernel in enumerate([temporal_kernel, spatial_kernel, spatial_kernel]):
|
|
padding_size = kernel.shape[0] // 2
|
|
# define the padding differently for each dimension
|
|
if dim == 0: # depth (temporal)
|
|
padding = nn.ReplicationPad3d((0, 0, 0, 0, padding_size, padding_size))
|
|
elif dim == 1: # height
|
|
padding = nn.ReplicationPad3d((0, 0, padding_size, padding_size, 0, 0))
|
|
elif dim == 2: # width
|
|
padding = nn.ReplicationPad3d((padding_size, padding_size, 0, 0, 0, 0))
|
|
image_stack = padding(image_stack)
|
|
if dim == 0: # depth/temporal
|
|
conv = nn.Conv3d(in_channels=3, out_channels=3, kernel_size=(kernel.shape[0], 1, 1), groups=3,
|
|
bias=False)
|
|
reshaped_kernel = kernel[None, None, :, None, None]
|
|
elif dim == 1: # height
|
|
conv = nn.Conv3d(in_channels=3, out_channels=3, kernel_size=(1, kernel.shape[0], 1), groups=3,
|
|
bias=False)
|
|
reshaped_kernel = kernel[None, None, None, :, None]
|
|
elif dim == 2: # width
|
|
conv = nn.Conv3d(in_channels=3, out_channels=3, kernel_size=(1, 1, kernel.shape[0]), groups=3,
|
|
bias=False)
|
|
reshaped_kernel = kernel[None, None, None, None, :]
|
|
conv.weight = nn.Parameter(reshaped_kernel.repeat(3, 1, 1, 1, 1), requires_grad=False)
|
|
image_stack = conv(image_stack)
|
|
|
|
# remove the batch dimension and rearrange dimensions back to: (depth, height, width, channel)
|
|
image_stack = image_stack.squeeze(0).permute(1, 2, 3, 0)
|
|
|
|
return (image_stack,)
|
|
|
|
|
|
import torch
|
|
from torchvision.models.optical_flow import raft_large, Raft_Large_Weights, raft_small, Raft_Small_Weights
|
|
from torchvision.utils import flow_to_image
|
|
from torchvision.transforms.functional import resize
|
|
|
|
|
|
class TemporalOpticalFlowSmoothing(metaclass=WidgetMetaclass):
|
|
"""Smooths the image stack in time using RAFT for optical flow estimation"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required":
|
|
{
|
|
"image_stack": ("IMAGE",),
|
|
"k": ("INT", {"default": 3, "min": 1, "max": 120}),
|
|
"device": ("STRING", ["cpu", "cuda"]),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "video"
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "handler"
|
|
|
|
@staticmethod
|
|
def make_grid(flow):
|
|
# Assumes flow is in shape Bx2xHxW
|
|
B, _, H, W = flow.size()
|
|
grid_y, grid_x = torch.meshgrid(torch.linspace(-1, 1, W), torch.linspace(-1, 1, H))
|
|
grid = torch.stack((grid_y.t(), grid_x.t()), 2).unsqueeze(0)
|
|
grid = grid.repeat(B, 1, 1, 1) # Make it BxHxWx2
|
|
grid = grid.permute(0, 3, 1, 2) # Convert to Bx2xHxW
|
|
grid = grid.type_as(flow) # Make sure the data type is the same
|
|
return grid + flow
|
|
|
|
@staticmethod
|
|
def worker(args):
|
|
import torch.nn.functional as F
|
|
k, image_stack, device = args
|
|
weights = Raft_Large_Weights.DEFAULT
|
|
model = raft_large(weights=weights, progress=False).to(device)
|
|
model = model.eval()
|
|
|
|
img1 = image_stack[0].unsqueeze(0)
|
|
img2 = image_stack[1].unsqueeze(0)
|
|
|
|
# Remove batch dimension if present
|
|
while len(img1.shape) > 4:
|
|
print(f"Batch dimension detected. for image shape {img1.shape}")
|
|
img1 = img1.squeeze(0)
|
|
img2 = img2.squeeze(0)
|
|
|
|
# Permute tensor dimensions to match expected input for model
|
|
img1 = img1.permute(0, 3, 1, 2)
|
|
img2 = img2.permute(0, 3, 1, 2)
|
|
|
|
# Move images to the specified device
|
|
img1, img2 = img1.to(device), img2.to(device)
|
|
|
|
flow = model(img1, img2)[-1]
|
|
|
|
# Create intermediary frames
|
|
frames = []
|
|
for i in range(k):
|
|
factor = i / k
|
|
#factor = factor * 2 - 1
|
|
|
|
#normalize flow
|
|
n_flow = flow.clone()
|
|
print(n_flow.min(), n_flow.max(), n_flow.dtype)
|
|
#n_flow[:,0,:,:] = n_flow[:,0,:,:] / (n_flow[:,0,:,:].max() - n_flow[:,0,:,:].min())
|
|
#n_flow[:,1,:,:] = n_flow[:,1,:,:] / (n_flow[:,1,:,:].max() - n_flow[:,1,:,:].min())
|
|
#print(n_flow.min(), n_flow.max(), n_flow.dtype)
|
|
#n_flow[:, 0, :, :] = n_flow[:, 0, :, :] - n_flow[:, 0, :, :].min()
|
|
#n_flow[:, 1, :, :] = n_flow[:, 1, :, :] - n_flow[:, 1, :, :].min()
|
|
#print (n_flow.min(), n_flow.max(),n_flow.dtype)
|
|
n_flow= n_flow/n_flow.shape[2]
|
|
print(n_flow.min(), n_flow.max(), n_flow.dtype)
|
|
|
|
|
|
flow_interpolated = n_flow * factor
|
|
grid = TemporalOpticalFlowSmoothing.make_grid(flow_interpolated)
|
|
grid = grid.permute(0, 2, 3, 1)
|
|
warped_image = F.grid_sample(img1, grid, mode='bicubic', padding_mode='zeros')
|
|
warped_image = warped_image.permute(0, 2, 3, 1).detach().cpu() # Convert it back to BxHxWxC
|
|
frames.append(warped_image)
|
|
|
|
return frames
|
|
|
|
@staticmethod
|
|
def calculate_optical_flow(image_stack, device,k):
|
|
# from torch.multiprocessing import Pool, set_start_method
|
|
|
|
# set_start_method('spawn') # set the start method to spawn
|
|
# with Pool(2) as pool:
|
|
# # create the tasks in an intellegent way, detaching each stack, with each stack being only the frames needed
|
|
# tasks = []
|
|
# for t in range(image_stack.shape[0] - 1):
|
|
# needed_frames = [t, t + 1] # the frames needed for this task
|
|
# stack = image_stack[needed_frames, :, :, :] # the stack of frames needed for this task
|
|
# stack = stack.clone()
|
|
# tasks.append((2, stack.detach().cpu(), device)) # add the task to the list of tasks
|
|
# flows = pool.map(TemporalOpticalFlowSmoothing.worker, tasks)
|
|
# # print(f"flows: {flows}")
|
|
# # print(f"flows[0]: {flows[0]}")
|
|
|
|
##map(TemporalOpticalFlowSmoothing.worker, tasks)
|
|
|
|
flows = []
|
|
for i in range(image_stack.shape[0] - 1):
|
|
needed_frames = [i, i + 1] # the frames needed for this task
|
|
keys = image_stack[needed_frames, :, :, :]
|
|
flow = TemporalOpticalFlowSmoothing.worker((k, keys, device))
|
|
|
|
flows.append(keys[0].cpu())
|
|
for f in flow:
|
|
flows.append(f[0].cpu())
|
|
#flows.append(keys[1].cpu())
|
|
|
|
|
|
return flows
|
|
|
|
def handler(self, **kwargs):
|
|
from copy import deepcopy
|
|
# image stacks are actually (b/t,w,h,c) so we need to rearrange the dimensions
|
|
k = kwargs["k"]
|
|
image_stack = deepcopy(kwargs["image_stack"])
|
|
device = kwargs.get("device", "cpu")
|
|
|
|
image_stack = image_stack.permute(0, 2, 1, 3)
|
|
|
|
|
|
# make sure image_stack is a torch tensor and has the right dimensions
|
|
if not torch.is_tensor(image_stack):
|
|
image_stack = torch.tensor(image_stack, dtype=torch.float32)
|
|
|
|
if len(image_stack.shape) < 4: # if it doesn't have a channel dimension
|
|
image_stack = image_stack.unsqueeze(0) # add a batch dimension
|
|
|
|
# Move image stack to the specified device
|
|
image_stack = image_stack.to(device)
|
|
|
|
# Calculate optical flow
|
|
flows = self.calculate_optical_flow(image_stack, device,k)
|
|
stack = torch.stack(flows, dim=0)
|
|
# Convert optical flows to images for visualization
|
|
# flow_images = [flow_to_image(flow) for flow in flows]
|
|
#change back to (b or t,w,h,c)
|
|
|
|
stack = stack.permute(0, 2, 1, 3)
|
|
return (stack,)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import doctest
|
|
|
|
doctest.testmod()
|