core/stereo_datasets.py

# Data loading based on https://github.com/NVIDIA/flownet2-pytorch

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 plane
from core.utils import frame_utils
from core.utils.ddp import get_loader
from core.utils.augmentor import FlowAugmentor, SparseFlowAugmentor
DATASET_ROOT = os.getenv('DATASET_ROOT')


class StereoDataset(data.Dataset):
    def __init__(self, aug_params=None, sparse=False, reader=None, args=None):
        self.augmentor = None
        self.sparse = sparse
        self.img_pad = aug_params.pop("img_pad", None) if aug_params is not None else None
        if aug_params is not None and "crop_size" in aug_params:
            if sparse:
                self.augmentor = SparseFlowAugmentor(**aug_params)
            else:
                self.augmentor = FlowAugmentor(**aug_params)

        if reader is None:
            self.disparity_reader = frame_utils.read_gen
        else:
            self.disparity_reader = reader        

        # if args is not None:
        #     # self.plane = args.plane_datset
        #     self.slant = args.slant 
        #     self.slant_norm = args.slant_norm
        # else:
        #     # self.plane = False
        #     self.slant = None 
        #     self.slant_norm = False

        self.is_test = args.is_test if hasattr(args, "is_test") and args.is_test else False
        self.init_seed = False
        self.flow_list = []
        self.disparity_list = []
        self.image_list = []
        self.extra_info = {}

    def __getitem__(self, index):

        if self.is_test:
            img1 = frame_utils.read_gen(self.image_list[index][0])
            img2 = frame_utils.read_gen(self.image_list[index][1])
            img1 = np.array(img1).astype(np.uint8)[..., :3]
            img2 = np.array(img2).astype(np.uint8)[..., :3]
            img1 = torch.from_numpy(img1).permute(2, 0, 1).float()
            img2 = torch.from_numpy(img2).permute(2, 0, 1).float()
            return self.image_list[index] + [self.disparity_list[index]], \
                   img1, img2, torch.zeros_like(torch.zeros_like(img1))[:1], torch.ones_like(torch.zeros_like(img1))[:1]

        if not self.init_seed:
            worker_info = torch.utils.data.get_worker_info()
            if worker_info is not None:
                torch.manual_seed(worker_info.id)
                np.random.seed(worker_info.id)
                random.seed(worker_info.id)
                self.init_seed = True

        try:
            index = index % len(self.image_list)
            intrinsic = self.extra_info["intrinsics"][index] if "intrinsics" in self.extra_info else None
            disp = self.disparity_reader(self.disparity_list[index])
            if isinstance(disp, tuple):
                disp, valid = disp
            else:
                valid = disp < 512

            img1 = frame_utils.read_gen(self.image_list[index][0])
            img2 = frame_utils.read_gen(self.image_list[index][1])
        
            img1 = np.array(img1).astype(np.uint8)
            img2 = np.array(img2).astype(np.uint8)

            disp = np.array(disp).astype(np.float32)
            flow = np.stack([-disp, np.zeros_like(disp)], axis=-1)

        except Exception as err:
            raise Exception(err, "{}, {}, {}".format(self.image_list[index][0], 
                                                     self.image_list[index][1], 
                                                     self.disparity_list[index] ))

        # grayscale images
        if len(img1.shape) == 2:
            img1 = np.tile(img1[...,None], (1, 1, 3))
            img2 = np.tile(img2[...,None], (1, 1, 3))
        else:
            img1 = img1[..., :3]
            img2 = img2[..., :3]

        if self.augmentor is not None:
            if self.sparse:
                img1, img2, flow, valid, intrinsic = self.augmentor(img1, img2, flow, valid, intrinsic)
            else:
                img1, img2, flow, intrinsic = self.augmentor(img1, img2, flow, intrinsic)

        try:
            img1 = torch.from_numpy(img1).permute(2, 0, 1).float()
            img2 = torch.from_numpy(img2).permute(2, 0, 1).float()
            flow = torch.from_numpy(flow).permute(2, 0, 1).float()
            intrinsic = torch.from_numpy(np.array(intrinsic)).float() if intrinsic is not None else torch.from_numpy(np.eye(3)).float()
        except Exception as err:
            raise Exception(err, "{}, {}, {}".format(self.image_list[index][0], 
                                                     self.image_list[index][1], 
                                                     self.disparity_list[index]),
                            "{}, {}, {}".format(img1.shape, img2.shape, flow.shape), )

        if self.sparse:
            valid = torch.from_numpy(valid)
        else:
            valid = (flow[0].abs() < 512) & (flow[1].abs() < 512)

        if self.img_pad is not None:
            padH, padW = self.img_pad
            img1 = F.pad(img1, [padW]*2 + [padH]*2)
            img2 = F.pad(img2, [padW]*2 + [padH]*2)

        flow = flow[:1]

        return self.image_list[index] + [self.disparity_list[index]], \
               img1, img2, flow, valid.float(), intrinsic


    def __mul__(self, v):
        copy_of_self = copy.deepcopy(self)
        copy_of_self.flow_list = v * copy_of_self.flow_list
        copy_of_self.image_list = v * copy_of_self.image_list
        copy_of_self.disparity_list = v * copy_of_self.disparity_list
        if isinstance(copy_of_self.extra_info, list):
            copy_of_self.extra_info = v * copy_of_self.extra_info
        else:
            copy_of_self.extra_info = {key: val*v for key, val in copy_of_self.extra_info.items()}
        return copy_of_self
        
    def __len__(self):
        return len(self.image_list)


class SceneFlowDatasets(StereoDataset):
    def __init__(self, aug_params=None, root='', dstype='frames_cleanpass', 
                 things_test=False, caching=False, args=None, eval=False):
        super(SceneFlowDatasets, self).__init__(aug_params, args=args)
        self.eval = args.eval if args is not None else eval
        self.root = root if len(root)>0 else DATASET_ROOT
        self.dstype = dstype
        self.caching = caching
        self.extra_info["intrinsics"] = []
        assert os.path.exists(self.root), "check the existence: {}".format(self.root)

        if things_test:
            self._add_things("TEST")
        else:
            self._add_things("TRAIN")
            self._add_monkaa()
            self._add_driving()

    def _add_things(self, split='TRAIN'):
        """ Add FlyingThings3D data """

        original_length = len(self.disparity_list)
        cache_file = osp.join(self.root, 'flying3d'+"-"+self.dstype+"-"+split+".npz")
        if self.caching and os.path.exists(cache_file):
            cache = np.load(cache_file)
            root = cache["root"]
            left_images = cache["left_images"]
            right_images = cache["right_images"]
            disparity_images = cache["disparity_images"]
        else :
            root = osp.join(self.root, 'flying3d')
            left_images = sorted( glob(osp.join(root, self.dstype, split, '*/*/left/*.png')) )
            right_images = [ im.replace('left', 'right') for im in left_images ]
            disparity_images = [ im.replace(self.dstype, 'disparity').replace('.png', '.pfm') for im in left_images ]
            if self.caching :
                np.savez(cache_file, 
                        root=root,
                        left_images=left_images, 
                        right_images=right_images, 
                        disparity_images=disparity_images)

        # Choose a random subset of 400 images for validation
        state = np.random.get_state()
        np.random.seed(1000)
        if not self.eval:
            val_idxs = set(np.random.permutation(len(left_images))[:400])
        else:
            val_idxs = set(np.random.permutation(len(left_images)))
        np.random.set_state(state)

        for idx, (img1, img2, disp) in enumerate(zip(left_images, right_images, disparity_images)):
            if (split == 'TEST' and idx in val_idxs) or split == 'TRAIN':
                self.image_list += [ [img1, img2] ]
                self.disparity_list += [ disp ]
                self.extra_info["intrinsics"] += [ [1050, 1050, 479.5, 269.5] ]
        
        logging.info(f"Added {len(self.disparity_list) - original_length} from FlyingThings {self.dstype}")

    def _add_monkaa(self):
        """ Add FlyingThings3D data """

        original_length = len(self.disparity_list)
        root = osp.join(self.root, 'monkaa')
        left_images = sorted( glob(osp.join(root, self.dstype, '*/left/*.png')) )
        right_images = [ image_file.replace('left', 'right') for image_file in left_images ]
        disparity_images = [ im.replace(self.dstype, 'disparity').replace('.png', '.pfm') for im in left_images ]

        for img1, img2, disp in zip(left_images, right_images, disparity_images):
            self.image_list += [ [img1, img2] ]
            self.disparity_list += [ disp ]
            self.extra_info["intrinsics"] += [ [1050, 1050, 479.5, 269.5] ]
        logging.info(f"Added {len(self.disparity_list) - original_length} from Monkaa {self.dstype}")


    def _add_driving(self):
        """ Add FlyingThings3D data """

        original_length = len(self.disparity_list)
        root = osp.join(self.root, 'driving')
        left_images = sorted( glob(osp.join(root, self.dstype, '*/*/*/left/*.png')) )
        right_images = [ image_file.replace('left', 'right') for image_file in left_images ]
        disparity_images = [ im.replace(self.dstype, 'disparity').replace('.png', '.pfm') for im in left_images ]

        for img1, img2, disp in zip(left_images, right_images, disparity_images):
            self.image_list += [ [img1, img2] ]
            self.disparity_list += [ disp ]
            if img1.find("15mm_focallength") != -1:
                self.extra_info["intrinsics"] += [ [450, 450, 479.5, 269.5] ]
            elif img1.find("35mm_focallength") != -1:
                self.extra_info["intrinsics"] += [ [1050, 1050, 479.5, 269.5] ]
            else:
                raise Exception(f"Unknown intrinsics: {im1}")
        logging.info(f"Added {len(self.disparity_list) - original_length} from Driving {self.dstype}")


class ETH3D(StereoDataset):
    def __init__(self, aug_params=None, root='datasets/ETH3D', split='training', args=None):
        super(ETH3D, self).__init__(aug_params, sparse=True, args=args)
        root = root if len(root)>0 else DATASET_ROOT
        assert os.path.exists(root), "check the existence: {}".format(root)

        image1_list = sorted( glob(osp.join(root, f'two_view_{split}/*/im0.png')) )
        image2_list = sorted( glob(osp.join(root, f'two_view_{split}/*/im1.png')) )
        disp_list = sorted( glob(osp.join(root, 'two_view_training/*/disp0GT.pfm')) ) if split == 'training' else [osp.join(root, 'two_view_training_gt/playground_1l/disp0GT.pfm')]*len(image1_list)

        for img1, img2, disp in zip(image1_list, image2_list, disp_list):
            self.image_list += [ [img1, img2] ]
            self.disparity_list += [ disp ]

class SintelStereo(StereoDataset):
    def __init__(self, aug_params=None, root='datasets/SintelStereo', args=None):
        super().__init__(aug_params, sparse=True, reader=frame_utils.readDispSintelStereo, args=args)
        root = root if len(root)>0 else DATASET_ROOT

        image1_list = sorted( glob(osp.join(root, 'training/*_left/*/frame_*.png')) )
        image2_list = sorted( glob(osp.join(root, 'training/*_right/*/frame_*.png')) )
        disp_list = sorted( glob(osp.join(root, 'training/disparities/*/frame_*.png')) ) * 2

        for img1, img2, disp in zip(image1_list, image2_list, disp_list):
            assert img1.split('/')[-2:] == disp.split('/')[-2:]
            self.image_list += [ [img1, img2] ]
            self.disparity_list += [ disp ]

class FallingThings(StereoDataset):
    def __init__(self, aug_params=None, root='datasets/FallingThings', args=None):
        super().__init__(aug_params, reader=frame_utils.readDispFallingThings, args=args)
        root = root if len(root)>0 else DATASET_ROOT
        assert os.path.exists(root)

        with open(os.path.join(root, 'filenames.txt'), 'r') as f:
            filenames = sorted(f.read().splitlines())

        image1_list = [osp.join(root, e) for e in filenames]
        image2_list = [osp.join(root, e.replace('left.jpg', 'right.jpg')) for e in filenames]
        disp_list = [osp.join(root, e.replace('left.jpg', 'left.depth.png')) for e in filenames]

        for img1, img2, disp in zip(image1_list, image2_list, disp_list):
            self.image_list += [ [img1, img2] ]
            self.disparity_list += [ disp ]

class TartanAir(StereoDataset):
    def __init__(self, aug_params=None, root='datasets', keywords=[]):
        super().__init__(aug_params, reader=frame_utils.readDispTartanAir)
        root = root if len(root)>0 else DATASET_ROOT
        assert os.path.exists(root)

        with open(os.path.join(root, 'tartanair_filenames.txt'), 'r') as f:
            filenames = sorted(list(filter(lambda s: 'seasonsforest_winter/Easy' not in s, f.read().splitlines())))
            for kw in keywords:
                filenames = sorted(list(filter(lambda s: kw in s.lower(), filenames)))

        image1_list = [osp.join(root, e) for e in filenames]
        image2_list = [osp.join(root, e.replace('_left', '_right')) for e in filenames]
        disp_list = [osp.join(root, e.replace('image_left', 'depth_left').replace('left.png', 'left_depth.npy')) for e in filenames]

        for img1, img2, disp in zip(image1_list, image2_list, disp_list):
            self.image_list += [ [img1, img2] ]
            self.disparity_list += [ disp ]

class KITTI(StereoDataset):
    def __init__(self, aug_params=None, root='datasets/KITTI', image_set='training', args=None):
        super(KITTI, self).__init__(aug_params, sparse=True, reader=frame_utils.readDispKITTI, args=args)
        root = root if len(root)>0 else DATASET_ROOT
        assert os.path.exists(root), "check the existence: {}".format(root)

        image1_list = sorted(glob(os.path.join(root, image_set, 'image_2/*_10.png')))
        image2_list = sorted(glob(os.path.join(root, image_set, 'image_3/*_10.png')))
        disp_list = sorted(glob(os.path.join(root, 'training', 'disp_occ_0/*_10.png'))) if image_set == 'training' else [osp.join(root, 'training/disp_occ_0/000085_10.png')]*len(image1_list)

        for idx, (img1, img2, disp) in enumerate(zip(image1_list, image2_list, disp_list)):
            self.image_list += [ [img1, img2] ]
            self.disparity_list += [ disp ]


class KITTI2012(StereoDataset):
    def __init__(self, aug_params=None, root='datasets/KITTI2012', image_set='training', args=None):
        super(KITTI2012, self).__init__(aug_params, sparse=True, reader=frame_utils.readDispKITTI, args=args)
        root = root if len(root)>0 else DATASET_ROOT
        assert os.path.exists(root), "check the existence: {}".format(root)

        image1_list = sorted(glob(os.path.join(root, image_set, 'image_0/*_10.png')))
        image2_list = sorted(glob(os.path.join(root, image_set, 'image_1/*_10.png')))
        disp_list = sorted(glob(os.path.join(root, 'training', 'disp_occ/*_10.png'))) if image_set == 'training' else [osp.join(root, 'training/disp_occ_0/000085_10.png')]*len(image1_list)

        for idx, (img1, img2, disp) in enumerate(zip(image1_list, image2_list, disp_list)):
            self.image_list += [ [img1, img2] ]
            self.disparity_list += [ disp ]


class Middlebury(StereoDataset):
    def __init__(self, aug_params=None, root='datasets/Middlebury', split='F', image_set='training', args=None):
        super(Middlebury, self).__init__(aug_params, sparse=True, reader=frame_utils.readDispMiddlebury, args=args)
        root = root if len(root)>0 else DATASET_ROOT
        assert os.path.exists(root), "check the existence: {}".format(root)
        assert split in ["F", "H", "Q", "2014"]
        if split == "2014": # datasets/Middlebury/2014/Pipes-perfect/im0.png
            scenes = list((Path(root) / "2014").glob("*"))
            for scene in scenes:
                for s in ["E","L",""]:
                    self.image_list += [ [str(scene / "im0.png"), str(scene / f"im1{s}.png")] ]
                    self.disparity_list += [ str(scene / "disp0.pfm") ]
        else:
            lines = list(map(osp.basename, glob(os.path.join(root, f"MiddEval3/{image_set}{split}/*"))))
            image1_list = sorted([os.path.join(root, "MiddEval3", f'{image_set}{split}', f'{name}/im0.png') for name in lines])
            image2_list = sorted([os.path.join(root, "MiddEval3", f'{image_set}{split}', f'{name}/im1.png') for name in lines])
            disp_list = sorted([os.path.join(root, "MiddEval3", f'{image_set}{split}', f'{name}/disp0GT.pfm') for name in lines])
            if image_set=="training":
                assert len(image1_list) == len(image2_list) == len(disp_list) > 0, [image1_list, root, image_set, split]
            else:
                assert len(image1_list) == len(image2_list) > 0, [image1_list, root, image_set, split]
            for img1, img2, disp in zip(image1_list, image2_list, disp_list):
                self.image_list += [ [img1, img2] ]
                self.disparity_list += [ disp ]


class Booster(StereoDataset):
    def __init__(self, aug_params=None, root='datasets/booster/train/balanced', image_set='train', args=None):
        super(Booster, self).__init__(aug_params, sparse=True, reader=frame_utils.readDispBooster)
        assert os.path.exists(root), print(root)
        # image1_list = sorted(glob(os.path.join(root, image_set, "**/camera_00/im*.png"), recursive=True))
        image2_list = sorted(glob(os.path.join(root, image_set, "**/camera_02/im*.png"), recursive=True))
        image1_list = [img.replace("camera_02", "camera_00") for img in image2_list]

        disp_list = [os.path.join(os.path.split(x)[0].replace("camera_00", ""), 'disp_00.npy') for x in image1_list]
        mask_list = [os.path.join(os.path.split(x)[0].replace("camera_00", ""), 'mask_cat.png') for x in image1_list]
        right_disp_list = [os.path.join(os.path.split(x)[0].replace("camera_00", ""), 'disp_02.npy') for x in image1_list]
        
        for img1, img2, disp, disp_r, mask in zip(image1_list, image2_list, disp_list, right_disp_list,mask_list):
            self.image_list += [[img1, img2]]
            self.disparity_list += [disp]
            # self.trans_mask += [mask]


class NerfStereoDataset(StereoDataset):
    def __init__(self, aug_params=None, root='datasets/NerfStereo', image_set='training', args=None, txt_root=None):
        super(NerfStereoDataset, self).__init__(aug_params, sparse=True, reader=frame_utils.readDispNerfS, args=args)
        root = root if len(root)>0 else DATASET_ROOT
        assert os.path.exists(root), "check the existence: {}".format(root)
        
        if txt_root is None: 
            left_list = sorted(glob(os.path.join(root, "*/*/baseline_*/left/*.jpg"), recursive=True))
            image1_list = []
            for path in left_list:
                match = re.search(r"(.*?/Q/)", path)
                prefix = match.group(1)  # prefix
                suffix = os.path.basename(path)  # file name
                path_new = f"{prefix}center/{suffix}"
                image1_list.append( path_new )
            image2_list = sorted(glob(os.path.join(root, "*/*/baseline_*/right/*.jpg"), recursive=True))
            disp_list = sorted(glob(os.path.join(root, "*/*/baseline_*/disparity/*.png"), recursive=True))
            # dispr_list = sorted(glob(os.path.join(root, "**/*_right.disp.png"), recursive=True))
        else:
            image1_list = np.load( os.path.join(txt_root, 'image1_list.npy') )
            image2_list = np.load( os.path.join(txt_root, 'image2_list.npy') )
            disp_list = np.load( os.path.join(txt_root, 'disp_list.npy') )

        for idx, (img1, img2, disp) in enumerate(zip(image1_list, image2_list, disp_list)):
            self.image_list += [ [img1, img2] ]
            self.disparity_list += [ disp ]


class CREStereoDataset(StereoDataset):
    def __init__(self, aug_params=None, root='datasets/CREStereo_dataset', image_set='training', args=None, txt_root=None):
        super(CREStereoDataset, self).__init__(aug_params, sparse=True, reader=frame_utils.readDispCRES, args=args)
        root = root if len(root)>0 else DATASET_ROOT
        assert os.path.exists(root), "check the existence: {}".format(root)

        if txt_root is None: 
            image1_list = sorted(glob(os.path.join(root, "**/*_left.jpg"), recursive=True))
            image2_list = sorted(glob(os.path.join(root, "**/*_right.jpg"), recursive=True))
            disp_list = sorted(glob(os.path.join(root, "**/*_left.disp.png"), recursive=True))
        else:
            image1_list = np.load( os.path.join(txt_root, 'image1_list.npy') )
            image2_list = np.load( os.path.join(txt_root, 'image2_list.npy') )
            disp_list = np.load( os.path.join(txt_root, 'disp_list.npy') )
        # dispr_list = sorted(glob(os.path.join(root, "**/*_right.disp.png"), recursive=True))

        for idx, (img1, img2, disp) in enumerate(zip(image1_list, image2_list, disp_list)):
            self.image_list += [ [img1, img2] ]
            self.disparity_list += [ disp ]

class Trans(StereoDataset):
    def __init__(self, aug_params=None, root='./datasets/Trans', things_test=False, args=None):
        super(Trans, self).__init__(aug_params)
        self.root = root if len(root)>0 else DATASET_ROOT
        self.args = args
        self.extra_info["intrinsics"] = []
        
        if things_test:
            self._add_things("TEST")
        else:
            self._add_things("TRAIN")

    def _add_things(self, split='TRAIN'):
        original_length = len(self.disparity_list)

        left_images = sorted(glob(osp.join(self.root, split, '*/*/left/img/*.jpg')) )
        assert len(left_images)>0, f"Loaded 0 images from {self.root}"

        right_images = [ im.replace('left', 'right') for im in left_images ]
        disparity_images = [ im.replace('img', 'disparity').replace('.jpg', '.pfm') for im in left_images ]
        disparity_images_noTran = [im.replace('img', 'disparity_without_trans').replace('.jpg', '.pfm') for im in left_images ]

        for idx, (img1, img2, disp, disp_noTran) in enumerate(zip(left_images, right_images, disparity_images, disparity_images_noTran)):
            self.image_list += [ [img1, img2] ]
            self.disparity_list += [ disp ]
            # self.multi_label.append([disp, disp_noTran])
            self.extra_info["intrinsics"] += [ [933.3333333333334, 787.5, 480.0, 270.0] ]
        logging.info("-"*10 + f"Added {len(self.disparity_list) - original_length} from Trans")

class Fooling3DDataset(StereoDataset):
    def __init__(self, aug_params=None, root='datasets/Fooling3D', image_set='training', args=None):
        super(Fooling3DDataset, self).__init__(aug_params, sparse=True, reader=frame_utils.readDispFooling3D)
        assert os.path.exists(root)
        self.root = root
        self.image_set = image_set
        self.video_frames_info = {}
        
        self._add_mono()
        self._build_video_frames_info()

    def _add_mono(self):
        origin_length = len(self.disparity_list)
        print(f"using {self.image_set} in fooling3D")
        
        if self.image_set=="training":
            df = pd.read_csv(os.path.join(self.root, 'meta_data/scale_factors.csv'), header=None)

            # df.columns = ['path', 'scale']
            # video_name = "Service_Cars_1_deleted_scene_3d_remake_Servio_Comunitrio"
            # df = df[df['path'].str.contains(video_name, case=False, na=False)]

            self.scale_factor = dict(zip(
                df.iloc[:, 0].apply(lambda x: x.replace('/data2', './datasets')),
                df.iloc[:, 1]
            ))
            # right_images = sorted(glob(os.path.join(self.root, 'video_frame_sequence_right/*/*/*.png')))
            right_images = df.iloc[:, 0].apply(lambda x: x.replace('/data2', './datasets')).tolist()
            disp_list =  [ im.replace('video_frame_sequence_right', 'depth_rect') for im in right_images ]
            left_images = [ im.replace('video_frame_sequence_right', 'video_frame_sequence') for im in right_images ]

            assert len(left_images) == len(right_images) == len(disp_list) > 0, [len(left_images), len(right_images), len(disp_list)]
            for img1, img2, disp in zip(left_images, right_images, disp_list):
                self.image_list += [ [img1, img2] ]
                self.disparity_list += [ disp ]
        
        elif self.image_set=="testing":
            with open(os.path.join(self.root, 'meta_data/testing_enter.pkl'), 'rb') as f:
                data = pickle.load(f)
            
            self.extra_info["mask"] = []
            for key, frame_dict in data.items():
                left_image_path  = os.path.join(self.root, "real_data/testing", frame_dict["left"])
                right_image_path = os.path.join(self.root, "real_data/testing", frame_dict["right"])
                disp_image_path  = os.path.join(self.root, "real_data/testing", frame_dict["disp"])
                mask_image_path  = os.path.join(self.root, "real_data/testing", frame_dict["mask"])

                self.image_list += [ [left_image_path, right_image_path] ]
                self.disparity_list += [ disp_image_path ]
                self.extra_info["mask"] += [ mask_image_path ]

            assert len(self.image_list) == len(self.disparity_list) == len(self.extra_info["mask"]) > 0, \
                   [len(self.image_list), len(self.disparity_list), len(self.extra_info["mask"])]

        else:
            raise Exception(f"{self.image_set} is not in ['training', 'testing']")
        
        logging.info(f"Added {len(self.disparity_list) - origin_length} from Fooling3D Mono")
    
    def _build_video_frames_info(self):
        for idx, img_path in enumerate(self.disparity_list):
            parts = img_path.split('/')
            video_name = parts[-2]
            frame_name = parts[-1]

            if video_name not in self.video_frames_info:
                self.video_frames_info[video_name] = []

            self.video_frames_info[video_name].append(idx)
        self.video_frames_info = list(self.video_frames_info.values())




class Fooling3DBatchSampler(data.Sampler):
    def __init__(self, dataset, batch_size):
        """
        Args:
            dataset (Dataset): The dataset to sample from.
            batch_size (int): The size of each batch (how many frames from the same video).
        """
        self.dataset = dataset
        self.batch_size = batch_size

    def __iter__(self):
        """
        This will return indices of frames in a single video folder, ensuring batch contains only frames from that video.
        """
        for video_idx in range(len(self.dataset.video_frames_info)):
            frames_info = self.dataset.video_frames_info[video_idx]
            num_frames = len(frames_info)
            frame_idx_list = list(np.arange(num_frames))

            # # Shuffle the frame indices if shuffle is True
            # if self.shuffle:
            #     np.random.shuffle(frame_idx_list)

            # If frames count is not divisible by batch size, repeat the last frame
            if num_frames % self.batch_size != 0:
                num_repeat = self.batch_size - (num_frames % self.batch_size)
                frame_idx_list += [frame_idx_list[-1]] * num_repeat  # Add last frame to fill up batch

            # Yield frames in batches of batch_size
            for i in range(0, len(frame_idx_list), self.batch_size):
                batch_info = [frames_info[frame_idx] for frame_idx in frame_idx_list[i:i + self.batch_size]]
                yield batch_info

    def __len__(self):
        """
        The length of the sampler is the number of total batches in all videos.
        """
        total_batches = 0
        for frames_info in self.dataset.video_frames_info:
            total_batches += len(frames_info) // self.batch_size + (1 if len(frames_info) % self.batch_size != 0 else 0)
        return total_batches


from torch.utils.data.distributed import DistributedSampler
class DistributedFooling3DBatchSampler(DistributedSampler):
    def __init__(self, dataset, batch_size, num_replicas=None, rank=None):
        """
        Args:
            dataset (Dataset): The dataset to sample from.
            batch_size (int): The size of each batch (how many frames from the same video).
            num_replicas (int): Total number of processes (GPUs) across all nodes.
            rank (int): Rank of the current process (GPU) in the group of workers.
        """
        self.dataset = dataset
        self.batch_size = batch_size
        self.num_replicas = num_replicas if num_replicas is not None else torch.distributed.get_world_size()
        self.rank = rank if rank is not None else torch.distributed.get_rank()

    def __iter__(self):
        """
        This will return indices of frames in a single video folder, ensuring batch contains only frames from that video.
        Distributes the frames across different processes.
        """
        for video_idx in range(len(self.dataset.video_frames_info)):
            frames_info = self.dataset.video_frames_info[video_idx]
            num_frames = len(frames_info)
            frame_idx_list = list(np.arange(num_frames))

            # # Shuffle the frame indices if shuffle is True
            # if self.shuffle:
            #     np.random.shuffle(frame_idx_list)

            # If frames count is not divisible by batch size, repeat the last frame
            if num_frames % self.batch_size != 0:
                num_repeat = self.batch_size - (num_frames % self.batch_size)
                frame_idx_list += [frame_idx_list[-1]] * num_repeat  # Add last frame to fill up batch

            # Total number of batches across all replicas
            num_batches = len(frame_idx_list) // self.batch_size + (1 if len(frame_idx_list) % self.batch_size != 0 else 0)
            
            # Divide the dataset into chunks and ensure each rank gets its share
            # Find out how many batches each rank should process
            chunks_per_rank = num_batches // self.num_replicas
            remainder = num_batches % self.num_replicas
            start_idx = self.rank * chunks_per_rank + min(self.rank, remainder)
            end_idx = (self.rank + 1) * chunks_per_rank + min(self.rank + 1, remainder)
            
            # Generate the frames indices for the current process's portion of the data
            for i in range(start_idx, end_idx):
                batch_info = [frames_info[frame_idx] for frame_idx in frame_idx_list[i * self.batch_size:(i + 1) * self.batch_size]]
                yield batch_info

    def __len__(self):
        """
        The length of the sampler is the total number of batches divided across all processes.
        """
        total_batches = 0
        for frames_info in self.dataset.video_frames_info:
            total_batches += len(frames_info) // self.batch_size + (1 if len(frames_info) % self.batch_size != 0 else 0)
        
        # Divide the total batches by the number of processes
        return total_batches // self.num_replicas + (1 if total_batches % self.num_replicas > self.rank else 0)

  
def fetch_dataloader(args):
    """ Create the data loader for the corresponding trainign set """

    aug_params = {'crop_size': args.image_size, 'min_scale': args.spatial_scale[0], 'max_scale': args.spatial_scale[1], 'do_flip': False, 'yjitter': not args.noyjitter}
    if hasattr(args, "saturation_range") and args.saturation_range is not None:
        aug_params["saturation_range"] = args.saturation_range
    if hasattr(args, "img_gamma") and args.img_gamma is not None:
        aug_params["gamma"] = args.img_gamma
    if hasattr(args, "do_flip") and args.do_flip is not None:
        aug_params["do_flip"] = args.do_flip

    train_dataset = None
    for dataset_name in args.train_datasets:
        if dataset_name.startswith("middlebury_"):
            new_dataset = Middlebury(aug_params, split=dataset_name.replace('middlebury_',''), args=args)
            logging.info(f"Adding {len(new_dataset)} samples from Middlebury")
        elif dataset_name == 'sceneflow':
            clean_dataset = SceneFlowDatasets(aug_params, dstype='frames_cleanpass', args=args)
            final_dataset = SceneFlowDatasets(aug_params, dstype='frames_finalpass', args=args)
            new_dataset = (clean_dataset*4) + (final_dataset*4)
            logging.info(f"Adding {len(new_dataset)} samples from SceneFlow")
        elif 'kitti' in dataset_name:
            new_dataset = KITTI(aug_params, split=dataset_name, args=args)
            logging.info(f"Adding {len(new_dataset)} samples from KITTI")
        elif dataset_name == 'sintel_stereo':
            new_dataset = SintelStereo(aug_params, args=args)*140
            logging.info(f"Adding {len(new_dataset)} samples from Sintel Stereo")
        elif dataset_name == 'falling_things':
            new_dataset = FallingThings(aug_params, args=args)*5
            logging.info(f"Adding {len(new_dataset)} samples from FallingThings")
        elif dataset_name.startswith('tartan_air'):
            new_dataset = TartanAir(aug_params, keywords=dataset_name.split('_')[2:])
            logging.info(f"Adding {len(new_dataset)} samples from Tartain Air")
        elif 'nerfstereo' in dataset_name:
            new_dataset = NerfStereoDataset(aug_params, args=args, root='./datasets/NerfStereo', txt_root='./datasets/NerfStereo/../')
            logging.info(f"Adding {len(new_dataset)} samples from NerfStereoDataset")
        elif 'crestereo' in dataset_name:
            new_dataset = CREStereoDataset(aug_params, args=args, txt_root='./datasets/CREStereo_dataset/../')
            logging.info(f"Adding {len(new_dataset)} samples from CREStereoDataset")
        elif dataset_name == 'Trans':
            new_dataset = Trans(aug_params, args=args)
            logging.info(f"Adding {len(new_dataset)} samples from Trans")
        elif dataset_name.lower() == 'fooling3d':
            new_dataset = Fooling3DDataset(aug_params, args=args, root='./datasets/Fooling3D')
            # print("+"*10, hasattr(args, 'enable_sampler') and args.enable_sampler)
            if hasattr(args, 'enable_sampler') and args.enable_sampler:
                # sampler = Fooling3DBatchSampler(new_dataset, args.batch_size)
                sampler = DistributedFooling3DBatchSampler(new_dataset, args.batch_size)
            logging.info(f"Adding {len(new_dataset)} samples from Fooling3DDataset")
            # TODO: Add Fooling3D dataset with only one sampler may cause conflict with other datasets
        train_dataset = new_dataset if train_dataset is None else train_dataset + new_dataset

    # train_loader = data.DataLoader(train_dataset, batch_size=args.batch_size, 
    #     pin_memory=True, shuffle=True, num_workers=int(os.environ.get('SLURM_CPUS_PER_TASK', 6))-2, drop_last=True)
    train_loader = get_loader(train_dataset, args)
    train_loader.sampler.set_epoch(0)

    logging.info('Training with %d image pairs' % len(train_dataset))
    return train_loader