|
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 |
|
|
|
|
|
|
