Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,008 Bytes
319886d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import numpy as np
from PIL import Image
from io import BytesIO
def compression_jpeg(img, severity=1):
"""
JPEG compression on a NumPy array.
severity=[1,2,3,4,5] corresponding to quality=[25,18,15,10,7].
from https://github.com/bethgelab/imagecorruptions/blob/master/imagecorruptions/corruptions.py
@param img: Input image as NumPy array, H x W x C, value range [0, 255]
@param severity: Severity of distortion, [1, 5]
@return: Degraded image as NumPy array, H x W x C, value range [0, 255]
"""
assert img.dtype == np.uint8, "Image array should have dtype of np.uint8"
assert severity in [1, 2, 3, 4, 5], 'Severity must be an integer between 1 and 5.'
quality = [25, 18, 12, 8, 5][severity - 1]
output = BytesIO()
gray_scale = False
if img.shape[2] == 1: # Check if the image is grayscale
gray_scale = True
# Convert NumPy array to PIL Image
img = Image.fromarray(img)
if gray_scale:
img = img.convert('L')
else:
img = img.convert('RGB')
# Save image to a bytes buffer using JPEG compression
img.save(output, 'JPEG', quality=quality)
output.seek(0)
# Load the compressed image from the bytes buffer
img_lq = Image.open(output)
# Convert PIL Image back to NumPy array
if gray_scale:
img_lq = np.array(img_lq.convert('L'))
img_lq = img_lq.reshape((img_lq.shape[0], img_lq.shape[1], 1)) # Maintaining the original shape (H, W, 1)
else:
img_lq = np.array(img_lq.convert('RGB'))
return img_lq
def compression_jpeg_2000(img, severity=1):
"""
JPEG2000 compression on a NumPy array.
severity=[1,2,3,4,5] corresponding to quality=[29,27.5,26,24.5,23], quality_mode='dB'.
@param x: Input image as NumPy array, H x W x C, value range [0, 255]
@param severity: Severity of distortion, [1, 5]
@return: Degraded image as NumPy array, H x W x C, value range [0, 255]
"""
assert img.dtype == np.uint8, "Image array should have dtype of np.uint8"
assert severity in [1, 2, 3, 4, 5], 'Severity must be an integer between 1 and 5.'
quality = [29, 27.5, 26, 24.5, 23][severity - 1]
output = BytesIO()
gray_scale = False
if img.shape[2] == 1: # Check if the image is grayscale
gray_scale = True
# Convert NumPy array to PIL Image
img = Image.fromarray(img)
if gray_scale:
img = img.convert('L')
else:
img = img.convert('RGB')
# Save image to a bytes buffer using JPEG compression
img.save(output, 'JPEG2000', quality_mode='dB', quality_layers=[quality])
output.seek(0)
# Load the compressed image from the bytes buffer
img_lq = Image.open(output)
# Convert PIL Image back to NumPy array
if gray_scale:
img_lq = np.array(img_lq.convert('L'))
img_lq = img_lq.reshape((img_lq.shape[0], img_lq.shape[1], 1)) # Maintaining the original shape (H, W, 1)
else:
img_lq = np.array(img_lq.convert('RGB'))
return img_lq
|