You need to enable JavaScript to run this app.
最新活动
大模型
产品
解决方案
定价
生态与合作
支持与服务
开发者
了解我们

视频帧AI评分去重需求:相似帧仅保留最高分条目

解决相近帧重复评分的问题

我明白你现在的困扰:视频里大量内容接近的帧,导致生成的评分列表里堆满了评分近乎一致、帧号也相近的冗余条目。咱们可以通过按视频分组→聚类相近帧→保留每组最高分帧的思路来解决这个问题,下面是具体的修改方案:

核心思路拆解

  1. 统一数据类型:把predict返回的字符串评分转成浮点数,方便后续比较和排序。
  2. 按视频分组:你的帧名称格式是视频号_帧号,可以把同一视频的帧归为一组单独处理,避免跨视频的帧干扰。
  3. 聚类相近帧:设定一个帧差阈值(比如参考你用的步长500,设为1000——也就是两帧间隔不超过2个步长就算相近,这个值你可以按需调整),把帧号接近的帧划分为同一簇。
  4. 保留最优帧:在每个簇里,挑选评分最高的那个帧,其余冗余条目直接舍弃。

修改后的完整代码

import os
import torch
import torchvision.models
import torchvision.transforms as transforms
from PIL import Image
import json
import cv2
from collections import defaultdict

def prepare_image(image):
    if image.mode != 'RGB':
        image = image.convert("RGB")
    Transform = transforms.Compose([
        transforms.Resize([224, 224]),
        transforms.ToTensor(),
    ])
    image = Transform(image)
    image = image.unsqueeze(0)
    return image

def predict(image, model):
    image = prepare_image(image)
    with torch.no_grad():
        preds = model(image)
        score = preds.detach().numpy().item()
        rounded_score = round(score, 2)
        print(f"Picture score: {rounded_score} | frames left: {framesToDo}")
        return rounded_score  # 直接返回浮点数,避免后续类型转换麻烦

if __name__ == '__main__':
    model = torchvision.models.resnet50()
    model.fc = torch.nn.Linear(in_features=2048, out_features=1)
    model.load_state_dict(torch.load('model/model-resnet50.pth', map_location=torch.device('cpu')))
    model.eval()
    result = []
    # 视频文件夹中存储着命名为1至23的视频
    for i in range(1, 23):
        vidcap = cv2.VideoCapture('./video/' + str(i) + '.mp4')
        succes, vidcap_image = vidcap.read()
        count = 0
        framestep = 500 # Stackoverflow示例用步长
        framesToDo = vidcap.get(cv2.CAP_PROP_FRAME_COUNT)
        # 循环读取帧
        while succes and count < int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT)):
            name = str(i) + '_' + str(count)
            cv2.imwrite("./frames_saved/" + 'vid' + '_' + name + ".jpg", vidcap_image)
            count += framestep
            framesToDo = framesToDo - framestep
            cv2_image = cv2.cvtColor(vidcap_image, cv2.COLOR_BGR2RGB)
            pil_image = Image.fromarray(cv2_image)
            # 存储(浮点数评分, 帧名称)
            result.append((predict(pil_image, model), name))
            succes, vidcap_image = vidcap.read()

    # ------------------- 新增的去重逻辑 -------------------
    # 1. 按视频分组,key是视频编号,value是该视频的所有(评分, 帧号, 帧名称)
    video_groups = defaultdict(list)
    for score, name in result:
        vid_id, frame_num = name.split('_')
        frame_num = int(frame_num)
        video_groups[vid_id].append( (score, frame_num, name) )

    filtered_result = []
    # 设定相近帧的阈值:帧号差小于这个值就算相近,可根据需求调整
    frame_threshold = framestep * 2

    for vid_id, frames in video_groups.items():
        # 2. 按帧编号从小到大排序,方便聚类
        frames_sorted = sorted(frames, key=lambda x: x[1])
        # 3. 聚类相近帧
        clusters = []
        current_cluster = [frames_sorted[0]]
        for frame in frames_sorted[1:]:
            last_frame_num = current_cluster[-1][1]
            if frame[1] - last_frame_num <= frame_threshold:
                current_cluster.append(frame)
            else:
                clusters.append(current_cluster)
                current_cluster = [frame]
        clusters.append(current_cluster)  # 加入最后一个未完成的簇

        # 4. 在每个簇里选评分最高的帧
        for cluster in clusters:
            best_frame = max(cluster, key=lambda x: x[0])
            filtered_result.append( (best_frame[0], best_frame[2]) )

    # ------------------- 去重逻辑结束 -------------------

    # 按评分排序(reverse=False是升序,想从高到低就改成True)
    filtered_result.sort(reverse=False)
    print(filtered_result)
    with open('result.json', 'w') as filehandle:
        filehandle.write(json.dumps(filtered_result))

关键修改点说明

  1. predict函数返回值:把原来的字符串评分改成浮点数,避免后续比较时的类型错误。
  2. 视频分组:用defaultdict把同一视频的帧归类,确保只在同视频内处理相近帧。
  3. 相近帧聚类:通过帧号差阈值判断是否属于同一簇,这个阈值frame_threshold你可以灵活调整——如果觉得当前的2倍步长太宽松,就改小一点,反之则改大。
  4. 保留最优帧:用max函数直接取每个簇里评分最高的帧,逻辑简单高效。

处理后,结果列表里就只会保留每个相似帧组的最优帧,不会再出现冗余的重复条目了~

内容的提问来源于stack exchange,提问作者BlinxX

火山引擎 最新活动