视频帧AI评分去重需求:相似帧仅保留最高分条目
解决相近帧重复评分的问题
我明白你现在的困扰:视频里大量内容接近的帧,导致生成的评分列表里堆满了评分近乎一致、帧号也相近的冗余条目。咱们可以通过按视频分组→聚类相近帧→保留每组最高分帧的思路来解决这个问题,下面是具体的修改方案:
核心思路拆解
- 统一数据类型:把
predict返回的字符串评分转成浮点数,方便后续比较和排序。 - 按视频分组:你的帧名称格式是
视频号_帧号,可以把同一视频的帧归为一组单独处理,避免跨视频的帧干扰。 - 聚类相近帧:设定一个帧差阈值(比如参考你用的步长500,设为1000——也就是两帧间隔不超过2个步长就算相近,这个值你可以按需调整),把帧号接近的帧划分为同一簇。
- 保留最优帧:在每个簇里,挑选评分最高的那个帧,其余冗余条目直接舍弃。
修改后的完整代码
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))
关键修改点说明
predict函数返回值:把原来的字符串评分改成浮点数,避免后续比较时的类型错误。- 视频分组:用
defaultdict把同一视频的帧归类,确保只在同视频内处理相近帧。 - 相近帧聚类:通过帧号差阈值判断是否属于同一簇,这个阈值
frame_threshold你可以灵活调整——如果觉得当前的2倍步长太宽松,就改小一点,反之则改大。 - 保留最优帧:用
max函数直接取每个簇里评分最高的帧,逻辑简单高效。
处理后,结果列表里就只会保留每个相似帧组的最优帧,不会再出现冗余的重复条目了~
内容的提问来源于stack exchange,提问作者BlinxX




