Spaces:
Running
on
T4
Running
on
T4
| 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) | |
| 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 | |