241 lines
10 KiB
Python
241 lines
10 KiB
Python
from torchvision.models.optical_flow import Raft_Large_Weights, Raft_Small_Weights
|
|
from torchvision.models.optical_flow import raft_large, raft_small
|
|
import torch
|
|
from .flow_utils import apply_warp, get_flow_and_mask, mix_cc
|
|
|
|
raft_weights = Raft_Large_Weights.C_T_SKHT_V1
|
|
raft_device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
|
|
class ExtractOpticalFlow:
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
return {"required":
|
|
{
|
|
"current_frame": ("IMAGE",),
|
|
"previous_frame": ("IMAGE",),
|
|
"num_flow_updates": ("INT", {"default": 20, "min": 5, "max": 100})
|
|
}
|
|
}
|
|
|
|
CATEGORY = "WarpFusion"
|
|
RETURN_TYPES = ("BACKWARD_FLOW", "MASK", "MASK", "MASK", "IMAGE")
|
|
RETURN_NAMES = ("Flow", "Motion edge mask", "Occlusion mask", "Border mask", "Flow preview")
|
|
FUNCTION = "get_flow"
|
|
|
|
raft_model = raft_large(weights=raft_weights, progress=False).to(raft_device).half()
|
|
|
|
def get_flow(self, current_frame, previous_frame, num_flow_updates):
|
|
flow, flow_imgs, edge_mask, occlusion_mask, border_mask = get_flow_and_mask(previous_frame, current_frame, num_flow_updates=num_flow_updates, raft_model=self.raft_model, edge_width=11, dilation=2)
|
|
return (flow, edge_mask, occlusion_mask, border_mask, flow_imgs, )
|
|
|
|
|
|
class WarpFrame:
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
return {"required":
|
|
{
|
|
"previous_frame": ("IMAGE",),
|
|
"flow": ("BACKWARD_FLOW",),
|
|
}
|
|
}
|
|
|
|
CATEGORY = "WarpFusion"
|
|
RETURN_TYPES = ("IMAGE", )
|
|
FUNCTION = "warp"
|
|
|
|
def warp(self, previous_frame, flow):
|
|
warped_frame = apply_warp(previous_frame, flow, padding=0.2)
|
|
|
|
return (warped_frame, )
|
|
|
|
class MixConsistencyMaps:
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
return {"required":
|
|
{
|
|
"missed_consistency": ("MASK",),
|
|
"overshoot_consistency": ("MASK",),
|
|
"edge_consistency": ("MASK",),
|
|
"blur": ("INT", {"default": 1, "min": 0, "max": 100}),
|
|
"dilate":("INT", {"default": 2, "min": 0, "max": 100}),
|
|
"force_binary":("BOOLEAN", {"default": True}),
|
|
"missed_consistency_weight":("FLOAT", {"default": 1.0, "min": 0.0, "max": 1, "step": 0.01}),
|
|
"overshoot_consistency_weight":("FLOAT", {"default": 1.0, "min": 0.0, "max": 1, "step": 0.01}),
|
|
"edges_consistency_weight":("FLOAT", {"default": 1.0, "min": 0.0, "max": 1, "step": 0.01}),
|
|
|
|
}
|
|
}
|
|
|
|
CATEGORY = "WarpFusion"
|
|
RETURN_TYPES = ("MASK", )
|
|
FUNCTION = "get_mixed_cc"
|
|
|
|
def get_mixed_cc(self, missed_consistency, overshoot_consistency, edge_consistency, blur, dilate,
|
|
force_binary, missed_consistency_weight, overshoot_consistency_weight, edges_consistency_weight):
|
|
mixed = mix_cc(missed_consistency, overshoot_consistency, edge_consistency, blur=blur, dilate=dilate, missed_consistency_weight=missed_consistency_weight,
|
|
overshoot_consistency_weight=overshoot_consistency_weight, edges_consistency_weight=edges_consistency_weight, force_binary=force_binary)
|
|
|
|
return (mixed, )
|
|
|
|
|
|
class ExtractFlowAndMixConsistencyMaps:
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
return {"required":
|
|
{
|
|
"current_frame": ("IMAGE",),
|
|
"previous_frame": ("IMAGE",),
|
|
"num_flow_updates": ("INT", {"default": 20, "min": 5, "max": 100}),
|
|
"blur": ("INT", {"default": 1, "min": 0, "max": 100}),
|
|
"dilate":("INT", {"default": 2, "min": 0, "max": 100}),
|
|
"force_binary":("BOOLEAN", {"default": True}),
|
|
"missed_consistency_weight":("FLOAT", {"default": 1.0, "min": 0.0, "max": 1, "step": 0.01}),
|
|
"overshoot_consistency_weight":("FLOAT", {"default": 1.0, "min": 0.0, "max": 1, "step": 0.01}),
|
|
"edges_consistency_weight":("FLOAT", {"default": 1.0, "min": 0.0, "max": 1, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
CATEGORY = "WarpFusion"
|
|
RETURN_TYPES = ("BACKWARD_FLOW", "MASK", "MASK", "MASK", "IMAGE", "MASK", "IMAGE", "IMAGE")
|
|
RETURN_NAMES = ("Flow", "Motion edge mask", "Occlusion mask", "Border mask", "Flow preview", "Mixed consistency map", "Current frame", "Previous frame")
|
|
FUNCTION = "get_flow_and_mixed_cc"
|
|
|
|
raft_model = raft_large(weights=raft_weights, progress=False).to(raft_device).half()
|
|
|
|
def get_flow_and_mixed_cc(self, current_frame, previous_frame, num_flow_updates, blur, dilate, force_binary, missed_consistency_weight, overshoot_consistency_weight, edges_consistency_weight):
|
|
flow, flow_imgs, edge_mask, occlusion_mask, border_mask = get_flow_and_mask(previous_frame, current_frame, num_flow_updates=num_flow_updates, raft_model=self.raft_model, edge_width=11, dilation=2)
|
|
mixed = mix_cc(missed_cc=occlusion_mask, overshoot_cc=border_mask, edge_cc=edge_mask, blur=blur, dilate=dilate, missed_consistency_weight=missed_consistency_weight,
|
|
overshoot_consistency_weight=overshoot_consistency_weight, edges_consistency_weight=edges_consistency_weight, force_binary=force_binary)
|
|
|
|
return (flow, edge_mask, occlusion_mask, border_mask, flow_imgs, mixed, current_frame, previous_frame)
|
|
|
|
class KeyframedFlowApplication:
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
return {"required": {
|
|
"motion_source_frames": ("IMAGE",),
|
|
"frames_to_warp": ("IMAGE",),
|
|
"keyframe_weights": ("STRING", {
|
|
"multiline": True,
|
|
"default": '{"0": 1.0, "10": 2.0}'
|
|
}),
|
|
"keyframe_repeats": ("STRING", {
|
|
"multiline": True,
|
|
"default": '{"0": 1, "10": 3}'
|
|
}),
|
|
"num_flow_updates": ("INT", {"default": 20, "min": 5, "max": 100})
|
|
}}
|
|
|
|
CATEGORY = "WarpFusion"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("Processed Frame",)
|
|
FUNCTION = "process_frames"
|
|
|
|
raft_model = raft_large(weights=Raft_Large_Weights.C_T_SKHT_V1, progress=False).to("cuda" if torch.cuda.is_available() else "cpu").half()
|
|
|
|
def process_frames(self, motion_source_frames, frames_to_warp, keyframe_weights, keyframe_repeats, num_flow_updates):
|
|
# Parse the keyframe dictionaries
|
|
print('keyframe_weights', keyframe_weights)
|
|
print('keyframe_repeats', keyframe_repeats)
|
|
weights = eval(keyframe_weights)
|
|
repeats = eval(keyframe_repeats)
|
|
|
|
if type(weights) == list:
|
|
weights = {str(i): v for i, v in enumerate(weights) if v > 1}
|
|
if type(repeats) == list:
|
|
repeats = {str(i): v for i, v in enumerate(repeats) if v > 1}
|
|
|
|
weights = {str(k): v for k, v in weights.items()}
|
|
repeats = {str(k): int(v) for k, v in repeats.items()}
|
|
|
|
print('parsed weights', weights)
|
|
print('parsed repeats', repeats)
|
|
|
|
# Convert frames to list if they're not already
|
|
processed_frames = []
|
|
num_frames = len(motion_source_frames)
|
|
|
|
flow_dict = {}
|
|
|
|
# Sort keyframes to process them in order
|
|
|
|
weights = [weights.get(str(frame_number), 1.0) for frame_number in range(num_frames-1)]
|
|
repeats = [repeats.get(str(frame_number), 1) for frame_number in range(num_frames-1)]
|
|
|
|
flow_map = {}
|
|
frame_number = 0
|
|
while frame_number < num_frames-1:
|
|
repeat_count = repeats[frame_number]
|
|
if repeat_count <= 1:
|
|
frame_number += 1
|
|
continue
|
|
if repeat_count > 1:
|
|
for i in range(repeat_count):
|
|
flow_map[frame_number + i] = frame_number
|
|
frame_number += repeat_count - 1
|
|
|
|
# print('flow_map', flow_map)
|
|
# print('weights', weights)
|
|
# print('repeats', repeats)
|
|
|
|
from tqdm import trange
|
|
for frame_number in trange(num_frames-1):
|
|
# Find the active keyframe
|
|
weight = weights[frame_number]
|
|
if (frame_number == 0) or (frame_number == num_frames-1) or ((frame_number not in flow_map) and(weight <= 1.0)):
|
|
processed_frames.append(frames_to_warp[frame_number:frame_number+1])
|
|
continue
|
|
|
|
# print('\napplying flow for frame', frame_number)
|
|
|
|
# Extract flow between source frames
|
|
flow_frame = flow_map.get(frame_number, frame_number)
|
|
if flow_frame != frame_number:
|
|
print('applying flow for frame', frame_number, 'from frame', flow_frame)
|
|
|
|
if flow_frame not in flow_dict:
|
|
flow, _, _, _, _ = get_flow_and_mask(
|
|
motion_source_frames[flow_frame:flow_frame+1],
|
|
motion_source_frames[flow_frame+1:flow_frame+2],
|
|
num_flow_updates=num_flow_updates,
|
|
raft_model=self.raft_model
|
|
)
|
|
flow_dict[frame_number] = flow
|
|
else:
|
|
flow = flow_dict[flow_frame]
|
|
# print('flow', flow.max(), flow.min())
|
|
|
|
if flow_frame != frame_number:
|
|
warped_frame = processed_frames[-1]
|
|
else:
|
|
warped_frame = frames_to_warp[frame_number:frame_number+1]
|
|
|
|
warped_frame = apply_warp(
|
|
warped_frame,
|
|
flow * weight,
|
|
padding=0.2
|
|
)
|
|
processed_frames.append(warped_frame)
|
|
|
|
# Concatenate all processed frames
|
|
output_frames = torch.cat(processed_frames, dim=0)
|
|
return (output_frames,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ExtractOpticalFlow": ExtractOpticalFlow,
|
|
"WarpFrame":WarpFrame,
|
|
"MixConsistencyMaps":MixConsistencyMaps,
|
|
"ExtractFlowAndMixConsistencyMaps":ExtractFlowAndMixConsistencyMaps,
|
|
"KeyframedFlowApplication": KeyframedFlowApplication
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ExtractOpticalFlow": "ExtractOpticalFlow",
|
|
"WarpFrame":"WarpFrame",
|
|
"MixConsistencyMaps":"MixConsistencyMaps",
|
|
"ExtractFlowAndMixConsistencyMaps":"ExtractFlowAndMixConsistencyMaps",
|
|
"KeyframedFlowApplication": "KeyframedFlowApplication"
|
|
}
|
|
|