def __getitem__(self, idx):
flow_dir = self.data_items[idx][0]
video_class_no = self.data_items[idx][1]
if self.config.get_mask:
mask_dir = self.data_items[idx][2]
if self.isTest:
output_dirs = self.data_items[idx][-1]
flow_set = []
mask_set = []
flow_mask_cat_set = []
flow_masked_set = []
if self.config.MASK_MODE == 'bbox':
tmp_bbox = im.random_bbox(self.config)
tmp_mask = im.bbox2mask(self.config, tmp_bbox)
tmp_mask = tmp_mask[0, 0, :, :]
fix_mask = np.expand_dims(tmp_mask, axis=2)
elif self.config.MASK_MODE == 'mid-bbox':
tmp_mask = im.mid_bbox_mask(self.config)
tmp_mask = tmp_mask[0, 0, :, :]
fix_mask = np.expand_dims(tmp_mask, axis=2)
for i in range(11):
tmp_flow = cvb.read_flow(flow_dir[i])
if self.config.get_mask:
tmp_mask = cv2.imread(mask_dir[i],
cv2.IMREAD_UNCHANGED)
tmp_mask = self._mask_tf(tmp_mask)
else:
if self.config.FIX_MASK:
tmp_mask = fix_mask.copy()
else:
tmp_bbox = im.random_bbox(self.config)
tmp_mask = im.bbox2mask(self.config, tmp_bbox)
tmp_mask = tmp_mask[0, 0, :, :]
tmp_mask = np.expand_dims(tmp_mask, axis=2)
tmp_flow = self._flow_tf(tmp_flow)
tmp_flow_masked = tmp_flow * (1. - tmp_mask)
if self.config.INITIAL_HOLE:
tmp_flow_resized = cv2.resize(tmp_flow, (self.size[1] // 2, self.size[0] // 2))
tmp_mask_resized = cv2.resize(tmp_mask, (self.size[1] // 2, self.size[0] // 2), cv2.INTER_NEAREST)
tmp_flow_masked_small = tmp_flow_resized
tmp_flow_masked_small[:, :, 0] = rf.regionfill(tmp_flow_resized[:, :, 0], tmp_mask_resized)
tmp_flow_masked_small[:, :, 1] = rf.regionfill(tmp_flow_resized[:, :, 1], tmp_mask_resized)
tmp_flow_masked = tmp_flow_masked + \
tmp_mask * cv2.resize(tmp_flow_masked_small, (self.size[1], self.size[0]))
flow_masked_set.append(tmp_flow_masked)
flow_set.append(tmp_flow)
mask_set.append(tmp_mask)
mask_set.append(tmp_mask)
tmp_flow_mask_cat = np.concatenate((tmp_flow_masked, tmp_mask), axis=2)
flow_mask_cat_set.append(tmp_flow_mask_cat)
flow_mask_cat = np.concatenate(flow_mask_cat_set, axis=2)
flow_masked = np.concatenate(flow_masked_set, axis=2)
gt_flow = np.concatenate(flow_set, axis=2)
mask = np.concatenate(mask_set, axis=2)
flow_mask_cat = torch.from_numpy(flow_mask_cat).permute(2, 0, 1).contiguous().float()
flow_masked = torch.from_numpy(flow_masked).permute(2, 0, 1).contiguous().float()
gt_flow = torch.from_numpy(gt_flow).permute(2, 0, 1).contiguous().float()
mask = torch.from_numpy(mask).permute(2, 0, 1).contiguous().float()
if self.isTest:
return flow_mask_cat, flow_masked, gt_flow, mask, output_dirs
return flow_mask_cat, flow_masked, gt_flow, mask
def __getitem__(self, idx):
flow_dir = self.data_items[idx][0]
video_class_no = self.data_items[idx][1]
gt_dir = self.data_items[idx][2]
if self.config.get_mask:
mask_dirs = self.data_items[idx][3]
if self.isTest:
output_dirs = self.data_items[idx][-1]
mask_set = []
flow_mask_cat_set = []
flow_masked_set = []
gt_flow_set = []
gt_dir_set = [gt_dir[5], gt_dir[16]]
for p in gt_dir_set:
tmp_flow = cvb.read_flow(p)
tmp_flow = self._flow_tf(tmp_flow)
gt_flow_set.append(tmp_flow)
if self.config.MASK_MODE == 'bbox':
tmp_bbox = im.random_bbox(self.config)
tmp_mask = im.bbox2mask(self.config, tmp_bbox)
tmp_mask = tmp_mask[0, 0, :, :]
fix_mask = np.expand_dims(tmp_mask, axis=2)
elif self.config.MASK_MODE == 'mid-bbox':
tmp_mask = im.mid_bbox_mask(self.config)
tmp_mask = tmp_mask[0, 0, :, :]
fix_mask = np.expand_dims(tmp_mask, axis=2)
f_flow_dir = flow_dir[:11]
r_flow_dir = flow_dir[11:]
for i in range(11):
tmp_flow = cvb.read_flow(f_flow_dir[i])
if self.config.get_mask:
tmp_mask = cv2.imread(mask_dirs[i], cv2.IMREAD_UNCHANGED)
tmp_mask = self._mask_tf(tmp_mask)
else:
if self.config.FIX_MASK:
tmp_mask = fix_mask.copy()
else:
tmp_bbox = im.random_bbox(self.config)
tmp_mask = im.bbox2mask(self.config, tmp_bbox)
tmp_mask = tmp_mask[0, 0, :, :]
tmp_mask = np.expand_dims(tmp_mask, axis=2)
tmp_flow = self._flow_tf(tmp_flow)
tmp_flow_masked = tmp_flow
flow_masked_set.append(tmp_flow_masked)
mask_set.append(tmp_mask)
mask_set.append(tmp_mask)
tmp_flow_mask_cat = np.concatenate((tmp_flow_masked, tmp_mask),
axis=2)
flow_mask_cat_set.append(tmp_flow_mask_cat)
for i in range(11):
tmp_flow = cvb.read_flow(r_flow_dir[i])
tmp_flow = self._flow_tf(tmp_flow)
if self.config.get_mask:
tmp_mask = cv2.imread(mask_dirs[i + 11], cv2.IMREAD_UNCHANGED)
tmp_mask = self._mask_tf(tmp_mask)
else:
if self.config.FIX_MASK:
tmp_mask = fix_mask.copy()
else:
tmp_bbox = im.random_bbox(self.config)
tmp_mask = im.bbox2mask(self.config, tmp_bbox)
tmp_mask = tmp_mask[0, 0, :, :]
tmp_mask = np.expand_dims(tmp_mask, axis=2)
tmp_flow_masked = tmp_flow
flow_masked_set.append(tmp_flow_masked)
mask_set.append(tmp_mask)
mask_set.append(tmp_mask)
tmp_flow_mask_cat = np.concatenate((tmp_flow_masked, tmp_mask),
axis=2)
flow_mask_cat_set.append(tmp_flow_mask_cat)
flow_mask_cat = np.concatenate(flow_mask_cat_set, axis=2)
flow_masked = np.concatenate(flow_masked_set, axis=2)
gt_flow = np.concatenate(gt_flow_set, axis=2)
mask = np.concatenate(mask_set, axis=2)
flow_mask_cat = torch.from_numpy(flow_mask_cat).permute(
2, 0, 1).contiguous().float()
flow_masked = torch.from_numpy(flow_masked).permute(
2, 0, 1).contiguous().float()
gt_flow = torch.from_numpy(gt_flow).permute(2, 0,
1).contiguous().float()
mask = torch.from_numpy(mask).permute(2, 0, 1).contiguous().float()
if self.isTest:
return flow_mask_cat, flow_masked, gt_flow, mask, output_dirs
else:
return flow_mask_cat, flow_masked, gt_flow, mask