多媒体实验3：基于 BOF 进行相似图片搜索

import numpy as np
import cv2
import glob
import json
from PIL import Image
from scipy.cluster.vq import vq
from sklearn.cluster import KMeans
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import TfidfTransformer

# debug
corel_path = './multi/corel/'
query_img_path = './multi/query.jpg'
dict_file_path = "./multi/dictionary.json"
bof_file_path = "./multi/bof.json"
tfidf_file_path = "./multi/tf-idf.json"

bag_of_features = []
features = []
tfidf_feature = []


# region preProcess
def get_dictionary(mode):
    """ 根据图像生成特征字典：一共有这么多的特征。\\
        包含SIFT和KMeans两步
    """
    if (mode == "load"):
        with open(dict_file_path, 'r') as f:
            dictionary = np.array(json.load(f))
        return dictionary

    progress = 0
    for imgpath in glob.glob(corel_path+'*/*.jpg'):
        img = cv2.imread(imgpath)
        # 显示进度
        print("load a pic success:"+(progress+1).__str__()+"/"+"100")
        progress += 1

        # 获取img的所有描述子,不需要关键点
        _, des = cv2.SIFT_create().detectAndCompute(img, None)

        features.append(des)

    # vertical stack,垂直堆叠二维数组
    all_features = np.vstack(features)

    # KMeans对象对 all_features 聚类
    k = 64
    kmeans = KMeans(n_clusters=k, max_iter=300, n_init=10)
    kmeans.fit(all_features)

    # 词典即聚类中心
    # k*128的二维数组
    dictionary = kmeans.cluster_centers_

    if (mode == "save"):
        with open(dict_file_path, 'w') as f:
            # ndarray不能直接输出为json
            json.dump(dictionary.tolist(), f)

    return dictionary


def corel_bof(dictionary, mode):
    """ 表示出每个图片包含特征字典中的哪些特征"""
    bag_of_features = []
    if (mode == 'load'):
        with open(bof_file_path, 'r') as f:
            bag_of_features = np.array(json.load(f))
        return bag_of_features

    # 遍历所有描述子
    for des in features:
        print("gennerated a des\n")
        # 使用vq获得图片每个描述子最近的聚类中心（属于哪个聚类中心）
        code, _ = vq(des, dictionary)

        # 生成到直方图并归一化
        # .shape[0]聚类中心的数量
        hist, _ = np.histogram(code, bins=range(dictionary.shape[0]+1))
        hist = hist/np.sum(hist)

        # 添加到列表中1000*64维
        bag_of_features.append(hist)

    if (mode == "save"):
        with open(bof_file_path, 'w') as f:
            json.dump([arr.tolist() for arr in bag_of_features], f)
    return bag_of_features


def tf_idf():
    """"
    计算TF-IDF\\
    TF值 = 词在文档中出现的次数 / 文档总词数\\
    逆向文件频率：IDF值 = log（语料库中包含某个词的文档总数 / 语料库中文档总数）\\
    每个词在每篇文档中的TF-IDF值：TF值*IDF\\
    在这篇“文章”中出现得多的且其他文章出现得少的“词”，才是能标明它身份的“词”\\
    大部分人都有的特征不算\\
    """

    # list[NDArray]->NDArray->稀疏矩阵->NDArray->list
    # 通过NDArray转化为list以输出
    tfidf_feature = TfidfTransformer().fit_transform(
        np.array(bag_of_features)).toarray().tolist()

    with open(tfidf_file_path, 'w') as f:
        json.dump(tfidf_feature, f)
    return tfidf_feature

# endregion


def searchSimilar(queryIMG, dictionary):

    # 对单个图像的ti-idf处理
    _, m_des = cv2.SIFT_create().detectAndCompute(queryIMG, None)
    query_code, _ = vq(m_des, dict)
    query_hist, _ = np.histogram(query_code, bins=range(dictionary.shape[0]+1))
    query_hist = query_hist/np.sum(query_hist)
    query_hist = np.array(query_hist).reshape(1, -1)
    query_tfidf = TfidfTransformer()
    query_tfidf = query_tfidf.fit_transform(query_hist).toarray()

    # reshape(1, -1)：转化为长度为一行的二维数组
    # 返回1*1000的二维矩阵similarity
    similarity = cosine_similarity(query_tfidf.reshape(1, -1), tfidf_feature)
    top10_indices = np.argsort(similarity[0])[-11:][::-1]
    print(top10_indices)
    return top10_indices


def show_result(result_array):
    imgs = []
    total_width = 0
    max_height = 0
    for idx in result_array:
        img_path = corel_path+str(idx//100)+"/"+str(idx)+".jpg"
        img = Image.open(img_path)
        imgs.append(img)
        total_width += img.width
        max_height = max(max_height, img.height)

    # 创建拼接后的图像
    result_image = Image.new('RGB', (total_width, max_height))

    # 拼接所有图像
    x_offset = 0
    for image in imgs:
        result_image.paste(image, (x_offset, 0))
        x_offset += image.size[0]

    # 显示拼接后的图像
    result_image.show()


if __name__ == '__main__':

    # load or save
    file_op = "load"

    # 字典、bof和tf-idf的获取
    dict = get_dictionary(mode=file_op)
    bag_of_features = corel_bof(dict, mode=file_op)
    tfidf_feature = tf_idf()

    # 查询
    query = cv2.imread(query_img_path, cv2.IMREAD_COLOR)
    result = searchSimilar(query, dict)

    # 结果
    show_result(result)