解决linux下读取文件顺序不一样导致后续处理出问题

__init__.py

@@ -479,13 +479,6 @@ class VideoCutFromDir:
         return {
             "required": {
                 "frame_dir": ("STRING", {"default": None}),
-                "threshold": ("FLOAT", {
-                    "default": 0.7,
-                    "min": 0.01,
-                    "max": 1.0,
-                    "step": 0.01,
-                    "display": "number"
-                }),
                 "min_frame": ("INT", {
                     "default": 16,
                     "min": 1,
@@ -508,8 +501,8 @@ class VideoCutFromDir:
 
     CATEGORY = "badger"
 
-    def video_cut_from_dir(self, frame_dir,threshold, min_frame, max_frame):
-        cutList = getCutList(frame_dir,threshold, min_frame, max_frame)
+    def video_cut_from_dir(self, frame_dir,min_frame, max_frame):
+        cutList = getCutList(frame_dir, min_frame, max_frame)
         dirPathString = cutToDir(frame_dir, cutList)
 
         return (dirPathString,)

videoCut.py

@@ -1,75 +1,79 @@
 import os
 import subprocess
 import shutil
+import torch
+import open_clip
 import cv2
-import skimage
+from sentence_transformers import util
+from PIL import Image
+from skimage import metrics
+
+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 calculate_image_similarity(img_path1, img_path2):
-    # 读取图片
-    img1 = cv2.imread(img_path1, cv2.IMREAD_GRAYSCALE)
-    img2 = cv2.imread(img_path2, cv2.IMREAD_GRAYSCALE)
-
-    # 检查图片是否有效
-    if img1 is None or img2 is None:
-        raise ValueError("One of the images couldn't be loaded.")
-
-    # 缩放图片到相同大小
-    img2 = cv2.resize(img2, (img1.shape[1], img1.shape[0]))
-
-    # 方法1: 直方图比较
-    hist1 = cv2.calcHist([img1], [0], None, [256], [0, 256])
-    hist2 = cv2.calcHist([img2], [0], None, [256], [0, 256])
-    hist_similarity = cv2.compareHist(hist1, hist2, cv2.HISTCMP_CORREL)
-
-    # 方法2: 结构相似性指数 (SSIM)
-    ssim_similarity = skimage.metrics.structural_similarity(img1, img2)
-
-    # 方法3: 特征点匹配 (使用ORB)
-    orb = cv2.ORB_create()
-    kp1, des1 = orb.detectAndCompute(img1, None)
-    kp2, des2 = orb.detectAndCompute(img2, None)
-    bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
-    matches = bf.match(des1, des2)
-    feature_similarity = len(matches) / float(min(len(des1), len(des2)))
-    # 综合评分
-    combined_score = (hist_similarity + ssim_similarity + feature_similarity) / 3.0
-
-    return combined_score
+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)
+    # Convert images to grayscale
+    image1_gray = cv2.cvtColor(image1, cv2.COLOR_BGR2GRAY)
+    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
 
 
-def getCutList(imagePath, threshold, min_frame, max_frame):
-    pngList = os.listdir(imagePath)
+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)
+    return score
+
+def getCutList(imagepath, min_frame, max_frame):
+    pngList = sorted(os.listdir(imagepath))
     cutList = []
-    indexList = []
     resList = []
+    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])
-        similarity = calculate_image_similarity(imgPath0, imgPath1)
-        print("切割画面(" + str(i + 1) + "/" + str(len(pngList) - 1) + ") 相似度:" + str(similarity))
+        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)
-        resList.append(similarity)
-        i += 1
-
-    i = min_frame - 1
-    avg = sum(resList) / len(resList)
-    threshold = avg * threshold
-    while i < len(resList) - 1:
         if num >= max_frame:
             num = 0
             cutList.append(pngList[i])
             i += min_frame
-        elif resList[i] < threshold:
-            num = 0
-            cutList.append(pngList[i])
-            i += min_frame
+        elif res < 95:
+            res2 = SSIM(imgPath0, imgPath1)
+            if res2 < 60:
+                res3 = (95-res)**2+(60-res2)**2
+                print(res3)
+                if (res3 > 150):
+                    num = 0
+                    cutList.append(pngList[i])
+                    i += min_frame
+                else:
+                    i += 1
+            else:
+                i += 1
         else:
             i += 1
-    print("平均计算阈值："+str(avg)+"   处理阈值："+str(threshold))
     print(cutList)
     return cutList