diff --git a/IGEV-Stereo/dataset_utils/create_crestereo_subsets.py b/IGEV-Stereo/dataset_utils/create_crestereo_subsets.py new file mode 100644 index 0000000..21ea954 --- /dev/null +++ b/IGEV-Stereo/dataset_utils/create_crestereo_subsets.py @@ -0,0 +1,49 @@ +import os +import numpy as np +import torch +import torch.utils.data as data +import torch.nn.functional as F +import logging +import os +import re +import copy +import math +import random +from pathlib import Path +from glob import glob +import os.path as osp + +from core.utils import frame_utils + +class CREStereo(): + def __init__(self, aug_params=None, root='/data/CREStereo'): + self.root = root + assert os.path.exists(root) + + disp_list = self.selector('_left.disp.png') + image1_list = self.selector('_left.jpg') + image2_list = self.selector('_right.jpg') + + assert len(image1_list) == len(image2_list) == len(disp_list) > 0 + for img1, img2, disp in zip(image1_list, image2_list, disp_list): + # if random.randint(1, 20000) != 1: + # continue + self.image_list += [[img1, img2]] + self.disparity_list += [disp] + + def selector(self, suffix): + files = list(glob(os.path.join(self.root, f"hole/*{suffix}"))) + files += list(glob(os.path.join(self.root, f"shapenet/*{suffix}"))) + files += list(glob(os.path.join(self.root, f"tree/*{suffix}"))) + files += list(glob(os.path.join(self.root, f"reflective/*{suffix}"))) + return sorted(files) + + def disparity_distribution(self): + disp_lists = self.selector('_left.disp.png') + disparities = [] + for filename in disp_lists: + disp_gt = frame_utils.readDispCREStereo(filename) + print(disp_gt.shape) + + +CREStereo.disparity_distribution() \ No newline at end of file