Spaces:
Runtime error
Runtime error
| 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 | |