文章目录
需求
做深度学习需要自己收集图片，其中一种是收集视频，然后将视频转换成图片。在视频转图片过程中，会存在大量的高度相似帧，对于模型训练无用，而且增加标注成本，如何选取有足够差异的图片是我们需要的。
方法
基于图片相似度来选取不同的图片进行保存，相似度计算方法主要参考https://aistudio.baidu.com/projectdetail/4185629?channelType=0\&channel=0 这篇中的方法。
代码
直接上代码，内容简单，很容易看明白。代码中提供基于hash的三种方法和一种结构相似性方法，需要手动改代码来切换方法及相关阈值。
python
import os
import cv2
import numpy as np
import sys
import shutil
from datetime import datetime
from skimage.metrics import structural_similarity as compare_ssim
# 均值哈希算法
def ahash(image):
# 将图片缩放为8*8的
image = cv2.resize(image, (8, 8), interpolation=cv2.INTER_CUBIC)
# 将图片转化为灰度图
gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
# s为像素和初始灰度值，hash_str为哈希值初始值
s = 0
# 遍历像素累加和
for i in range(8):
for j in range(8):
s = s + gray[i, j]
# 计算像素平均值
avg = s / 64
# 灰度大于平均值为1相反为0，得到图片的平均哈希值，此时得到的hash值为64位的01字符串
ahash_str = ''
for i in range(8):
for j in range(8):
if gray[i, j] > avg:
ahash_str = ahash_str + '1'
else:
ahash_str = ahash_str + '0'
result = ''
for i in range(0, 64, 4):
result += ''.join('%x' % int(ahash_str[i: i + 4], 2))
# print("ahash值：",result)
return result
# phash
def phash(img):
# 加载并调整图片为32*32的灰度图片
img1 = cv2.resize(img, (32, 32),cv2.COLOR_RGB2GRAY)
# 创建二维列表
h, w = img.shape[:2]
vis0 = np.zeros((h, w), np.float32)
vis0[:h, :w] = img1
# DCT二维变换
# 离散余弦变换，得到dct系数矩阵
img_dct = cv2.dct(cv2.dct(vis0))
img_dct.resize(8,8)
# 把list变成一维list
img_list = np.array().flatten(img_dct.tolist())
# 计算均值
img_mean = cv2.mean(img_list)
avg_list = ['0' if i<img_mean else '1' for i in img_list]
return ''.join(['%x' % int(''.join(avg_list[x:x+4]),2) for x in range(0,64,4)])
#差异值哈希算法
def dhash(image):
#将图片resize 到8x8
image = cv2.resize(image,(9,8),interpolation=cv2.INTER_CUBIC)
#转成灰度图
gray = cv2.cvtColor(image,cv2.COLOR_RGB2GRAY)
#计算dhash 二进制
dhash_str =""
for i in range(8):
for j in range(8):
if gray[i,j]>gray[i,j+1]:
dhash_str = dhash_str+"1"
else:
dhash_str = dhash_str+"0"
#二进制转十六近制
result = ""
for i in range(0,64,4):
result += "".join("%x" %int(dhash_str[i:i+4],2))
return result
# 计算两个哈希值之间的差异
def campHash(hash1, hash2):
n = 0
# hash长度不同返回-1,此时不能比较
if len(hash1) != len(hash2):
return -1
# 如果hash长度相同遍历长度
for i in range(len(hash1)):
if hash1[i] != hash2[i]:
n = n + 1
return n
def extract_frames(video_path, similarity_threshold, output_dir):
# 读取视频文件
cap = cv2.VideoCapture(video_path)
# 创建输出文件夹
if os.path.exists(output_dir):
shutil.rmtree(output_dir)
os.makedirs(output_dir)
#要保存的图片
previous_image=None
frame_count=0
# 遍历视频帧
while True:
# 读取一帧
ret, frame = cap.read()
# 如果读取到最后一帧，退出循环
if not ret:
break
# 将帧转换为图像
image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
if previous_image is None:
previous_image=cv2.resize(image,(128,128))
# 获取当前时间
now = datetime.now()
# 格式化成指定的时间格式
formatted_time = now.strftime("%Y_%m_%d-%H_%M_%S")
# 保存帧
cv2.imwrite(os.path.join(output_dir,f"{formatted_time}_{frame_count}.jpg"),frame)
continue
else:
# 计算图像之间的相似度
current_image = cv2.resize(image,(128,128))
#ssim
#similarity = compare_ssim(current_image, previous_image,channel_axis=2)
#差异hash
hash1 = ahash(previous_image)
hash2 = ahash(current_image)
similarity = campHash(hash1,hash2)
# ssim如果相似度小于阈值，则不够相似，则抽取帧
#if similarity < similarity_threshold:
# dhash如果相似度大于阈值，则不够相似，则抽取帧
if similarity > similarity_threshold:
# 获取当前时间
now = datetime.now()
# 格式化成指定的时间格式
formatted_time = now.strftime("%Y_%m_%d-%H_%M_%S")
# 保存帧
cv2.imwrite(os.path.join(output_dir,f"{formatted_time}_{frame_count}.jpg"),frame)
# 更新上一帧
previous_image = current_image
frame_count += 1
print(".",end="")
sys.stdout.flush()
cap.release()
if __name__ == "__main__":
# 视频路径
video_path = "../jiabo/20230829/跳远30fps_20230829194849_CH01.avi"
# ssim相似度阈值
#similarity_threshold = 0.9
# dhash 相似度阈值
similarity_threshold = 10
# 输出文件夹
output_dir = "split_output_ahash"
# 抽取帧
extract_frames(video_path, similarity_threshold, output_dir)