def __getitem__(self, idx):
image_name = self.image_names[idx]
image = self.__load(image_name, self.images_path)
image = self.__process_image(image)
gt = self.gts[image_name - 1]
gt_rand = self.gts[(image_name + random.randint(10, 500)) %
len(self.image_names)]
gs, gt_gs = [], []
if self.log_transform:
image, gs = transform_to_log(image)
gt, gt_gs = transform_to_log(gt)
gt_rand, _ = transform_to_log(gt_rand)
augmented = self.transforms(image=image, mask=gs)
img, gs = augmented['image'], augmented['mask']
if self.preprocessing:
preprocessed = self.preprocessing(image=img, mask=gs)
img, gs = preprocessed['image'], preprocessed['mask']
if self.illumination_known:
shape = img.shape
gt_img = np.array([[np.array(gt) for i in range(shape[1])]
for j in range(shape[0])])
gt_rand_img = np.array(
[[np.array(gt_rand) for i in range(shape[1])]
for j in range(shape[0])])
img = np.dstack((img, gt_img, gt_rand_img))
img = torch.tensor(img.transpose(2, 0, 1),
dtype=torch.float32,
device="cuda")
gt = torch.tensor(gt, dtype=torch.float32, device='cuda')
return img, gt, gs, gt_gs
def __getitem__(self, idx):
patches_per_img = ceil(1 / (self.patch_height_ratio * self.patch_width_ratio))
image_idx = int(idx / patches_per_img)
patch_idx = idx % patches_per_img
image_name = self.image_names[image_idx]
image, gt, mask = load_img_and_gt_crf_dataset(image_name, self.path, self.folder, use_mask=self.use_mask,
use_corrected=self.use_corrected, dataset=self.dataset)
image = get_patch_with_index(image, patch_idx, self.patch_height_ratio, self.patch_width_ratio)
mask = get_patch_with_index(mask, patch_idx, self.patch_height_ratio, self.patch_width_ratio)
gt = get_patch_with_index(gt, patch_idx, self.patch_height_ratio, self.patch_width_ratio)
augmented = self.transforms(image=image, mask=mask)
img = augmented['image']
mask = augmented['mask']
if self.preprocessing:
preprocessed = self.preprocessing(image=img, mask=mask)
img = preprocessed['image']
mask = preprocessed['mask']
gs, gt_gs = [], []
if self.log_transform:
img, gs = transform_to_log(img)
gt, gt_gs = transform_to_log(gt)
gt, gt_gs = gt.mean(0).mean(0), gt_gs.mean(0).mean(0)
img = torch.tensor(img.transpose(2, 0, 1), dtype=torch.float32, device="cuda")
gt = torch.tensor(gt, dtype=torch.float32, device="cuda")
mask = np.array([[(np.array([1]) if pixel[0] > 128 else np.array([0])) for pixel in row] for row in mask])
mask = torch.tensor(mask.transpose(2, 0, 1), dtype=torch.float32, device="cuda")
return (img, gs), \
mask, \
(gt, gt_gs)
def __getitem__(self, idx):
patches_per_img = ceil(
1 / (self.patch_height_ratio * self.patch_width_ratio))
image_idx = int(idx / patches_per_img)
patch_idx = idx % patches_per_img
image_name = self.image_names[idx]
image = self.__load(image_name, self.images_path).astype(int)
image = self.__process_image(image, resize=False)
gt = self.gts[image_name - 1]
gt_rand = self.gts[(image_name + random.randint(10, 500)) %
len(self.image_names)]
image_patch = get_patch_with_index(image, patch_idx,
self.patch_height_ratio,
self.patch_width_ratio)
image = image_patch
gs, gt_gs = [], []
if self.log_transform:
image, gs = transform_to_log(image)
gt, gt_gs = transform_to_log(gt)
gt_rand, _ = transform_to_log(gt_rand)
augmented = self.transforms(image=image, mask=gs)
img, gs = augmented['image'], augmented['mask']
if self.preprocessing:
preprocessed = self.preprocessing(image=img, mask=gs)
img, gs = preprocessed['image'], preprocessed['mask']
if self.illumination_known:
shape = img.shape
gt_img = np.array([[np.array(gt) for i in range(shape[1])]
for j in range(shape[0])])
gt_rand_img = np.array(
[[np.array(gt_rand) for i in range(shape[1])]
for j in range(shape[0])])
if random.randint(0, 1) == 0:
img = np.dstack((img, gt_img, gt_rand_img))
else:
img = np.dstack((img, gt_rand_img, gt_img))
img = torch.tensor(img.transpose(2, 0, 1),
dtype=torch.float32,
device="cuda")
gt = torch.tensor(gt, dtype=torch.float32, device='cuda')
return img, gt, gs, gt_gs
def get_test_item(self, idx):
image_name = self.image_names[idx]
image, _, _ = load_img_and_gt_crf_dataset(image_name, self.path, self.folder, use_mask=self.use_mask,
use_corrected=self.use_corrected, dataset=self.dataset, load_any_mask=self.load_any_mask)
gs = []
if self.log_transform:
image, gs = transform_to_log(image)
aug = self.transforms(image=image, mask=gs)
img, gs = aug['image'], aug['mask'].squeeze()
else:
aug = self.transforms(image=image, mask=image)
img = aug['image']
if self.preprocessing:
if self.log_transform:
aug = self.preprocessing(image=img, mask=gs)
img, gs = aug['image'], aug['mask'].squeeze()
else:
aug = self.preprocessing(image=img, mask=image)
img = aug['image']
img = torch.tensor(img.transpose(2, 0, 1), dtype=torch.float32, device="cuda")
return image_name, img, gs
def get_train_item(self, idx):
image_name = self.image_names[idx]
image, gt, mask = load_img_and_gt_crf_dataset(image_name, self.path, self.folder, use_mask=self.use_mask,
use_corrected=self.use_corrected, dataset=self.dataset, load_any_mask=self.load_any_mask)
gt = cv2.cvtColor(gt, cv2.COLOR_RGB2LUV)
if not self.load_any_mask:
mask = gt
gs = []
if self.log_transform:
image, gs = transform_to_log(image)
gs = np.expand_dims(gs, axis=2)
inputs = {'image': image, 'mask': mask, 'image2': gs, 'image3': gt}
augmented = self.transforms(**inputs)
gs = augmented['image2']
gs = gs.squeeze()
else:
inputs = {'image': image, 'mask': mask, 'image3': gt}
augmented = self.transforms(**inputs)
img = augmented['image']
mask = augmented['mask']
gt = augmented['image3']
if self.preprocessing:
if self.log_transform:
preprocessed = self.preprocessing(image=img, mask=mask, image3=gt, image2=gs)
gs = preprocessed['image2']
else:
preprocessed = self.preprocessing(image=img, mask=mask, image2=gt)
img = preprocessed['image']
mask = preprocessed['mask']
gt = preprocessed['image3']
img = torch.tensor(img.transpose(2, 0, 1), dtype=torch.float32, device="cuda")
gt = torch.tensor(gt.transpose(2, 0, 1), dtype=torch.float32, device="cuda")
mask = np.expand_dims(np.where(mask > 128, np.ones(1), np.zeros(1))[:, :, 0], -1)
mask = torch.tensor(mask.transpose(2, 0, 1), dtype=torch.float32, device="cuda")
return (img, gs), \
mask, \
(gt)
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import utils.dataset_utils as du
from utils.HDRDataset import HDRPatchedDataset
from utils.load_mat import get_gt_for_image
from utils.transformation_utils import color_correct_fast, transform_from_log, adjust_gamma, transform_to_log
if __name__ == '__main__':
filename = './data/dataset_568_shi_gehler/cs/chroma/data/canon_dataset/568_dataset/png/IMG_0293.png'
im = cv2.imread(filename, cv2.COLOR_BGR2RGB)
im2 = adjust_gamma(im, gamma=2.2)
im3 = adjust_gamma(im, gamma=3.5)
im4 = adjust_gamma(im, gamma=5)
log, g = transform_to_log(im3)
im4 = transform_from_log(log, g).astype(int)
du.visualize(im, im2, im3, im4)
filename_cc = './data/dataset_hdr/HDR/cs/chroma/data/Nikon_D700/HDR_MATLAB_3x3/S0560.hdr'
im_cc = cv2.imread(filename_cc, cv2.IMREAD_ANYDEPTH)
im_cc2 = np.array(im_cc.astype(int), dtype=np.uint8)
im2 = adjust_gamma(im_cc2, gamma=2.2)
log, g = transform_to_log(im2)
im2 = transform_from_log(log, g).astype(int)
gt = get_gt_for_image('S0560_real_illum')
im_corrected = color_correct_fast(im_cc, gt)
im_corrected_2 = color_correct_fast(im2, (gt * 255).astype(int))
du.visualize(im_cc, im2, im_corrected_2, im_corrected)
log_rg, log_bg = cv2.split(log)
plt.scatter(log_rg, log_bg, s=2)