哈夫曼编码

import os
import cv2
from queue import PriorityQueue
import numpy as np
import math
import struct
import matplotlib.pyplot as plt
class HuffmanNode(object):
    def __init__(self,value,key=None,symbol='',left_child=None,right_child=None):
        self.left_child=left_child
        self.right_child=right_child
        self.value=value
        self.key=key
        assert symbol==''
        self.symbol=symbol
    def __eq__(self,other):
        return self.value==other.value
    def __gt__(self,other):
        return self.value>other.value
    def __lt__(self,other):
        return self.value < other.value
def createTree(hist_dict:dict)->HuffmanNode:
    q=PriorityQueue()
    for k, v in hist_dict.items():
        q.put(HuffmanNode(value=v,key=k))
    while q.qsize()>1:
        l_freq, r_freq=q.get(),q.get()
        node=HuffmanNode(value=l_freq.value+r_freq.value,left_child=l_freq,right_child=r_freq)
        q.put(node)
    return q.get()
def walkTree_VLR(root_node:HuffmanNode,symbol=''):
    global Huffman_encode_dict
    if isinstance(root_node,HuffmanNode):
        root_node.symbol+=symbol
        if root_node.key!=None:
            Huffman_encode_dict[root_node.key]=root_node.symbol
        walkTree_VLR(root_node.left_child,symbol=root_node.symbol+'0')
        walkTree_VLR(root_node.right_child, symbol=root_node.symbol+'1')
    return
def encodeImage(src_img: np.ndarray, encode_dict: dict):
    img_encode=""
    assert len(src_img.shape)==1, '`src_img` must be a vector'
    for pixel in src_img:
        img_encode+=encode_dict[pixel]
    return img_encode
def writeBinImage(img_encode:str,huffman_file:str):
    with open(huffman_file,'wb') as f:
        for i in range(0,len(img_encode),8):
            img_encode_dec=int(img_encode[i:i+8],2)
            img_encode_bin=struct.pack('>B',img_encode_dec)
            f.write(img_encode_bin)
def readBinImage(huffman_file: str,img_encode_len:int):
    code_bin_str=""
    with open(huffman_file,'rb') as f:
        content=f.read()
        code_dec_tuple=struct.unpack('>'+'B'*len(content),content)
        for code_dec in code_dec_tuple:
            code_bin_str+=bin(code_dec)[2:].zfill(8)
        len_diff=len(code_bin_str)-img_encode_len
        code_bin_str=code_bin_str[:-8]+code_bin_str[-(8-len_diff):]
    return code_bin_str
def decodeHuffman(img_encode:str,huffman_tree_root:HuffmanNode):
    img_src_val_list=[]
    root_node=huffman_tree_root
    for code in img_encode:
        if code=='0':
            root_node=root_node.left_child
        elif code=='1':
            root_node=root_node.right_child
        if root_node.key!=None:
            img_src_val_list.append(root_node.key)
            root_node=huffman_tree_root
    return np.asarray(img_src_val_list)
def decodeHuffmanByDict(img_encode:str,encode_dict:dict):
    img_src_val_list=[]
    decode_dict={}
    for k, v in encode_dict.items():
        decode_dict[v]=k
    s=0
    while len(img_encode)>s+1:
        for k in decode_dict.keys():
            if k==img_encode[s:s+len(k)]:
                img_src_val_list.append(decode_dict[k])
                s+=len(k)
                break
    return np.asarray(img_src_val_list)
def put(path):
    src_img=cv2.imread(path,cv2.IMREAD_GRAYSCALE)
    src_img_w, src_img_h=src_img.shape[:2]
    src_img_ravel=src_img.ravel()
    hist_dict={}
    for p in src_img_ravel:
        if p not in hist_dict:
            hist_dict[p]=1
        else:
            hist_dict[p]+=1
    #构造哈夫曼树
    huffman_root_node=createTree(hist_dict)
    #遍历哈夫曼树
    walkTree_VLR(huffman_root_node)
    global Huffman_encode_dict
    print('哈夫曼编码字典：',Huffman_encode_dict)
    img_encode=encodeImage(src_img_ravel,Huffman_encode_dict)
    writeBinImage(img_encode,'huffman_bin_img_file.bin')
    img_read_code=readBinImage('huffman_bin_img_file.bin',len(img_encode))
    #解码二进制编码数据字符串得到原始图像展开的向量
    #根据哈夫曼树进行解码的方式
    img_src_val_array = decodeHuffman(img_read_code, huffman_root_node)
    assert len(img_src_val_array)==src_img_w*src_img_h
    #恢复原始图像
    img_decode=np.reshape(img_src_val_array,[src_img_w,src_img_h])
    #计算平均编码长度和编码效率
    total_code_len=0
    total_code_num=sum(hist_dict.values())
    avg_code_len=0
    I_entropy=0
    for key in hist_dict.keys():
        count=hist_dict[key]
        code_len=len(Huffman_encode_dict[key])
        prob=count/total_code_num
        avg_code_len+=prob*code_len
        I_entropy+=-(prob*math.log2(prob))
    S_eff=I_entropy/avg_code_len
    print("平均编码长度为{:.3f}".format(avg_code_len))
    print("编码效率为{:.6f}".format(S_eff))
    #压缩率
    ori_size=src_img_w*src_img_h*8/(1024*8)
    comp_size=len(img_encode)/(1024*8)
    comp_rate=1-comp_size/ori_size
    print('原图灰度图大小',ori_size,'KB  压缩后大小',comp_size,'KB  压缩率',comp_rate,'%')
    plt.rcParams['font.sans-serif']=['SimHei']
    plt.subplot(121),plt.imshow(src_img,plt.cm.gray),plt.title('原图灰度图像'),plt.axis('off')
    plt.subplot(122),plt.imshow(img_decode,plt.cm.gray),plt.title('解压后'),plt.axis('off')
    plt.show()
if __name__=='__main__':
    #哈夫曼编码字典{pixel_value:code}
    Huffman_encode_dict={}
    put(r'C:/Users/xpp/Desktop/Lena.png')