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)