DL-Art-School/codes/data/multi_frame_dataset.py
James Betker 11155aead4 Directly use dataset keys
This has been a long time coming. Cleans up messy "GT" nomenclature and simplifies ExtensibleTraner.feed_data
2020-12-04 20:14:53 -07:00

107 lines
4.3 KiB
Python

from data.base_unsupervised_image_dataset import BaseUnsupervisedImageDataset
import numpy as np
import torch
from bisect import bisect_left
import os.path as osp
class MultiFrameDataset(BaseUnsupervisedImageDataset):
def __init__(self, opt):
super(MultiFrameDataset, self).__init__(opt)
self.num_frames = opt['num_frames']
def chunk_name(self, i):
return osp.basename(self.chunks[i].path)
def get_sequential_image_paths_from(self, chunk_index, chunk_offset):
im_name = self.chunk_name(chunk_index)
source_name = im_name[:-12]
# Search backwards for the frames needed. We are assuming that every video in the dataset has at least frames_needed frames.
frames_needed = self.num_frames
search_idx = chunk_index
while frames_needed > 0:
frames_needed -= 1
search_idx -= 1
if source_name not in self.chunk_name(search_idx):
search_idx += 1
break
# Now build num_frames starting from search_idx.
hqs, refs, masks, centers = [], [], [], []
for i in range(self.num_frames):
idx = search_idx + i
if idx < 0 or idx >= len(self.chunks) or chunk_offset < 0 or chunk_offset >= len(self.chunks[idx]):
print("Chunk reference indexing failed for %s." % (im_name,), search_idx, i, chunk_offset, self.num_frames)
h, r, c, m, p = self.chunks[search_idx + i][chunk_offset]
hqs.append(h)
refs.append(r)
masks.append(m)
centers.append(c)
path = p
return hqs, refs, masks, centers, path
def __getitem__(self, item):
chunk_ind = bisect_left(self.starting_indices, item)
chunk_ind = chunk_ind if chunk_ind < len(self.starting_indices) and self.starting_indices[chunk_ind] == item else chunk_ind-1
hqs, refs, masks, centers, path = self.get_sequential_image_paths_from(chunk_ind, item-self.starting_indices[chunk_ind])
hs, hrs, hms, hcs = self.resize_hq(hqs, refs, masks, centers)
ls, lrs, lms, lcs = self.synthesize_lq(hs, hrs, hms, hcs)
# Convert to torch tensor
hq = torch.from_numpy(np.ascontiguousarray(np.transpose(np.stack(hs), (0, 3, 1, 2)))).float()
hq_ref = torch.from_numpy(np.ascontiguousarray(np.transpose(np.stack(hrs), (0, 3, 1, 2)))).float()
hq_mask = torch.from_numpy(np.ascontiguousarray(np.stack(hms))).unsqueeze(dim=1)
hq_ref = torch.cat([hq_ref, hq_mask], dim=1)
lq = torch.from_numpy(np.ascontiguousarray(np.transpose(np.stack(ls), (0, 3, 1, 2)))).float()
lq_ref = torch.from_numpy(np.ascontiguousarray(np.transpose(np.stack(lrs), (0, 3, 1, 2)))).float()
lq_mask = torch.from_numpy(np.ascontiguousarray(np.stack(lms))).unsqueeze(dim=1)
lq_ref = torch.cat([lq_ref, lq_mask], dim=1)
return {'GT_path': path, 'lq': lq, 'hq': hq, 'gt_fullsize_ref': hq_ref, 'lq_fullsize_ref': lq_ref,
'lq_center': torch.tensor(lcs, dtype=torch.long), 'gt_center': torch.tensor(hcs, dtype=torch.long)}
if __name__ == '__main__':
opt = {
'name': 'amalgam',
'paths': ['F:\\4k6k\\datasets\\images\\fullvideo\\ge_fv_256_tiled'],
'weights': [1],
'target_size': 128,
'force_multiple': 32,
'scale': 2,
'eval': False,
'fixed_corruptions': ['jpeg-medium'],
'random_corruptions': [],
'num_corrupts_per_image': 0,
'num_frames': 10
}
ds = MultiFrameDataset(opt)
import os
os.makedirs("debug", exist_ok=True)
bs = 0
batch = None
for i in range(len(ds)):
import random
k = 'lq'
element = ds[random.randint(0,len(ds))]
base_file = osp.basename(element["GT_path"])
o = element[k].unsqueeze(0)
if bs < 32:
if batch is None:
batch = o
else:
batch = torch.cat([batch, o], dim=0)
bs += 1
continue
if 'path' not in k and 'center' not in k:
b, fr, f, h, w = batch.shape
for j in range(fr):
import torchvision
base=osp.basename(base_file)
torchvision.utils.save_image(batch[:, j], "debug/%i_%s_%i__%s.png" % (i, k, j, base))
bs = 0
batch = None