def __getitem__(self, index):
folder = self.folders[index]
frame_files = glob(
os.path.join(folder, 'frame/' + self.opt.image_pattern))
frame_count = len(frame_files)
frames = torch.zeros(
[frame_count, 3, self.opt.content_height, self.opt.content_width])
optflow = torch.zeros([
frame_count - 1, 2, self.opt.content_height, self.opt.content_width
])
downflow = torch.zeros([
frame_count - 1, 2, self.opt.content_height // 4,
self.opt.content_width // 4
])
mask = torch.zeros([
frame_count - 1, 1, self.opt.content_height, self.opt.content_width
])
downmask = torch.zeros([
frame_count - 1, 1, self.opt.content_height // 4,
self.opt.content_width // 4
])
for i in range(frame_count):
imageName = os.path.join(folder, 'frame/' + str(i) + '.jpg')
# img = Image.open(imageName).convert('L')
img = Image.open(imageName)
img = transforms.ToTensor()(img)
img = transforms.Lambda(lambda x: x.mul(255))(img)
frames[i] = img
for i in range(frame_count - 1):
# flow: (h, w, c) -> (c, h, w)
flowName = os.path.join(folder, 'flow/' + str(i) + '.flo')
flow = torch.Tensor(readFlow(flowName)).permute(2, 0, 1)
optflow[i] = flow
downflow[i] = nn.functional.interpolate(
flow.unsqueeze(0),
scale_factor=(0.25, 0.25),
mode='bilinear',
align_corners=True).squeeze(0)
for i in range(frame_count - 1):
maskName = os.path.join(folder, 'mask/' + str(i))
cmask = torch.Tensor(
cv2.imread(maskName, cv2.IMREAD_GRAYSCALE) / 255.).unsqueeze(0)
mask[i] = cmask
downmask[i] = nn.functional.interpolate(
cmask.unsqueeze(0),
scale_factor=(0.25, 0.25),
mode='bilinear',
align_corners=True).squeeze(0)
return {
'frames': frames,
'optflow': optflow,
'downflow': downflow,
'mask': mask,
'downmask': downmask,
}
def __getitem__(self, idx):
HR = []
FLOW = []
for i in range(self.input_frame):
img_hr = Image.open(self.dir + '/sequences/' + '/' +
self.train_list[idx] + '/im' + str(i + 1) +
'.png') #'/sequences/' +
img_hr = np.array(img_hr, dtype=np.float32) / 255.0
img_hr = img_hr.transpose(2, 0, 1)
if i != (self.input_frame - 1):
flo = readFlow(
'./../../../../../../data2/anchen/data/opticflow/' + '/' +
self.train_list[idx] + '/im' + str(i + 1) + '.flo')
flo = flo.transpose(2, 0, 1)
FLOW.append(flo)
HR.append(img_hr)
HR = np.stack(HR, 1)
FLOW.append(flo) # for augumentation
FLOW = np.stack(FLOW, 1)
combine_data = np.concatenate((HR, FLOW), axis=0)
combine_data = random_crop(combine_data, self.patch_size)
#combine_data = self.tranform(combine_data)
combine_data = torch.from_numpy(np.ascontiguousarray(combine_data))
HR = combine_data[0:3, :, :, :]
#print(np.shape(HR))
FLOW = combine_data[3:5, 0:self.input_frame - 1, :, :]
#print(np.shape(FLOW))
return HR, FLOW
def read_gen(file_name):
ext = splitext(file_name)[-1]
if ext == '.png' or ext == '.jpeg' or ext == '.ppm' or ext == '.jpg':
im = imread(file_name)
h = im.shape[0]
w = im.shape[1]
if ((h % 64 != 0) or (w % 64 != 0)):
im = im_rescale(im)
if im.shape[2] > 3:
return im[:, :, :3]
else:
return im
elif ext == '.bin' or ext == '.raw':
return np.load(file_name)
elif ext == '.flo':
flo = flow_utils.readFlow(file_name).astype(np.float32)
h = flo.shape[0]
w = flo.shape[1]
h_new = 64 * int(round(h / 64))
w_new = 64 * int(round(w / 64))
if ((h % 64 != 0) or (w % 64 != 0)):
flo = np.resize(flo, (h_new, w_new, 2))
#print(flo.shape)
return flo
return []
def read_gen(file_name):
ext = splitext(file_name)[-1]
if ext == '.png' or ext == '.jpeg' or ext == '.ppm' or ext == '.jpg':
im = imread(file_name)
if im.shape[2] > 3:
return im[:,:,:3]
else:
return im
elif ext == '.bin' or ext == '.raw':
return np.load(file_name)
elif ext == '.flo':
return readFlow(file_name).astype(np.float32)
return []
def read_gen(file_name):
ext = splitext(file_name)[-1]
if ext == '.png' or ext == '.jpeg' or ext == '.ppm' or ext == '.jpg':
im = imread(file_name)
if len(im.shape) < 3:
im = im[:, :, np.newaxis]
if im.shape[2] > 3:
return im[:, :, :3]
else:
return im
elif ext == '.bin' or ext == '.raw':
return np.load(file_name)
elif ext == '.flo':
return flow_utils.readFlow(file_name).astype(np.float32)
elif ext == '.pfm':
im = readPFM(file_name)
if len(im.shape) >= 3:
return im[:, :, :2]
else:
return im[:, :, np.newaxis]
return []
import numpy as np
import os
import flow_utils
new_width = 224
new_height = 224
num = 25
path = r"./flow_v_CricketShot_g04_c01_resized"
filenames = os.listdir(path)
filenames.sort(key=lambda x: int(x[:-4]))
flows = []
normed_flows = []
for i in range(len(filenames) / num):
my_filename = filenames[num * i:num * (i + 1)]
for filename in list(my_filename):
flow = flow_utils.readFlow(os.path.join(path, filename))
x = flow[:, :, 0]
y = flow[:, :, 1]
normed_x = ((x - x.min()) / (x.max() - x.min()) - 0.5) * 2
normed_y = ((y - y.min()) / (y.max() - y.min()) - 0.5) * 2
width, height, channel = flow.shape
normed_flow = np.zeros((width, height, 2))
normed_flow[:, :, 0] = normed_x
normed_flow[:, :, 1] = normed_y
left = (width - new_width) / 2
right = (width + new_width) / 2
top = (height - new_height) / 2
bottom = (height + new_height) / 2
croped_flow = flow[left:right, top:bottom, :]
croped_normed_flow = normed_flow[left:right, top:bottom, :]
flows.append(np.array(croped_flow))