基于图片相似度对视频进行抽帧

需求

做深度学习需要自己收集图片，其中一种是收集视频，然后将视频转换成图片。在视频转图片过程中，会存在大量的高度相似帧，对于模型训练无用，而且增加标注成本，如何选取有足够差异的图片是我们需要的。

方法

基于图片相似度来选取不同的图片进行保存，相似度计算方法主要参考https://aistudio.baidu.com/projectdetail/4185629?channelType=0&channel=0 这篇中的方法。

代码

直接上代码，内容简单，很容易看明白。代码中提供基于hash的三种方法和一种结构相似性方法，需要手动改代码来切换方法及相关阈值。

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)