from PIL import Image
import numpy as np
import cv2
import io


def compress_image(img: "np.ndarray", target_size: "int"):
    cur_size = img.shape[0] * img.shape[1]
    if cur_size <= target_size:
        return img

    image = Image.fromarray(img)
    output = io.BytesIO()
    image.save(output, format="JPEG", quality=int(target_size / cur_size * 100))
    output.seek(0)
    compressed_image = Image.open(output)

    return np.array(compressed_image)


def main():
    input_path = r'C:\Users\huimv\Pictures\Saved Pictures\to-compress.jpg'
    max_size = 3800 * 1500  # 长 × 宽
    # 读取图像为NumPy数组
    image_array = cv2.imread(input_path)
    # 压缩图像
    compressed_image_array = compress_image(image_array, max_size)
    # 保存压缩后的图像
    output_path = 'img/compress-res.jpg'
    cv2.imwrite(output_path, compressed_image_array)


if __name__ == "__main__":
    main()