添加视频尺寸参数

__init__.py

@@ -5,7 +5,7 @@ from PIL import Image
 import numpy as np
 import torch
 import comfy.utils
-from .videoCut import getCutList, videoToPng, cutToDir, frames_to_video
+from .videoCut import getCutList, video_to_frames, cutToDir, frames_to_video
 from .seg import get_masks
 from .thick_lines_from_canny import fill_white_segments, find_largest_white_component
 
@@ -447,6 +447,13 @@ class VideoToFrame:
             "required": {
                 "video_path": ("STRING", {"default": None}),
                 "save_name": ("STRING", {"default": "temp"}),
+                "min_side_length": ("INT", {
+                    "default": 512,
+                    "min": 1,
+                    "max": 4096,
+                    "step": 1,
+                    "display": "number"
+                }),
                 "frame_rate": ("INT", {
                     "default": 24,
                     "min": 1,
@@ -462,9 +469,9 @@ class VideoToFrame:
 
     CATEGORY = "badger"
 
-    def video_to_frame(self, video_path, save_name, frame_rate):
+    def video_to_frame(self, video_path, save_name,min_side_length, frame_rate):
         videoPath = os.path.abspath(video_path)
-        imagePath = videoToPng(videoPath, frame_rate, save_name)
+        imagePath = video_to_frames(videoPath,min_side_length, frame_rate, save_name)
 
         return (imagePath,)
 
@@ -760,12 +767,12 @@ class ApplyMaskToImage:
         return (imgToTensor(result_image),)
 
 
-class deleteDir:
+class DeleteDir:
     @classmethod
     def INPUT_TYPES(s):
         return {
             "required": {
-                "start": ("VHS_FILENAMES",),
+                "start": ("STRING", {"default": None}),
                 "dir_path": ("STRING", {"default": ""}),
             }
         }
@@ -969,7 +976,7 @@ NODE_CLASS_MAPPINGS = {
     "SegmentToMaskByPoint-badger": SegmentToMaskByPoint,
     "CropImageByMask-badger": CropImageByMask,
     "ApplyMaskToImage-badger": ApplyMaskToImage,
-    "deleteDir-badger": deleteDir,
+    "deleteDir-badger": DeleteDir,
     "FindThickLinesFromCanny-badger": FindThickLinesFromCanny,
     "TrimTransparentEdges-badger": TrimTransparentEdges,
     "ExpandImageWithColor-badger": ExpandImageWithColor,

videoCut.py

@@ -7,13 +7,14 @@ import cv2
 from sentence_transformers import util
 from PIL import Image
 from skimage import metrics
+from moviepy.editor import VideoFileClip
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-16-plus-240', pretrained="laion400m_e32")
 model.to(device)
 
 
-def SSIM(imgPath0,imgPath1):
+def SSIM(imgPath0, imgPath1):
     image1 = cv2.imread(imgPath0)
     image2 = cv2.imread(imgPath1)
     image2 = cv2.resize(image2, (image1.shape[1], image1.shape[0]), interpolation=cv2.INTER_AREA)
@@ -22,24 +23,26 @@ def SSIM(imgPath0,imgPath1):
     image2_gray = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
     # Calculate SSIM
     ssim_score = metrics.structural_similarity(image1_gray, image2_gray, full=True)
-    return round(ssim_score[0], 2)*100
+    return round(ssim_score[0], 2) * 100
 
 
 def imageEncoder(img):
-
     img1 = Image.fromarray(img).convert('RGB')
     img1 = preprocess(img1).unsqueeze(0).to(device)
     img1 = model.encode_image(img1)
     return img1
+
+
 def generateScore(image1, image2):
     test_img = cv2.imread(image1, cv2.IMREAD_UNCHANGED)
     data_img = cv2.imread(image2, cv2.IMREAD_UNCHANGED)
     img1 = imageEncoder(test_img)
     img2 = imageEncoder(data_img)
     cos_scores = util.pytorch_cos_sim(img1, img2)
-    score = round(float(cos_scores[0][0])*100, 2)
+    score = round(float(cos_scores[0][0]) * 100, 2)
     return score
 
+
 def getCutList(imagepath, min_frame, max_frame):
     pngList = sorted(os.listdir(imagepath))
     cutList = []
@@ -47,12 +50,12 @@ def getCutList(imagepath, min_frame, max_frame):
     indexList = []
     i = min_frame - 1
     num = 0
-    while i < len(pngList)-1:
+    while i < len(pngList) - 1:
         num += 1
-        imgPath0 = os.path.join(imagepath,pngList[i])
-        imgPath1 = os.path.join(imagepath,pngList[i+1])
-        res = generateScore(imgPath0,imgPath1)
-        print("切割画面（"+str(i+1)+"/"+str(len(pngList)-1) + ")" + str(res))
+        imgPath0 = os.path.join(imagepath, pngList[i])
+        imgPath1 = os.path.join(imagepath, pngList[i + 1])
+        res = generateScore(imgPath0, imgPath1)
+        print("切割画面（" + str(i + 1) + "/" + str(len(pngList) - 1) + ")" + str(res))
         resList.append(res)
         indexList.append(i)
         if num >= max_frame:
@@ -62,7 +65,7 @@ def getCutList(imagepath, min_frame, max_frame):
         elif res < 95:
             res2 = SSIM(imgPath0, imgPath1)
             if res2 < 60:
-                res3 = (95-res)**2+(60-res2)**2
+                res3 = (95 - res) ** 2 + (60 - res2) ** 2
                 print(res3)
                 if (res3 > 150):
                     num = 0
@@ -78,26 +81,38 @@ def getCutList(imagepath, min_frame, max_frame):
     return cutList
 
 
-def videoToPng(videopath, rate, save_name):
+def video_to_frames(video_path, min_side_length, frame_rate, save_name):
     current_file_path = __file__
     absolute_path = os.path.abspath(current_file_path)
     directory = os.path.dirname(absolute_path)
 
     all_cut_dir = os.path.join(directory, "VideoCutDir")
+
+    video = VideoFileClip(video_path)
+    width, height = video.size
+
+    if height < width:
+        width = -1
+        height = min_side_length
+    else:
+        width = min_side_length
+        height = -1
+
     if not os.path.exists(all_cut_dir):
         os.mkdir(all_cut_dir)
     root_dir = os.path.join(all_cut_dir, save_name)
     if os.path.exists(root_dir):
         shutil.rmtree(root_dir)
     os.mkdir(root_dir)
+    size = str(width) + ":" + str(height)
     ffmpegCMD = [
         'ffmpeg',
-        '-i', videopath,  # 输入视频路径
-        '-vf', f'fps={rate}',  # 设置帧率
+        '-i', video_path,  # 输入视频路径
+        '-vf', f'fps={frame_rate},scale={size}',  # 设置帧率和缩放
         '-q:v', '2',  # 设置输出质量
         os.path.join(root_dir, 'frame_%05d.png')  # 输出文件名和格式
     ]
-    subprocess.run(ffmpegCMD)
+    subprocess.run(ffmpegCMD, check=True)
     print("视频转图片完成")
 
     return root_dir