68 lines
2.5 KiB
Python
68 lines
2.5 KiB
Python
import torch
|
|
|
|
class VideoFramePreprocessor:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",), # Input is a batch of video frames
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("processed_images",)
|
|
FUNCTION = "process_frames"
|
|
CATEGORY = "Stand-In" # Placed in a sub-category for organization
|
|
|
|
def process_frames(self, images: torch.Tensor):
|
|
# Input tensor shape: (batch_size/frames, height, width, channels)
|
|
if images.dim() != 4:
|
|
raise ValueError("Input must be a batch of images (video frames).")
|
|
|
|
total_frames, original_h, original_w, _ = images.shape
|
|
print(f"Original video specs: {total_frames} frames, {original_w}x{original_h}")
|
|
|
|
# 1. Trim frame count to be 4n+1
|
|
# We find the largest number <= total_frames that satisfies the condition.
|
|
new_total_frames = total_frames - ((total_frames - 1) % 4)
|
|
|
|
if new_total_frames != total_frames:
|
|
print(f"Trimming frames to be 4n+1: {total_frames} -> {new_total_frames}")
|
|
images = images[:new_total_frames, :, :, :]
|
|
else:
|
|
print("Frame count already meets 4n+1 requirement. No trimming needed.")
|
|
|
|
# 2. Crop dimensions to the nearest multiple of 8 (rounding down)
|
|
new_h = (original_h // 8) * 8
|
|
new_w = (original_w // 8) * 8
|
|
|
|
if new_h != original_h or new_w != original_w:
|
|
print(f"Cropping dimensions to a multiple of 8: {original_w}x{original_h} -> {new_w}x{new_h}")
|
|
|
|
# Calculate the amount to remove from each side for a center crop
|
|
h_to_remove = original_h - new_h
|
|
w_to_remove = original_w - new_w
|
|
|
|
h_start = h_to_remove // 2
|
|
w_start = w_to_remove // 2
|
|
|
|
# Perform the centered crop using tensor slicing
|
|
processed_images = images[:, h_start : h_start + new_h, w_start : w_start + new_w, :]
|
|
else:
|
|
print("Dimensions are already multiples of 8. No cropping needed.")
|
|
processed_images = images
|
|
|
|
print(f"Final video specs: {processed_images.shape[0]} frames, {processed_images.shape[2]}x{processed_images.shape[1]}")
|
|
|
|
return (processed_images,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"VideoFramePreprocessor": VideoFramePreprocessor,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"VideoFramePreprocessor": "Stand-In Trimmer & Cropper",
|
|
}
|