2021-06-20

图像二维熵计算

# coding=utf-8
import os

import cv2
import time
import reader
import numpy as np
from collections import Counter

image_A_path_lsit = []
image_B_path_lsit = []
image_A_entropy_lsit = []
image_B_entropy_lsit = []

def calcIJ(img_patch):
    total_p = img_patch.shape[0] * img_patch.shape[1]
    if total_p % 2 != 0:
        center_p = img_patch[int(img_patch.shape[0] / 2), int(img_patch.shape[1] / 2)]
        mean_p = round((np.sum(img_patch) - center_p) / (total_p - 1))
        return (center_p, mean_p)
    else:
        pass

def calcEntropy2dSpeedUp(img, mask, win_w = 3, win_h = 3, format_len = 5):
    # 边缘填充行列，值填0
    height = img.shape[0]
    ext_x = int(win_w / 2)
    ext_y = int(win_h / 2)
    ext_h_part = np.zeros([height, ext_x], img.dtype)
    tem_img = np.hstack((ext_h_part, img, ext_h_part))
    ext_v_part = np.zeros([ext_y, tem_img.shape[1]], img.dtype)
    final_img = np.vstack((ext_v_part, tem_img, ext_v_part))
    new_width = final_img.shape[1]
    new_height = final_img.shape[0]

    # 遍历计算二元组
    IJ = []
    for i in range(ext_x, new_width - ext_x):
        for j in range(ext_y, new_height - ext_y):
            # patch对应一个win_w*win_h领域
            if mask[j - ext_y][i - ext_x] == 0:
                continue
            patch = final_img[j - ext_y:j + ext_y + 1, i - ext_x:i + ext_x + 1]
            ij = calcIJ(patch)
            IJ.append(ij)

    # 计算F(i,j)出现的次数
    Fij = Counter(IJ).items()

    # 计算各二元组出现的概率
    totol_pixel = np.count_nonzero(mask)
    Pij = []
    for item in Fij:
        # tmp = format(item[1] * 1.0 / totol_pixel, format_len)
        # Pij.append(float(tmp))
        Pij.append(item[1] * 1.0 / totol_pixel)

    H_tem = []
    for item in Pij:
        h_tem = -item * (np.log(item) / np.log(2))
        H_tem.append(h_tem)

    H = np.sum(H_tem)
    return H

def get_median(src, mask):
    height = src.shape[0]
    width = src.shape[1]
    pixel_value = []
    for i in range(height):
        for j in range(width):
            if mask[i][j] != 0:
                pixel_value.append(src[i][j])
    return np.median(pixel_value)

def get_folder_img(folder):
    files = reader.get_files(folder, "jpg")
    if len(files) != 2:
        print(folder, " img num error")
        return -1, -1
    for item in files:
        # 找到路径中含有mask，即为mask图路径
        if item.find('mask') != -1:
            mask_img_path = item
        else:
            ori_img_path = item
    ori_img = cv2.imread(ori_img_path, cv2.IMREAD_GRAYSCALE)
    mask_img = cv2.imread(mask_img_path, cv2.IMREAD_GRAYSCALE)
    return ori_img, mask_img

def get_normalized(img, factor):
    return cv2.multiply(img, factor)

def compire_entropy(folders_A, folders_B, patch_size, format_len):
    result = []
    A_to_B = False
    if len(folders_A) != len(folders_B) :
        print("folder num not equal")
    curTime = 0
    totalTimes = len(folders_A)
    for index, folder in enumerate(folders_A):
        curTime += 1
        print("Current progress : ", curTime, "/", totalTimes)
        src_A, mask_A = get_folder_img(folders_A[index])
        src_B, mask_B = get_folder_img(folders_B[index])
        median_A = get_median(src_A, mask_A)
        median_B = get_median(src_B, mask_B)
        if median_A < median_B :
            A_to_B = True
        factor = min(median_A, median_B) / max(median_A, median_B)
        h_A = 0
        h_B = 0
        if A_to_B:
            normalized_A = get_normalized(src_A, factor)
            h_A = calcEntropy2dSpeedUp(normalized_A, mask_A, patch_size, patch_size, format_len)
            h_B = calcEntropy2dSpeedUp(src_B, mask_B, patch_size, patch_size, format_len)
        else:
            normalized_B = get_normalized(src_B, factor)
            h_A = calcEntropy2dSpeedUp(src_A, mask_A, patch_size, patch_size, format_len)
            h_B = calcEntropy2dSpeedUp(normalized_B, mask_B, patch_size, patch_size, format_len)

        if h_A < h_B:
            result.append(0)
        else:
            result.append(1)

        image_A_path_lsit.append(folders_A[index].split('\\')[-1])
        image_B_path_lsit.append(folders_B[index].split('\\')[-1])
        image_A_entropy_lsit.append(h_A)
        image_B_entropy_lsit.append(h_B)
    return result

def save_result(filename):
    with open(filename, 'w') as f_obj:
        for index, entropy in enumerate(image_A_path_lsit):
            str_output = image_A_path_lsit[index] + " " + str(image_A_entropy_lsit[index]) +\
                         " " + image_B_path_lsit[index] + " " + str(image_B_entropy_lsit[index])
            f_obj.write(str_output)
            f_obj.write("\n")

if __name__ == '__main__':
    folders_A = reader.get_folders(".\input\A", "jpg")
    folders_B = reader.get_folders(".\input\B", "jpg")

    patch_size = 3
    format_len = '.5f'
    result = compire_entropy(folders_A, folders_B, patch_size, format_len)
    output_file = "entropy_output.txt"
    save_result(output_file)
    print(result)
    print("Done")