Spaces:
Running
on
T4
Running
on
T4
File size: 1,945 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 66 67 68 |
import sys, os, random
import cv2, torch
from multiprocessing import Process, Queue
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 JPEG():
def __init__(self) -> None:
# Choose an image compression degradation
# self.jpeger = DiffJPEG(differentiable=False).cuda()
pass
def compress_and_store(self, np_frames, store_path, idx):
''' Compress and Store the whole batch as JPEG
Args:
np_frames (numpy): The numpy format of the data (Shape:?)
store_path (str): The store path
Return:
None
'''
# Preparation
single_frame = np_frames
# Compress as JPEG
jpeg_quality = random.randint(*opt['jpeg_quality_range2'])
encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality]
_, encimg = cv2.imencode('.jpg', single_frame, encode_param)
decimg = cv2.imdecode(encimg, 1)
# Store the image with quality
cv2.imwrite(store_path, decimg)
@staticmethod
def compress_tensor(tensor_frames):
''' Compress tensor input to JPEG and then return it
Args:
tensor_frame (tensor): Tensor inputs
Returns:
result (tensor): Tensor outputs (same shape as input)
'''
single_frame = tensor2np(tensor_frames)
# Compress as JPEG
jpeg_quality = random.randint(*opt['jpeg_quality_range1'])
encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality]
_, encimg = cv2.imencode('.jpg', single_frame, encode_param)
decimg = cv2.imdecode(encimg, 1)
# Store the image with quality
# cv2.imwrite(store_name, decimg)
result = np2tensor(decimg)
return result
|