Add files via upload

2023-11-10 16:00:40 -06:00
parent 9419e42652
commit 72342def15
1 changed files with 78 additions and 0 deletions
@@ -0,0 +1,78 @@
+import cv2
+import torch
+import torchvision.transforms as transforms
+import numpy as np
+import os
+
+class BreakFrames:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "file_input": ("STRING", {
+                    "multiline": False,
+                    "default": "C:/Videos/video.mp4"
+                }),
+                "num_keyframes": ("INT", {
+                        "default": 0, 
+                        "min": 0,
+                        "max": 4096,
+                        "step": 1
+                }),
+            },
+        }
+    RETURN_TYPES = ("IMAGE", "IMAGE")
+    RETURN_NAMES = ("Frames","Keyframes")
+
+    FUNCTION = "breakframes"
+    CATEGORY = "Frames"
+
+    def breakframes(self, file_input, num_keyframes):
+        if not os.path.exists(file_input):
+            raise FileNotFoundError(f"File '{file_input} cannot be found.'")
+        # Open the video file
+        try:
+            video_capture = cv2.VideoCapture(file_input)
+            # Check if the anim was loaded
+            if not video_capture.isOpened():
+                print(f"Error: Could not open video file {file_input}.")
+                return None
+        except:
+            print(f"Error: Could not open video file {file_input}.")
+            return None
+        # Read each frame from anim
+        frame_count = 0
+        transformer = transforms.ToTensor()
+        tensors = []
+        while True:
+            ret, frame = video_capture.read()
+            if ret:
+                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                tensors.append(transformer(frame).unsqueeze(0))
+                frame_count += 1
+            else:
+                break
+        video_capture.release()
+        if num_keyframes > 0:
+            N = np.clip(num_keyframes, 2, len(tensors)-1) # Save a spot for first frame
+            differences = [torch.norm(tensors[i+1] - tensors[i], p=2) for i in range(len(tensors)-1)]
+            _, top_indices = torch.topk(torch.tensor(differences), k=N, largest=True)
+            keyframe_indices = sorted([index.item() + 1 for index in top_indices])
+            keyframe_indices.insert(0, 0)
+            cat_keyframe_tensors = [tensors[i] for i in keyframe_indices]
+            cat_keyframe_tensors = torch.cat(cat_keyframe_tensors, dim = 0).permute(0, 2, 3, 1)
+        else:
+            cat_keyframe_tensors = None
+        cat_frame_tensors = torch.cat(tensors, dim = 0).permute(0, 2, 3, 1)
+        return (cat_frame_tensors, cat_keyframe_tensors)
+
+NODE_CLASS_MAPPINGS = {
+    "BreakFrames": BreakFrames
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "BreakFrames": "BreakFrames"
+}