1
import numpy as np
2
import matplotlib.pyplot as plt
3
import PIL.Image
4
import argparse
5

6
def load_img(path):
7
    img = PIL.Image.open(path)
8
    img = np.array(img).astype('float32')
9
    return img
10

11
def compress(path, rate):
12
    img = load_img(path)
13
    img = np.transpose(img, (2, 0, 1)) # (m, n, 3) -> (3, m, n)
14
    u, s, v = np.linalg.svd(img) # SVD分解
15
    if rate >= 1:
16
        return u, s, v
17
    m, n = img.size
18
    k = rate_to_k(m, n, rate)
19
    return truncate(u, s, v, k)
20

21
def truncate(u, s, v, k): # 截断矩阵
22
    u = u[..., :k] # 保留前k列
23
    s = s[:, :k] # 保留前k个奇异值
24
    v = v[:, :k] # 保留前k行
25
    return u, s, v
26

27
def rate_to_k(m, n, rate):
28
    '''
29
    设原图像size为m*n
30
    则占用空间为m*n*3
31
    设保留k个奇异值
32
    压缩后占用空间为(m+n+1)*k*3*4
33
    压缩率为rate=(m+n+1)*k*4/(m*n)
34
    k = rate*m*n/((m+n+1)*4)
35
    '''
36
    return int(rate*m*n/((m+n+1)*4))
37

38

39
def decompress(u, s, v):
40
    img = (u * s[:, np.newaxis]) @ v # (3, m, k) * (3, 1, k) @ (3, k, n) -> (3, m, n)
41
    img = np.transpose(img, (1, 2, 0)) # (3, m, n) -> (m, n, 3)
42
    img = np.round(img.clip(0, 255)).astype('uint8')
43
    return img
44

45
def preview(path, rates, col=5):
46
    row = (len(rates) + col - 1) // col
47
    fig, axes = plt.subplots(row, col)
48
    for i in axes.flat:
49
        i.axis('off')
50
    u, s, v = compress(path, 1)
51
    m, n = PIL.Image.open(path).size
52
    for i, rate in enumerate(rates):
53
        k = rate_to_k(m, n, rate)
54
        img = decompress(*truncate(u, s, v, k))
55
        ax = axes[i // col, i % col]
56
        ax.set_title(f'rate={rate}')
57
        ax.imshow(img)
58
    plt.show()
59

60
def save(path, u, s, v):
61
    np.savez_compressed(path, u=u, s=s, v=v)
62

63
def load_c(path):
64
    d = np.load(path)
65
    return d['u'], d['s'], d['v']
66

67
def main(): # main里的内容并不重要，这是使用ChatGPT自动生成的命令行界面，便于使用。
68
    parser = argparse.ArgumentParser(description="SVD Image Compression")
69

70
    # Compression options
71
    parser.add_argument('-c', '--compress', metavar='FILE', help='Compress an image')
72
    parser.add_argument('-o', '--output', metavar='FILE', help='Specify output file for compression')
73
    parser.add_argument('-r', '--rate', type=float, help='Compression rate')
74

75
    # Decompression options
76
    parser.add_argument('-d', '--decompress', metavar='FILE', help='Decompress a compressed file')
77

78
    # Preview options
79
    parser.add_argument('-p', '--preview', metavar='FILE', help='Preview the compressed images')
80
    parser.add_argument('--rates', type=float, nargs='+', help='Specify compression rates for preview')
81

82
    args = parser.parse_args()
83

84
    if args.compress:
85
        u, s, v = compress(args.compress, args.rate or 0.8)
86
        output = args.output or args.compress
87
        if not output.endswith('.npz'):
88
            output += '.npz'
89
        save(output, u, s, v)
90
        print(f'Image compressed and saved to {output}')
91

92
    elif args.decompress:
93
        u, s, v = load_c(args.decompress)
94
        output = args.output or args.decompress
95
        if output.endswith('.npz'):
96
            output = output[:-4]
97
        img = decompress(u, s, v)
98
        PIL.Image.fromarray(img).save(output)
99
        print(f'Image decompressed and saved to {output}')
100

101
    elif args.preview:
102
        rates = args.rates or [
103
            1, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1
104
        ]
105
        preview(args.preview, rates)
106

107
    else:
108
        parser.print_help()
109

110
if __name__ == '__main__':
111
    main()