Class ValidateData

core/datasets.py:207–235

        return len(self.images_list)
        
class ValidateData(data.Dataset):
    def __init__(self, args, split):
        self.root = os.path.join(args.validate_data_dir, split)
        if not os.path.exists(self.root):
            raise colored('[Error: ]', 'red') + 'file path {:s} is not exist!'.format(self.root)

        self.images_dir = os.path.join(self.root, 'images')
        self.flows_dir = os.path.join(self.root, 'flows')
        self.masks_dir = os.path.join(self.root, 'masks')


    def __getitem__(self, index):
        im1 = cv2.imread(os.path.join(self.images_dir, 'image_{:06d}_0.png'.format(index)))
        im2 = cv2.imread(os.path.join(self.images_dir, 'image_{:06d}_1.png'.format(index)))
        flow = readFlow(os.path.join(self.flows_dir, 'flow_{:06d}.flo'.format(index)))

        H, W = im1.shape[:2]

        im1 = torch.from_numpy(im1).permute([2, 0, 1])
        im2 = torch.from_numpy(im2).permute([2, 0, 1])
        flow = torch.from_numpy(flow).permute([2, 0, 1])
   
        valid_mask = np.ones([H, W]).astype(bool)
        if os.path.exists(self.masks_dir):
            valid_mask = cv2.imread(os.path.join(self.masks_dir, 'mask_{:06d}.png'.format(index)), cv2.IMREAD_UNCHANGED).astype(bool)
        return im1, im2, flow, valid_mask

    def __len__(self):
        return len(os.listdir(self.flows_dir))

def fetch_dataloader(args, split='train'):
    print(colored('[Data]: ', 'yellow') + 'Loading MegaDepth CAPS')