From e27a49454e48c4239bfe296b477b9c625620fabd Mon Sep 17 00:00:00 2001 From: James Betker Date: Wed, 27 May 2020 08:04:35 -0600 Subject: [PATCH] Enable vertical splitting on inference images to support very large resolutions. --- codes/data/LQ_dataset.py | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/codes/data/LQ_dataset.py b/codes/data/LQ_dataset.py index 54bc379f..35c20a06 100644 --- a/codes/data/LQ_dataset.py +++ b/codes/data/LQ_dataset.py @@ -19,6 +19,7 @@ class LQDataset(data.Dataset): self.start_at = self.opt['start_at'] else: self.start_at = 0 + self.vertical_splits = self.opt['vertical_splits'] self.paths_LQ, self.paths_GT = None, None self.LQ_env = None # environment for lmdb @@ -33,21 +34,29 @@ class LQDataset(data.Dataset): def __getitem__(self, index): if self.data_type == 'lmdb' and self.LQ_env is None: self._init_lmdb() - actual_index = index # int(index / 2) - is_left = (index % 2) == 0 + if self.vertical_splits > 0: + actual_index = int(index / self.vertical_splits) + else: + actual_index = index # get LQ image LQ_path = self.paths_LQ[actual_index] img_LQ = Image.open(LQ_path) - left = 0 if is_left else 1920 - # crop input if needed. - #img_LQ = F.crop(img_LQ, 5, left + 5, 1900, 1900) + if self.vertical_splits > 0: + w, h = img_LQ.size + split_index = (index % self.vertical_splits) + w_per_split = int(w / self.vertical_splits) + left = w_per_split * split_index + img_LQ = F.crop(img_LQ, 0, left, h, w_per_split) img_LQ = F.to_tensor(img_LQ) img_name = osp.splitext(osp.basename(LQ_path))[0] - LQ_path = LQ_path.replace(img_name, img_name + "_%i" % (index % 2)) + LQ_path = LQ_path.replace(img_name, img_name + "_%i" % (index % self.vertical_splits)) return {'LQ': img_LQ, 'LQ_path': LQ_path} def __len__(self): - return len(self.paths_LQ) # * 2 + if self.vertical_splits > 0: + return len(self.paths_LQ) * self.vertical_splits + else: + return len(self.paths_LQ)