import cv2import osimport face_recognitionimport sysimport imagehashfrom PIL import Imageimport numpy as np# 这是一个基于图片整体相似度的视频帧提取程序try:
# 强制使用CPU，禁用所有硬件加速 cv2.ocl.setUseOpenCL(False)
os.environ["CUDA_VISIBLE_DEVICES"] = "-1" os.environ["DLIB_USE_CUDA"] = "0" # 1. 加载参考图片（加载目录下所有图片） ref_dir = os.path.abspath("./file") # 参考图片所在目录 if not os.path.exists(ref_dir):
print(f"参考图片目录不存在: {ref_dir}")
sys.exit(1)
ref_encodings = [] # 存储所有参考图片的人脸编码 ref_hashes = [] # 存储所有参考图片的哈希值 # 遍历目录下所有图片文件 for filename in os.listdir(ref_dir):
if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
ref_path = os.path.join(ref_dir, filename)
# 用OpenCV加载图片 ref_img_cv = cv2.imread(ref_path)
ref_img_rgb = cv2.cvtColor(ref_img_cv, cv2.COLOR_BGR2RGB)
# 提取人脸编码用于匹配判断 encodings = face_recognition.face_encodings(ref_img_rgb)
if encodings:
ref_encodings.append(encodings[0])
# 计算图片哈希值用于整体相似度比较 pil_img = Image.fromarray(ref_img_rgb)
img_hash = imagehash.phash(pil_img) # 使用感知哈希算法 ref_hashes.append(img_hash)
print(f"加载参考图片: {filename}")
if not ref_encodings:
print("没有有效的参考人脸数据，程序退出")
sys.exit(1)
# 2. 准备视频目录和输出文件夹（使用绝对路径） video_dir = os.path.abspath("./file") # 视频所在目录 if not os.path.exists(video_dir):
print(f"视频目录不存在: {video_dir}")
sys.exit(1)
output_dir = "pic" os.makedirs(output_dir, exist_ok=True)
# 存储已保存图片的哈希值，用于整体相似度去重 saved_hashes = []
# 图片相似度阈值（哈希差异值，越小越相似，通常8-10为宜） IMAGE_SIMILARITY_THRESHOLD = 15 # 统计总匹配图片数量 total_count = 0 # 3. 遍历目录下所有视频文件 video_extensions = ('.mp4', '.avi', '.mov', '.mkv', '.flv')
for video_filename in os.listdir(video_dir):
if video_filename.lower().endswith(video_extensions):
video_path = os.path.join(video_dir, video_filename)
print(f"\n开始处理视频: {video_filename}")
# 打开视频（使用CV_CAP_FFMPEG后端） video = cv2.VideoCapture(video_path, cv2.CAP_FFMPEG)
if not video.isOpened():
print(f"无法打开视频，尝试更换解码器")
# 尝试不同的后端 video = cv2.VideoCapture(video_path, cv2.CAP_MSMF)
if not video.isOpened():
print(f"所有解码器都无法打开视频: {video_filename}，跳过该视频")
continue # 每段视频的帧计数（用于文件名） video_count = 0 # 一共查找多少帧 max_frames = 70000 # 从第几帧开始查找 current_frame = 4000 # 获取视频总帧数，校验初始帧是否合法 total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
if current_frame >= total_frames:
print(f"起始帧 {current_frame} 超过视频总帧数 {total_frames}，跳过该视频")
video.release()
continue # 跳转到设定的初始帧 video.set(cv2.CAP_PROP_POS_FRAMES, current_frame)
while current_frame < max_frames:
ret, frame = video.read()
if not ret:
break print(f"处理第{current_frame}帧...")
# 将图像的颜色通道从 BGR 格式转换为 RGB 格式 rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# 检测人脸 face_locations = face_recognition.face_locations(rgb_frame, model="hog", number_of_times_to_upsample=0)
if face_locations:
# 提取第一个人脸的特征 face_encodings = face_recognition.face_encodings(rgb_frame, [face_locations[0]], num_jitters=1)
if face_encodings:
# 与所有参考人脸比对，只要有一个匹配即可 matches = face_recognition.compare_faces(ref_encodings, face_encodings[0])
if True in matches: # 只要有一个参考人脸匹配 # 计算当前帧的哈希值 pil_frame = Image.fromarray(rgb_frame)
frame_hash = imagehash.phash(pil_frame)
# 判断是否与已保存的图片过于相似 is_duplicate = False if saved_hashes: # 如果已有保存的图片 # 计算与所有已保存图片的哈希差异 for saved_hash in saved_hashes:
hash_diff = abs(frame_hash - saved_hash)
if hash_diff < IMAGE_SIMILARITY_THRESHOLD:
is_duplicate = True print(f"检测到相似图片（哈希差异: {hash_diff}），不重复保存")
break # 不是相似图片才保存 if not is_duplicate:
# 文件名包含总计数，确保唯一 save_path = f"{output_dir}/{total_count}.jpg" cv2.imwrite(save_path, frame)
# 将新图片的哈希值加入已保存列表 saved_hashes.append(frame_hash)
video_count += 1 total_count += 1 print(f"找到新的匹配图片，已保存: {save_path}")
# 强制释放 del frame, rgb_frame, face_locations
current_frame += 500 video.set(cv2.CAP_PROP_POS_FRAMES, current_frame)
video.release()
print(f"视频{video_filename}处理完成，本视频找到{video_count}张匹配图片")
print(f"\n所有视频处理完成，共找到{total_count}张不重复的匹配图片")except Exception as e:
print(f"错误详情: {str(e)}")
sys.exit(1)