mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Enable force_multiple in LQ_dataset

Browse Source
This commit is contained in:
James Betker 2020-07-22 20:51:16 -06:00
parent 106b8da315
commit 8a0a1569f3
1 changed files with 9 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
codes/data/LQ_dataset.py
View File

@ -6,6 +6,7 @@ import data.util as util
import torchvision.transforms.functional as F import torchvision.transforms.functional as F
from PIL import Image from PIL import Image
import os.path as osp import os.path as osp
import cv2
class LQDataset(data.Dataset): class LQDataset(data.Dataset):
@ -22,6 +23,7 @@ class LQDataset(data.Dataset):
self.vertical_splits = self.opt['vertical_splits'] self.vertical_splits = self.opt['vertical_splits']
self.paths_LQ, self.paths_GT = None, None self.paths_LQ, self.paths_GT = None, None
self.LQ_env = None # environment for lmdb self.LQ_env = None # environment for lmdb
self.force_multiple = self.opt['force_multiple'] if 'force_multiple' in self.opt.keys() else 1
self.paths_LQ, self.sizes_LQ = util.get_image_paths(self.data_type, opt['dataroot_LQ']) self.paths_LQ, self.sizes_LQ = util.get_image_paths(self.data_type, opt['dataroot_LQ'])
self.paths_LQ = self.paths_LQ[self.start_at:] self.paths_LQ = self.paths_LQ[self.start_at:]
@ -48,6 +50,13 @@ class LQDataset(data.Dataset):
w_per_split = int(w / self.vertical_splits) w_per_split = int(w / self.vertical_splits)
left = w_per_split * split_index left = w_per_split * split_index
img_LQ = F.crop(img_LQ, 0, left, h, w_per_split) img_LQ = F.crop(img_LQ, 0, left, h, w_per_split)
# Enforce force_resize constraints.
h, w = img_LQ.size
if h % self.force_multiple != 0 or w % self.force_multiple != 0:
h, w = (w - w % self.force_multiple), (h - h % self.force_multiple)
img_LQ = img_LQ.resize((w, h))
img_LQ = F.to_tensor(img_LQ) img_LQ = F.to_tensor(img_LQ)
img_name = osp.splitext(osp.basename(LQ_path))[0] img_name = osp.splitext(osp.basename(LQ_path))[0]