| import os |
| import numpy as np |
| import random |
| import torch |
| from torch.utils.data import Dataset |
| from torchvision import transforms |
| import cv2 |
| import sys |
| from pathlib import Path |
| import glob |
| sys.path.append('..') |
| |
| from dataloader.data_utils import rggb_raw, random_crop, Rawread, path_replace, bayer2raw, read_rawpng |
| from natsort import ns, natsorted |
| import time |
| from tqdm import tqdm |
| from multiprocessing import Pool |
| import pdb |
| import json |
| from fractions import Fraction |
| from pathlib import Path |
| from json import JSONEncoder |
| from exifread.utils import Ratio |
| |
| class imageSet(Dataset): |
| def __init__(self,args): |
| super(imageSet).__init__() |
| self.args = args |
| self.test_dir= args.test_dir |
| self.debug = args.debug |
|
|
| self.paths = [] |
| for file in os.listdir(self.test_dir): |
| if '.png' in file: |
| self.img_path = os.path.join(self.test_dir, file) |
| self.json_path = self.img_path.replace('png', 'json') |
| self.paths.append(dict([(f"{'img_path'}", self.img_path), (f"{'json_path'}", self.json_path)])) |
| |
| def __getitem__(self, index): |
| img_path = self.paths[index]['img_path'] |
| json_path = self.paths[index]['json_path'] |
|
|
| metadata = json_read(json_path, object_hook=fraction_from_json) |
|
|
| input_img = read_rawpng(img_path, metadata) |
|
|
| return {'input': input_img, 'json_path': json_path} |
| |
| def __len__(self): |
| return len(self.paths) |
|
|
|
|
| def normalize(raw_image, black_level, white_level): |
| if type(black_level) is list and len(black_level) == 1: |
| black_level = float(black_level[0]) |
| if type(white_level) is list and len(white_level) == 1: |
| white_level = float(white_level[0]) |
| black_level_mask = black_level |
| if type(black_level) is list and len(black_level) == 4: |
| if type(black_level[0]) is Ratio: |
| black_level = ratios2floats(black_level) |
| if type(black_level[0]) is Fraction: |
| black_level = fractions2floats(black_level) |
| black_level_mask = np.zeros(raw_image.shape) |
| idx2by2 = [[0, 0], [0, 1], [1, 0], [1, 1]] |
| step2 = 2 |
| for i, idx in enumerate(idx2by2): |
| black_level_mask[idx[0]::step2, idx[1]::step2] = black_level[i] |
| normalized_image = raw_image.astype(np.float32) - black_level_mask |
| |
| normalized_image[normalized_image < 0] = 0 |
| normalized_image = normalized_image / (white_level - black_level_mask) |
| return normalized_image |
|
|
| def ratios2floats(ratios): |
| floats = [] |
| for ratio in ratios: |
| floats.append(float(ratio.num) / ratio.den) |
| return floats |
|
|
| def fractions2floats(fractions): |
| floats = [] |
| for fraction in fractions: |
| floats.append(float(fraction.numerator) / fraction.denominator) |
| return floats |
| |
| def json_read(fname, **kwargs): |
| with open(fname) as j: |
| data = json.load(j, **kwargs) |
| return data |
|
|
| def fraction_from_json(json_object): |
| if 'Fraction' in json_object: |
| return Fraction(*json_object['Fraction']) |
| return json_object |
|
|
|
|
|
|
