import cv2
import numpy as np
import imagehash
from PIL import Image
from skimage.metrics import structural_similarity as ssim

# 色度直方图相似度
def calculate_color_histogram_similarity(image_path1, image_path2):
    # 读取图像
    img1 = cv2.imread(image_path1)
    img2 = cv2.imread(image_path2)

    # 将图像转换为 HSV 色彩空间
    hsv1 = cv2.cvtColor(img1, cv2.COLOR_BGR2HSV)
    hsv2 = cv2.cvtColor(img2, cv2.COLOR_BGR2HSV)

    # 计算色度直方图
    hist1 = cv2.calcHist([hsv1], [0, 1], None, [50, 60], [0, 180, 0, 256])
    hist2 = cv2.calcHist([hsv2], [0, 1], None, [50, 60], [0, 180, 0, 256])

    # 归一化直方图
    hist1 = cv2.normalize(hist1, hist1).flatten()
    hist2 = cv2.normalize(hist2, hist2).flatten()

    # 计算相似性
    correlation = cv2.compareHist(hist1, hist2, cv2.HISTCMP_CORREL)
    chi_square = cv2.compareHist(hist1, hist2, cv2.HISTCMP_CHISQR)
    intersection = cv2.compareHist(hist1, hist2, cv2.HISTCMP_INTERSECT)
    bhattacharyya = cv2.compareHist(hist1, hist2, cv2.HISTCMP_BHATTACHARYYA)

    return correlation, chi_square, intersection, bhattacharyya


# SIFT特征相似度
def feature_matching(img1_path, img2_path):
    img1 = cv2.imread(img1_path, cv2.IMREAD_GRAYSCALE)
    img2 = cv2.imread(img2_path, cv2.IMREAD_GRAYSCALE)
    
    # 初始化SIFT检测器
    sift = cv2.SIFT_create()
    kp1, des1 = sift.detectAndCompute(img1, None)
    kp2, des2 = sift.detectAndCompute(img2, None)
    
    # FLANN匹配器
    FLANN_INDEX_KDTREE = 1
    index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
    search_params = dict(checks=50)
    flann = cv2.FlannBasedMatcher(index_params, search_params)
    
    matches = flann.knnMatch(des1, des2, k=2)
    
    # 筛选优质匹配
    good = []
    for m, n in matches:
        if m.distance < 0.7 * n.distance:
            good.append(m)
    
    return len(good) / min(len(kp1), len(kp2))  # 归一化匹配率


# 感知相似度
def phash_similarity(img1_path, img2_path):
    hash1 = imagehash.phash(Image.open(img1_path))
    hash2 = imagehash.phash(Image.open(img2_path))
    return 1 - (hash1 - hash2) / len(hash1.hash) ** 2  # 归一化为0-1


# SSIM相似度
def compare_ssim(img1_path, img2_path):
    img1 = cv2.resize(cv2.imread(img1_path), (256, 256))
    img2 = cv2.resize(cv2.imread(img2_path), (256, 256))
    # 转换为灰度图
    gray1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    gray2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
    # 计算 SSIM
    score, _ = ssim(gray1, gray2, full=True)
    return score

if __name__ == '__main__':
    # 示例用法
    image_path1 = '/data/data/luosy/project/oral/data/key_frame/frame_001.jpg'
    image_path2 = '/data/data/luosy/project/oral/data/key_frame/frame_002.jpg'
    similarity_scores = calculate_color_histogram_similarity(image_path1, image_path2)

    print(f'相似性（⬆）: {similarity_scores[0]}')
    print(f'卡方（⬇）: {similarity_scores[1]}')
    print(f'交集（⬆）: {similarity_scores[2]}')
    print(f'巴氏距离（⬇）: {similarity_scores[3]}')

    match_ratio = feature_matching(image_path1, image_path2)
    print(f"特征点匹配率: {match_ratio:.4f}")

    print("pHash Similarity:", phash_similarity(image_path1, image_path2))

    print("SSIM:", compare_ssim(image_path1, image_path2))