import cv2
import os
import face_recognition
import sys
import imagehash
from PIL import Image
import 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)