Spaces:
Running
on
T4
Running
on
T4
File size: 2,084 Bytes
561c629 |
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 |
import torch, sys, os, random
import torch.nn.functional as F
import numpy as np
import cv2
from multiprocessing import Process, Queue
from PIL import Image
root_path = os.path.abspath('.')
sys.path.append(root_path)
# Import files from the local folder
from opt import opt
from degradation.ESR.utils import tensor2np, np2tensor
class WEBP():
def __init__(self) -> None:
# Choose an image compression degradation
pass
def compress_and_store(self, np_frames, store_path, idx):
''' Compress and Store the whole batch as WebP (~ VP8)
Args:
np_frames (numpy): The numpy format of the data (Shape:?)
store_path (str): The store path
Return:
None
'''
single_frame = np_frames
# Choose the quality
quality = random.randint(*opt['webp_quality_range2'])
method = random.randint(*opt['webp_encode_speed2'])
# Transform to PIL and then compress
PIL_image = Image.fromarray(np.uint8(single_frame[...,::-1])).convert('RGB')
PIL_image.save(store_path, 'webp', quality=quality, method=method)
@staticmethod
def compress_tensor(tensor_frames, idx = 0):
''' Compress tensor input to WEBP and then return it
Args:
tensor_frame (tensor): Tensor inputs
Returns:
result (tensor): Tensor outputs (same shape as input)
'''
single_frame = tensor2np(tensor_frames)
# Choose the quality
quality = random.randint(*opt['webp_quality_range1'])
method = random.randint(*opt['webp_encode_speed1'])
# Transform to PIL and then compress
PIL_image = Image.fromarray(np.uint8(single_frame[...,::-1])).convert('RGB')
store_path = os.path.join("tmp", "temp_"+str(idx)+".webp")
PIL_image.save(store_path, 'webp', quality=quality, method=method)
# Read back
decimg = cv2.imread(store_path)
result = np2tensor(decimg)
os.remove(store_path)
return result |