To wy : 图像二维熵计算

reader.py

# -*- coding: utf-8 -*-
import os


def file_ext(filename, level=1):
    """
    return extension of filename

    Parameters:
    -----------
    filename: str
        name of file, path can be included
    level: int
        level of extension.
        for example, if filename is 'sky.png.bak', the 1st level extension
        is 'bak', and the 2nd level extension is 'png'

    Returns:
    --------
    extension of filename
    """
    return filename.split('.')[-level]


def _contain_file(path, extensions):
    """
    check whether path contains any file whose extension is in extensions list

    Parameters:
    -----------
    path: str
        path to be checked
    extensions: str or list/tuple of str
        extension or extensions list

    Returns:
    --------
    return True if contains, else return False
    """
    assert os.path.exists(path), 'path must exist'
    assert os.path.isdir(path), 'path must be dir'

    if isinstance(extensions, str):
        extensions = [extensions]

    for file in os.listdir(path):
        if os.path.isfile(os.path.join(path, file)):
            if (extensions is None) or (file_ext(file) in extensions):
                return True
    return False


def _process_extensions(extensions=None):
    """
    preprocess and check extensions, if extensions is str, convert it to list.

    Parameters:
    -----------
    extensions: str or list/tuple of str
        file extensions

    Returns:
    --------
    extensions: list/tuple of str
        file extensions
    """
    if extensions is not None:
        if isinstance(extensions, str):
            extensions = [extensions]
        assert isinstance(extensions, (list, tuple)), \
            'extensions must be str or list/tuple of str'
        for ext in extensions:
            assert isinstance(ext, str), 'extension must be str'
    return extensions


def get_files(path, extensions=None, is_recursive=True):
    """
    read files in path. if extensions is None, read all files, if extensions
    are specified, only read the files who have one of the extensions. if
    is_recursive is True, recursively read all files, if is_recursive is False,
    only read files in current path.

    Parameters:
    -----------
    path: str
        path to be read
    extensions: str or list/tuple of str
        file extensions
    is_recursive: bool
        whether read files recursively. read recursively is True, while just
        read files in current path if False

    Returns:
    --------
    files: the obtained files in path
    """
    extensions = _process_extensions(extensions)
    files = []
    # get files in current path
    if not is_recursive:
        for name in os.listdir(path):
            fullname = os.path.join(path, name)
            if os.path.isfile(fullname):
                if (extensions is None) or (file_ext(fullname) in extensions):
                    files.append(fullname)
        return files
    # get files recursively
    for main_dir, _, sub_file_list in os.walk(path):
        for filename in sub_file_list:
            fullname = os.path.join(main_dir, filename)
            if (extensions is None) or (file_ext(fullname) in extensions):
                files.append(fullname)
    return files


def get_folders(path, extensions=None, is_recursive=True):
    """
    read folders in path. if extensions is None, read all folders, if
    extensions are specified, only read the folders who contain any files that
    have one of the extensions. if is_recursive is True, recursively read all
    folders, if is_recursive is False, only read folders in current path.

    Parameters:
    -----------
    path: str
        path to be read
    extensions: str or list/tuple of str
        file extensions
    is_recursive: bool
        whether read folders recursively. read recursively is True, while just
        read folders in current path if False

    Returns:
    --------
    folders: the obtained folders in path
    """
    extensions = _process_extensions(extensions)
    folders = []
    # get folders in current path
    if not is_recursive:
        for name in os.listdir(path):
            fullname = os.path.join(path, name)
            if os.path.isdir(fullname):
                if (extensions is None) or \
                        (_contain_file(fullname, extensions)):
                    folders.append(fullname)
        return folders
    # get folders recursively
    for main_dir, _, _ in os.walk(path):
        if (extensions is None) or (_contain_file(main_dir, extensions)):
            folders.append(main_dir)
    return folders

main.py

# coding=utf-8
import cv2
import time
import reader
import numpy as np
from collections import Counter

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):
    # 边缘填充行列，值填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:
        # Pij.append(item[1] * 1.0 / (new_height * new_width))
        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 save_result(filename, img_name, output):
    with open(filename, 'w') as f_obj:
        for index, entropy in enumerate(output):
            str_output = img_name[index] + " " + str(entropy)
            f_obj.write(str_output)
            f_obj.write("\n")

def dowork(folders, filename):
    curTime = 0
    totalTimes = len(folders)
    imgs_list= []
    entropy_list = []
    for folder in folders:
        curTime += 1
        print("Current progress : ", curTime, "/", totalTimes)
        files = reader.get_files(folder, "jpg")
        if len(files) != 2:
            print(folder, "don't have two imgs")
            continue
        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)
        H1 = calcEntropy2dSpeedUp(ori_img, mask_img, 3, 3)
        name = ori_img_path.split('\\')[-1]
        imgs_list.append(name)
        entropy_list .append(H1)
    return imgs_list, entropy_list 

if __name__ == '__main__':
    folders = reader.get_folders(".\imgs", "jpg")
    output_file = "entropy_output.txt"
    imgs_list, entropy_list = dowork(folders, output_file)
    save_result(output_file, imgs_list, entropy_list)
    print("Done")