|
|
import cv2 as cv |
|
|
import numpy as np |
|
|
from PIL import Image |
|
|
|
|
|
DCT_SIZE = 8 |
|
|
TABLE_SIZE = DCT_SIZE ** 2 |
|
|
ZIG_ZAG = [ |
|
|
[0, 0], |
|
|
[0, 1], |
|
|
[1, 0], |
|
|
[2, 0], |
|
|
[1, 1], |
|
|
[0, 2], |
|
|
[0, 3], |
|
|
[1, 2], |
|
|
[2, 1], |
|
|
[3, 0], |
|
|
[4, 0], |
|
|
[3, 1], |
|
|
[2, 2], |
|
|
[1, 3], |
|
|
[0, 4], |
|
|
[0, 5], |
|
|
[1, 4], |
|
|
[2, 3], |
|
|
[3, 2], |
|
|
[4, 1], |
|
|
[5, 0], |
|
|
[6, 0], |
|
|
[5, 1], |
|
|
[4, 2], |
|
|
[3, 3], |
|
|
[2, 4], |
|
|
[1, 5], |
|
|
[0, 6], |
|
|
[0, 7], |
|
|
[1, 6], |
|
|
[2, 5], |
|
|
[3, 4], |
|
|
[4, 4], |
|
|
[5, 3], |
|
|
[6, 2], |
|
|
[7, 1], |
|
|
[7, 2], |
|
|
[6, 3], |
|
|
[5, 4], |
|
|
[4, 5], |
|
|
[3, 5], |
|
|
[2, 6], |
|
|
[1, 7], |
|
|
[2, 7], |
|
|
[3, 6], |
|
|
[4, 5], |
|
|
[5, 4], |
|
|
[6, 3], |
|
|
[7, 2], |
|
|
[7, 3], |
|
|
[6, 4], |
|
|
[5, 5], |
|
|
[4, 6], |
|
|
[3, 7], |
|
|
[4, 7], |
|
|
[5, 6], |
|
|
[6, 5], |
|
|
[7, 4], |
|
|
[7, 5], |
|
|
[6, 6], |
|
|
[5, 7], |
|
|
[6, 7], |
|
|
[7, 6], |
|
|
[7, 7], |
|
|
] |
|
|
|
|
|
|
|
|
def compress_jpg(image: Image.Image, quality, color=True): |
|
|
"""Compress a PIL image to JPEG format with specified quality. |
|
|
|
|
|
Args: |
|
|
image: Input PIL image (RGB format) |
|
|
quality: JPEG compression quality (1-100) |
|
|
color: Whether to preserve color (BGR format) |
|
|
|
|
|
Returns: |
|
|
np.ndarray: Decompressed image in BGR or grayscale format |
|
|
""" |
|
|
|
|
|
img_np = np.array(image) |
|
|
if color: |
|
|
img_np = cv.cvtColor(img_np, cv.COLOR_RGB2BGR) |
|
|
|
|
|
_, buffer = cv.imencode(".jpg", img_np, [cv.IMWRITE_JPEG_QUALITY, quality]) |
|
|
return cv.imdecode(buffer, cv.IMREAD_COLOR if color else cv.IMREAD_GRAYSCALE) |
|
|
|
|
|
|
|
|
def loss_curve(image: Image.Image, qualities=tuple(range(1, 101)), normalize=True): |
|
|
"""Calculate JPEG compression loss curve for quality estimation. |
|
|
|
|
|
Args: |
|
|
image: Input PIL image (RGB format) |
|
|
qualities: Quality values to test (1-100) |
|
|
normalize: Whether to normalize the output curve |
|
|
|
|
|
Returns: |
|
|
np.ndarray: Mean absolute difference values across quality levels |
|
|
""" |
|
|
|
|
|
img_np = np.array(image) |
|
|
if len(img_np.shape) == 3: |
|
|
x = cv.cvtColor(img_np, cv.COLOR_RGB2GRAY) |
|
|
else: |
|
|
x = img_np |
|
|
|
|
|
c = np.array( |
|
|
[cv.mean(cv.absdiff(compress_jpg(x, q, False), x))[0] for q in qualities] |
|
|
) |
|
|
if normalize: |
|
|
c = cv.normalize(c, None, 0, 1, cv.NORM_MINMAX).flatten() |
|
|
return c |
|
|
|
|
|
|
|
|
def estimate_qf(image): |
|
|
return np.argmin(loss_curve(image)) |
|
|
|
|
|
|
|
|
def get_tables(quality): |
|
|
luma = 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], |
|
|
] |
|
|
) |
|
|
chroma = 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], |
|
|
] |
|
|
) |
|
|
quality = np.clip(quality, 1, 100) |
|
|
if quality < 50: |
|
|
quality = 5000 / quality |
|
|
else: |
|
|
quality = 200 - quality * 2 |
|
|
tables = np.concatenate((luma[:, :, np.newaxis], chroma[:, :, np.newaxis]), axis=2) |
|
|
tables = (tables * quality + 50) / 100 |
|
|
return np.clip(tables, 1, 255).astype(int) |
|
|
|
