JPEG图像压缩算法的python实现

import numpy as np import os from PIL import Image import math class KJPEG: def __init__(self,Q=1): # 初始化DCT变换的A矩阵 self.__dctA = np.zeros(shape=(8, 8)) for i in range(8): c = 0 if i == 0: c = np.sqrt(1 / 8) else: c = np.sqrt(2 / 8) for j in range(8): self.__dctA[i, j] = c * np.cos(np.pi * i * (2 * j + 1) / (2 * 8)) # 亮度量化矩阵 self.__lq = np.array([ 16, 11, 10, 16, 24, 40, 51, 61, 12, 12, 14, 19, 26, 58, 60, 55, 14, 13, 16, 24, 40, 57, 69, 56, 14, 17, 22, 29, 51, 87, 80, 62, 18, 22, 37, 56, 68, 109, 103, 77, 24, 35, 55, 64, 81, 104, 113, 92, 49, 64, 78, 87, 103, 121, 120, 101, 72, 92, 95, 98, 112, 100, 103, 99, ]) # 色度量化矩阵 self.__cq = np.array([ 17, 18, 24, 47, 99, 99, 99, 99, 18, 21, 26, 66, 99, 99, 99, 99, 24, 26, 56, 99, 99, 99, 99, 99, 47, 66, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, ]) # 标记矩阵类型，lt是亮度矩阵，ct是色度矩阵 self.__lt = 0 self.__ct = 1 # Zig编码表 self.__zig = np.array([ 0, 1, 8, 16, 9, 2, 3, 10, 17, 24, 32, 25, 18, 11, 4, 5, 12, 19, 26, 33, 40, 48, 41, 34, 27, 20, 13, 6, 7, 14, 21, 28, 35, 42, 49, 56, 57, 50, 43, 36, 29, 22, 15, 23, 30, 37, 44, 51, 58, 59, 52, 45, 38, 31, 39, 46, 53, 60, 61, 54, 47, 55, 62, 63 ]) # Zag编码表 self.__zag = np.array([ 0, 1, 5, 6, 14, 15, 27, 28, 2, 4, 7, 13, 16, 26, 29, 42, 3, 8, 12, 17, 25, 30, 41, 43, 9, 11, 18, 24, 31, 40, 44, 53, 10, 19, 23, 32, 39, 45, 52, 54, 20, 22, 33, 38, 46, 41, 55, 60, 21, 34, 37, 47, 50, 56, 59, 61, 35, 36, 48, 49, 57, 58, 62, 63 ]) self.__Q = Q def rgb2ycbcr(self,rgb_image): if len(rgb_image.shape)!=3 or rgb_image.shape[2]!=3: raise ValueError("input image is not a rgb image") rgb_image = rgb_image.astype(np.float32) # 1：创建变换矩阵，和偏移量 transform_matrix = np.array([[0.257, 0.564, 0.098], [-0.148, -0.291, 0.439], [0.439, -0.368, -0.071]]) shift_matrix = np.array([16, 128, 128]) ycbcr_image = np.zeros(shape=rgb_image.shape) w, h, _ = rgb_image.shape # 2：遍历每个像素点的三个通道进行变换 for i in range(w): for j in range(h): ycbcr_image[i, j, :] = np.dot(transform_matrix, rgb_image[i, j, :]) + shift_matrix #print(ycbcr_image) return ycbcr_image def ycbcr2rgb(self,ycbcr_image): if len(ycbcr_image.shape)!=3 or ycbcr_image.shape[2]!=3: raise ValueError("input image is not a rgb image") ycbcr_image = ycbcr_image.astype(np.float32) transform_matrix = np.array([[0.257, 0.564, 0.098], [-0.148, -0.291, 0.439], [0.439, -0.368, -0.071]]) transform_matrix_inv = np.linalg.inv(transform_matrix) shift_matrix = np.array([16, 128, 128]) rgb_image = np.zeros(shape=ycbcr_image.shape) w, h, _ = ycbcr_image.shape for i in range(w): for j in range(h): rgb_image[i, j, :] = np.dot(transform_matrix_inv, ycbcr_image[i, j, :]) - np.dot(transform_matrix_inv, shift_matrix) return rgb_image.astype(np.uint8) def __Fill(self, matrix): # 图片的长宽都需要满足是16的倍数（采样长宽会缩小1/2和取块长宽会缩小1/8） fh, fw = 0, 0 if self.height % 16 != 0: fh = 16 - self.height % 16 if self.width % 16 != 0: fw = 16 - self.width % 16 res = np.pad(matrix, ((0, fh), (0, fw)), 'constant', constant_values=(0, 0)) return res def __Encode(self, matrix, tag): # 先对矩阵进行填充 matrix = self.__Fill(matrix) # 将图像矩阵切割成8*8小块 height, width = matrix.shape # 减少for循环语句，利用numpy的自带函数来提升算法效率 shape = (height // 8, width // 8, 8, 8) strides = matrix.itemsize * np.array([width * 8, 8, width, 1]) blocks = np.lib.stride_tricks.as_strided(matrix, shape=shape, strides=strides) res = [] for i in range(height // 8): for j in range(width // 8): res.append(self.__Quantize(self.__Dct(blocks[i, j]).reshape(64), tag)) return res def __Dct(self, block): # DCT变换 res = np.dot(self.__dctA, block) res = np.dot(res, np.transpose(self.__dctA)) #print(res) return res def __Quantize(self, block, tag): res = block if tag == self.__lt: res = np.round(res / (self.__lq*self.__Q)) elif tag == self.__ct: res = np.round(res / (self.__cq*self.__Q)) #print(res) return res def __Zig(self, blocks): ty = np.array(blocks) tz = np.zeros(ty.shape) for i in range(len(self.__zig)): tz[:, i] = ty[:, self.__zig[i]] tz = tz.reshape(tz.shape[0] * tz.shape[1]) #print(tz.tolist()) return tz.tolist() def __Rle(self, blist): res = [] cnt = 0 for i in range(len(blist)): if blist[i] != 0: res.append(cnt) res.append(int(blist[i])) cnt = 0 elif cnt == 15: res.append(cnt) res.append(int(blist[i])) cnt = 0 else: cnt += 1 # 末尾全是0的情况 if cnt != 0: res.append(cnt - 1) res.append(0) #print(res) return res def Compress(self, filename): # 根据路径image_path读取图片，并存储为RGB矩阵 image = Image.open(filename) self.width, self.height = image.size image = np.asarray(image) ycbcr = self.rgb2ycbcr(image) y = ycbcr[:, :, 0] cb = ycbcr[:, :, 1] cr = ycbcr[:, :, 2] # 对图像矩阵进行编码 对不同的通道分别进行编码 y_blocks = self.__Encode(y, self.__lt) cb_blocks = self.__Encode(cb, self.__ct) cr_blocks = self.__Encode(cr, self.__ct) # 对图像小块进行Zig编码和RLE编码 y_code = self.__Rle(self.__Zig(y_blocks)) u_code = self.__Rle(self.__Zig(cb_blocks)) v_code = self.__Rle(self.__Zig(cr_blocks)) # 计算VLI可变字长整数编码并写入文件，未实现Huffman部分 buff = 0 tfile = os.path.splitext(filename)[0] + '_%f'%self.__Q + ".gpj" if os.path.exists(tfile): os.remove(tfile) with open(tfile, 'wb') as o: o.write(self.height.to_bytes(2, byteorder='big')) o.flush() o.write(self.width.to_bytes(2, byteorder='big')) o.flush() o.write((len(y_code)).to_bytes(4, byteorder='big')) o.flush() o.write((len(u_code)).to_bytes(4, byteorder='big')) o.flush() o.write(