File size: 3,150 Bytes
6721043 |
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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
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 utils_ours import util
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
# if some values were smaller than black level
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
|