69 lines
2.4 KiB
Python
69 lines
2.4 KiB
Python
import numpy as np
|
|
|
|
from PIL import Image
|
|
from skimage.filters import threshold_multiotsu
|
|
|
|
|
|
|
|
def quantization_otsu(img, severity=1):
|
|
"""
|
|
Color Quantization using OTSU method.
|
|
severity=[1, 2, 3, 4, 5] corresponding to num_classes=[15, 11, 8, 5, 3].
|
|
severity mainly refer to KADID-10K and Imagecorruptions.
|
|
|
|
@param img: Input image, H x W x 3, value range [0, 255]
|
|
@param severity: Severity of distortion, [1, 5]
|
|
@return: Degraded image, H x W x 3, value range [0, 255]
|
|
"""
|
|
c = [15, 11, 8, 5, 3][severity - 1]
|
|
img = np.array(img).astype(np.float32)
|
|
for i in range(img.shape[2]):
|
|
img_gray = img[:, :, i]
|
|
thresholds = threshold_multiotsu(img_gray, classes=c, nbins=30) # modify skimage
|
|
v_max = img_gray.max()
|
|
v_min = img_gray.min()
|
|
img[:, :, i] = np.digitize(img[:, :, i], bins=thresholds) * (v_max - v_min) / c + v_min
|
|
img = np.clip(img, 0, 255)
|
|
return img
|
|
|
|
|
|
def quantization_median(img, severity=1):
|
|
"""
|
|
Color Quantization using Histogram Median.
|
|
severity=[1, 2, 3, 4, 5] corresponding to num_classes=[20, 15, 10, 6, 3].
|
|
severity mainly refer to KADID-10K and Imagecorruptions.
|
|
|
|
@param img: Input image, H x W x 3, value range [0, 255]
|
|
@param severity: Severity of distortion, [1, 5]
|
|
@return: Degraded image, H x W x 3, value range [0, 255]
|
|
"""
|
|
c = [20, 15, 10, 6, 3][severity - 1]
|
|
for i in range(img.shape[2]):
|
|
img_gray = Image.fromarray(img[:, :, i])
|
|
img_gray = img_gray.quantize(colors=c, method=Image.Quantize.MEDIANCUT).convert("L")
|
|
img[:, :, i] = np.array(img_gray)
|
|
img = np.clip(img, 0, 255)
|
|
return img
|
|
|
|
|
|
def quantization_hist(img, severity=1):
|
|
"""
|
|
Color Quantization using Histogram Equalization.
|
|
severity=[1, 2, 3, 4, 5] corresponding to num_classes=[24, 16, 8, 6, 4].
|
|
severity mainly refer to KADID-10K and Imagecorruptions.
|
|
|
|
@param img: Input image, H x W x 3, value range [0, 255]
|
|
@param severity: Severity of distortion, [1, 5]
|
|
@return: Degraded image, H x W x 3, value range [0, 255]
|
|
"""
|
|
c = [24, 16, 8, 6, 4][severity - 1]
|
|
hist, _ = np.histogram(img.flatten(), bins=c, range=[0, 255])
|
|
cdf = hist.cumsum()
|
|
cdf_m = np.ma.masked_equal(cdf, 0)
|
|
cdf_m = (cdf_m - cdf_m.min()) * 255 / (cdf_m.max() - cdf_m.min())
|
|
cdf = np.ma.filled(cdf_m, 0).astype('uint8')
|
|
img = np.uint8(np.round(img / 255 * (c - 1)))
|
|
img = cdf[img]
|
|
img = np.clip(img, 0, 255)
|
|
return img
|