a
    d                     @   s  d dl Z d dlZd dlZd dlZd dlZd dlmZmZ d dlmZ d dl	m
Z
 d dlmZmZmZmZmZmZ d dlZd dlmZ ddlmZmZmZ dd	lmZ eejejeej ejf Zeejejeej f Zd
Ze j eddZ!G dd deeZ"G dd de"Z#G dd de"Z$G dd de"Z%G dd de"Z&G dd de"Z'G dd de"Z(G dd de"Z)G dd de"Z*G dd de"Z+G d d! d!e"Z,dS )"    N)ABCabstractmethod)glob)Path)CallablecastListOptionalTupleUnion)Image   )	_read_pfmdownload_and_extract_archiveverify_str_arg)VisionDataset )Zslice_channelsc                       s   e Zd ZdZdZdeee dd fddZe	ee
f ejddd	Zdeee eeeee f  d
ddZeeeeej eej f dddZee	eef dddZedddZ  ZS )StereoMatchingDatasetz+Base interface for Stereo matching datasetsFNroot
transformsreturnc                    s$   t  j|d || _g | _g | _dS )a}  
        Args:
            root(str): Root directory of the dataset.
            transforms(callable, optional): A function/transform that takes in Tuples of
                (images, disparities, valid_masks) and returns a transformed version of each of them.
                images is a Tuple of (``PIL.Image``, ``PIL.Image``)
                disparities is a Tuple of (``np.ndarray``, ``np.ndarray``) with shape (1, H, W)
                valid_masks is a Tuple of (``np.ndarray``, ``np.ndarray``) with shape (H, W)
                In some cases, when a dataset does not provide disparities, the ``disparities`` and
                ``valid_masks`` can be Tuples containing None values.
                For training splits generally the datasets provide a minimal guarantee of
                images: (``PIL.Image``, ``PIL.Image``)
                disparities: (``np.ndarray``, ``None``) with shape (1, H, W)
                Optionally, based on the dataset, it can return a ``mask`` as well:
                valid_masks: (``np.ndarray | None``, ``None``) with shape (H, W)
                For some test splits, the datasets provides outputs that look like:
                imgaes: (``PIL.Image``, ``PIL.Image``)
                disparities: (``None``, ``None``)
                Optionally, based on the dataset, it can return a ``mask`` as well:
                valid_masks: (``None``, ``None``)
        r   N)super__init__r   _images_disparities)selfr   r   	__class__r   n/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/torchvision/datasets/_stereo_matching.pyr      s    zStereoMatchingDataset.__init__	file_pathr   c                 C   s"   t |}|jdkr|d}|S )NZRGB)r   openmodeconvert)r   r"   imgr   r   r    	_read_img:   s    


zStereoMatchingDataset._read_img)paths_left_patternpaths_right_patternr   c              
   C   s   t tt|}|r&t tt|}nt dd |D }|sJtd| |s\td| t|t|krtdt| dt| d| d| d	t d	d t||D }|S )
Nc                 s   s   | ]
}d V  qd S Nr   .0_r   r   r    	<genexpr>L       z4StereoMatchingDataset._scan_pairs.<locals>.<genexpr>z0Could not find any files matching the patterns: zFound z left files but z# right files using:
 left pattern: z
right pattern: 
c                 s   s   | ]\}}||fV  qd S r*   r   )r,   leftrightr   r   r    r.   [   r/   )listsortedr   FileNotFoundErrorlen
ValueErrorzip)r   r(   r)   Z
left_pathsZright_pathspathsr   r   r    _scan_pairs@   s$    z!StereoMatchingDataset._scan_pairsc                 C   s   d S r*   r   )r   r"   r   r   r    _read_disparity^   s    z%StereoMatchingDataset._read_disparityindexr   c                 C   s   |  | j| d }|  | j| d }| | j| d \}}| | j| d \}}||f}||f}	||f}
| jdur| ||	|
\}}	}
| js|
d dur|d |d |	d ttj|
d fS |d |d |	d fS dS )ao  Return example at given index.

        Args:
            index(int): The index of the example to retrieve

        Returns:
            tuple: A 3 or 4-tuple with ``(img_left, img_right, disparity, Optional[valid_mask])`` where ``valid_mask``
                can be a numpy boolean mask of shape (H, W) if the dataset provides a file
                indicating which disparity pixels are valid. The disparity is a numpy array of
                shape (1, H, W) and the images are PIL images. ``disparity`` is None for
                datasets on which for ``split="test"`` the authors did not provide annotations.
        r   r   N)	r'   r   r;   r   r   _has_built_in_disparity_maskr   npndarray)r   r=   Zimg_leftZ	img_rightZdsp_map_leftZvalid_mask_leftZdsp_map_rightZvalid_mask_rightimgsZdsp_mapsZvalid_masksr   r   r    __getitem__c   s     
$z!StereoMatchingDataset.__getitem__)r   c                 C   s
   t | jS r*   )r6   r   )r   r   r   r    __len__   s    zStereoMatchingDataset.__len__)N)N)__name__
__module____qualname____doc__r>   strr	   r   r   r   r   r   r'   r   r
   r:   r   r?   r@   r;   intT1T2rB   rC   __classcell__r   r   r   r    r      s   	 &#r   c                       s^   e Zd ZdZdeee dd fddZeee	j
df dddZeed	 fd
dZ  ZS )CarlaStereoai  
    Carla simulator data linked in the `CREStereo github repo <https://github.com/megvii-research/CREStereo>`_.

    The dataset is expected to have the following structure: ::

        root
            carla-highres
                trainingF
                    scene1
                        img0.png
                        img1.png
                        disp0GT.pfm
                        disp1GT.pfm
                        calib.txt
                    scene2
                        img0.png
                        img1.png
                        disp0GT.pfm
                        disp1GT.pfm
                        calib.txt
                    ...

    Args:
        root (string): Root directory where `carla-highres` is located.
        transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version.
    Nr   c           	         s   t  || t|d }t|d d d }t|d d d }| ||}|| _t|d d d }t|d d d }| ||}|| _d S )Nzcarla-highresZ	trainingF*im0.pngim1.pngdisp0GT.pfmzdisp1GT.pfmr   r   r   rH   r:   r   r   )	r   r   r   left_image_patternright_image_patternrA   left_disparity_patternright_disparity_patterndisparitiesr   r   r    r      s    zCarlaStereo.__init__r!   c                 C   s   t |}t|}d }||fS r*   _read_pfm_filer?   absr   r"   disparity_map
valid_maskr   r   r    r;      s    
zCarlaStereo._read_disparityr<   c                    s   t tt |S a  Return example at given index.

        Args:
            index(int): The index of the example to retrieve

        Returns:
            tuple: A 3-tuple with ``(img_left, img_right, disparity)``.
            The disparity is a numpy array of shape (1, H, W) and the images are PIL images.
            If a ``valid_mask`` is generated within the ``transforms`` parameter,
            a 4-tuple with ``(img_left, img_right, disparity, valid_mask)`` is returned.
        r   rJ   r   rB   r   r=   r   r   r    rB      s    zCarlaStereo.__getitem__)NrD   rE   rF   rG   rH   r	   r   r   r
   r?   r@   r;   rI   rJ   rB   rL   r   r   r   r    rM      s   rM   c                       sh   e Zd ZdZdZdeeee dd fddZee	ee
j df dd	d
Zeed fddZ  ZS )Kitti2012Stereoa  
    KITTI dataset from the `2012 stereo evaluation benchmark <http://www.cvlibs.net/datasets/kitti/eval_stereo_flow.php>`_.
    Uses the RGB images for consistency with KITTI 2015.

    The dataset is expected to have the following structure: ::

        root
            Kitti2012
                testing
                    colored_0
                        1_10.png
                        2_10.png
                        ...
                    colored_1
                        1_10.png
                        2_10.png
                        ...
                training
                    colored_0
                        1_10.png
                        2_10.png
                        ...
                    colored_1
                        1_10.png
                        2_10.png
                        ...
                    disp_noc
                        1.png
                        2.png
                        ...
                    calib

    Args:
        root (string): Root directory where `Kitti2012` is located.
        split (string, optional): The dataset split of scenes, either "train" (default) or "test".
        transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version.
    TtrainNr   splitr   r   c                    s   t  || t|ddd t|d |d  }t|d d }t|d d }| ||| _|d	krt|d
 d }| |d | _ntdd | jD | _d S )Nre   rc   testZvalid_valuesZ	Kitti2012ingZ	colored_0z*_10.pngZ	colored_1rc   Zdisp_noc*.pngc                 s   s   | ]
}d V  qdS NNNr   r+   r   r   r    r.     r/   z+Kitti2012Stereo.__init__.<locals>.<genexpr>	r   r   r   r   rH   r:   r   r   r3   )r   r   re   r   left_img_patternright_img_patterndisparity_patternr   r   r    r      s    zKitti2012Stereo.__init__r!   c                 C   sB   |d u rdS t t|d }|d d d d d f }d }||fS Nrl   g      p@r?   asarrayr   r#   r[   r   r   r    r;     s    zKitti2012Stereo._read_disparityr<   c                    s   t tt |S a  Return example at given index.

        Args:
            index(int): The index of the example to retrieve

        Returns:
            tuple: A 4-tuple with ``(img_left, img_right, disparity, valid_mask)``.
            The disparity is a numpy array of shape (1, H, W) and the images are PIL images.
            ``valid_mask`` is implicitly ``None`` if the ``transforms`` parameter does not
            generate a valid mask.
            Both ``disparity`` and ``valid_mask`` are ``None`` if the dataset split is test.
        r_   r`   r   r   r    rB     s    zKitti2012Stereo.__getitem__)rc   NrD   rE   rF   rG   r>   rH   r	   r   r   r
   r?   r@   r;   rI   rJ   rB   rL   r   r   r   r    rb      s
   &rb   c                       sh   e Zd ZdZdZdeeee dd fddZee	ee
j df dd	d
Zeed fddZ  ZS )Kitti2015Stereoa<  
    KITTI dataset from the `2015 stereo evaluation benchmark <http://www.cvlibs.net/datasets/kitti/eval_scene_flow.php>`_.

    The dataset is expected to have the following structure: ::

        root
            Kitti2015
                testing
                    image_2
                        img1.png
                        img2.png
                        ...
                    image_3
                        img1.png
                        img2.png
                        ...
                training
                    image_2
                        img1.png
                        img2.png
                        ...
                    image_3
                        img1.png
                        img2.png
                        ...
                    disp_occ_0
                        img1.png
                        img2.png
                        ...
                    disp_occ_1
                        img1.png
                        img2.png
                        ...
                    calib

    Args:
        root (string): Root directory where `Kitti2015` is located.
        split (string, optional): The dataset split of scenes, either "train" (default) or "test".
        transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version.
    Trc   Nrd   c                    s   t  || t|ddd t|d |d  }t|d d }t|d d }| ||| _|d	krt|d
 d }t|d d }| ||| _ntdd | jD | _d S )Nre   rf   rh   Z	Kitti2015ri   Zimage_2rj   Zimage_3rc   Z
disp_occ_0Z
disp_occ_1c                 s   s   | ]
}d V  qdS rk   r   r+   r   r   r    r.   Z  r/   z+Kitti2015Stereo.__init__.<locals>.<genexpr>rm   r   r   re   r   rn   ro   rU   rV   r   r   r    r   K  s    zKitti2015Stereo.__init__r!   c                 C   sB   |d u rdS t t|d }|d d d d d f }d }||fS rq   rr   r[   r   r   r    r;   \  s    zKitti2015Stereo._read_disparityr<   c                    s   t tt |S rt   r_   r`   r   r   r    rB   g  s    zKitti2015Stereo.__getitem__)rc   Nru   r   r   r   r    rv     s
   )rv   c                	       s   e Zd ZdZg dg dg ddZdZdeeee eee	 ed
d fddZ
eeef ejd fddZeeed eejejf f dddZed
dddZeed fddZ  ZS )Middlebury2014StereoaI	  Publicly available scenes from the Middlebury dataset `2014 version <https://vision.middlebury.edu/stereo/data/scenes2014/>`.

    The dataset mostly follows the original format, without containing the ambient subdirectories.  : ::

        root
            Middlebury2014
                train
                    scene1-{perfect,imperfect}
                        calib.txt
                        im{0,1}.png
                        im1E.png
                        im1L.png
                        disp{0,1}.pfm
                        disp{0,1}-n.png
                        disp{0,1}-sd.pfm
                        disp{0,1}y.pfm
                    scene2-{perfect,imperfect}
                        calib.txt
                        im{0,1}.png
                        im1E.png
                        im1L.png
                        disp{0,1}.pfm
                        disp{0,1}-n.png
                        disp{0,1}-sd.pfm
                        disp{0,1}y.pfm
                    ...
                additional
                    scene1-{perfect,imperfect}
                        calib.txt
                        im{0,1}.png
                        im1E.png
                        im1L.png
                        disp{0,1}.pfm
                        disp{0,1}-n.png
                        disp{0,1}-sd.pfm
                        disp{0,1}y.pfm
                    ...
                test
                    scene1
                        calib.txt
                        im{0,1}.png
                    scene2
                        calib.txt
                        im{0,1}.png
                    ...

    Args:
        root (string): Root directory of the Middleburry 2014 Dataset.
        split (string, optional): The dataset split of scenes, either "train" (default), "test", or "additional"
        use_ambient_views (boolean, optional): Whether to use different expose or lightning views when possible.
            The dataset samples with equal probability between ``[im1.png, im1E.png, im1L.png]``.
        calibration (string, optional): Whether or not to use the calibrated (default) or uncalibrated scenes.
        transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version.
        download (boolean, optional): Whether or not to download the dataset in the ``root`` directory.
    )
Z
AdirondackZ	JadeplantZ
MotorcycleZPianoZPipesZPlayroomZ	PlaytableZRecycleZShelvesZVintage)ZBackpackZBicycle1ZCableZ
Classroom1ZCouchZFlowersZMaskZShopvacZSticksZStorageZSword1ZSword2ZUmbrella)ZPlantsZClassroom2EZ
Classroom2Z	AustraliaZDjembeLZCrusadePZCrusadeZHoopsZBicycle2Z	StaircaseZNewkubaZ
AustraliaPZDjembeZ
LivingroomZComputer)rc   
additionalrg   Trc   perfectFN)r   re   calibrationuse_ambient_viewsr   downloadr   c                    s  t  || t|ddd || _|rFt|ddd |dkrdtdn|dkrdtd| d	| d
|rr| | t|d }tj	|| st
d| d| j|  t fddt|| D st
d| ddgdgdgddgd| }|D ]}d| }	t|| |	 d }
t|| |	 d }|  j| |
|7  _|dkrdtdd | jD | _qt|| |	 d }t|| |	 d }|  j| ||7  _q|| _d S )Nre   )rc   rg   ry   rh   r{   )rz   	imperfectbothNrg   zMSplit 'test' has only no calibration settings, please set `calibration=None`.zSplit 'zr' has calibration settings, however None was provided as an argument.
Setting calibration to 'perfect' for split 'zF'. Available calibration settings are: 'perfect', 'imperfect', 'both'.Middlebury2014zThe z7 directory was not found in the provided root directoryc                 3   s"   | ]} D ]}| |V  q
qd S r*   )
startswith)r,   scenesZsplit_scenesr   r    r.     s   z0Middlebury2014Stereo.__init__.<locals>.<genexpr>z:Provided root folder does not contain any scenes from the z split. z-perfectz
-imperfect)Nrz   r~   r   rN   rO   rP   c                 s   s   | ]
}d V  qdS rk   r   r+   r   r   r    r.     r/   z	disp0.pfmz	disp1.pfm)r   r   r   re   r7   _download_datasetr   ospathexistsr5   splitsanylistdirrH   r   r:   r3   r   r|   )r   r   re   r{   r|   r   r}   Zcalibrartion_suffixesZcalibration_suffixZscene_patternrn   ro   Zleft_dispartity_patternZright_dispartity_patternr   r   r    r     sR    	



zMiddlebury2014Stereo.__init__r!   c                    sp   t |tst|}|jdkrd| jrd|j t fdddD }ttdd |}|| t	|}t
 |S )a  
        Function that reads either the original right image or an augmented view when ``use_ambient_views`` is True.
        When ``use_ambient_views`` is True, the dataset will return at random one of ``[im1.png, im1E.png, im1L.png]``
        as the right image.
        rP   c                 3   s   | ]} | V  qd S r*   r   )r,   Z	view_name	base_pathr   r    r.   0  r/   z1Middlebury2014Stereo._read_img.<locals>.<genexpr>)zim1E.pngzim1L.pngc                 S   s   t j| S r*   )r   r   r   )pr   r   r    <lambda>2  r/   z0Middlebury2014Stereo._read_img.<locals>.<lambda>)
isinstancer   namer|   parentr3   filterappendrandomchoicer   r'   )r   r"   Zambient_file_pathsr   r   r    r'   "  s    


zMiddlebury2014Stereo._read_imgrl   c                 C   sB   |d u rdS t |}t|}d||tjk< |dkd}||fS )Nrl   r   )rY   r?   rZ   infZsqueezer[   r   r   r    r;   8  s    
z$Middlebury2014Stereo._read_disparity)r   r   c                    sJ  d}t  d  | j}|dkr| j| D ]h} | }dD ]V}| d| }| d| d}td|  ||  s8t|| dt|d	d
 q8q(nt d  t	 fdd| jd D rFd}	t|	t d	d t
t d D ]J\}
}}|D ]:} d }t |
| }tj|d	d tt|t| qqtt d  d S )Nz8https://vision.middlebury.edu/stereo/data/scenes2014/zipr   rg   )rz   r~   -/z.zipzDownloading T)urlfilenamedownload_rootremove_finishedc                 3   s    | ]}|t  d  vV  qdS )rg   N)r   r   )r,   r   r   r   r    r.   Z  r/   z9Middlebury2014Stereo._download_dataset.<locals>.<genexpr>zEhttps://vision.middlebury.edu/stereo/submit3/zip/MiddEval3-data-F.zip)r   r   r   zMiddEval3/testF)exist_okZ	MiddEval3)r   re   r   printr   r   rH   r   makedirsr   walkshutilmovermtree)r   r   base_urlZ
split_nameZsplit_sceneZ
split_rootr{   Z
scene_nameZ	scene_urlZtest_set_urlZ	scene_dirZscene_namesr-   r   Zscene_dst_dirZscene_src_dirr   r   r    r   C  s8    z&Middlebury2014Stereo._download_datasetr<   c                    s   t tt |S )az  Return example at given index.

        Args:
            index(int): The index of the example to retrieve

        Returns:
            tuple: A 4-tuple with ``(img_left, img_right, disparity, valid_mask)``.
            The disparity is a numpy array of shape (1, H, W) and the images are PIL images.
            ``valid_mask`` is implicitly ``None`` for `split=test`.
        r   rK   r   rB   r`   r   r   r    rB   j  s    z Middlebury2014Stereo.__getitem__)rc   rz   FNF)rD   rE   rF   rG   r   r>   rH   r	   boolr   r   r   r   r   r'   r
   r?   r@   r;   r   rI   rK   rB   rL   r   r   r   r    rx   w  s0   9/     A('rx   c                       sb   e Zd ZdZdZdeee dd fddZee	e
jdf ddd	Zeed
 fddZ  ZS )	CREStereoa  Synthetic dataset used in training the `CREStereo <https://arxiv.org/pdf/2203.11483.pdf>`_ architecture.
    Dataset details on the official paper `repo <https://github.com/megvii-research/CREStereo>`_.

    The dataset is expected to have the following structure: ::

        root
            CREStereo
                tree
                    img1_left.jpg
                    img1_right.jpg
                    img1_left.disp.jpg
                    img1_right.disp.jpg
                    img2_left.jpg
                    img2_right.jpg
                    img2_left.disp.jpg
                    img2_right.disp.jpg
                    ...
                shapenet
                    img1_left.jpg
                    img1_right.jpg
                    img1_left.disp.jpg
                    img1_right.disp.jpg
                    ...
                reflective
                    img1_left.jpg
                    img1_right.jpg
                    img1_left.disp.jpg
                    img1_right.disp.jpg
                    ...
                hole
                    img1_left.jpg
                    img1_right.jpg
                    img1_left.disp.jpg
                    img1_right.disp.jpg
                    ...

    Args:
        root (str): Root directory of the dataset.
        transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version.
    TNr   c                    s   t  || t|d }g d}|D ]x}t|| d }t|| d }| ||}|  j|7  _t|| d }t|| d }	| ||	}
|  j|
7  _q&d S )Nr   )ZshapenetZ
reflectivetreeZholez
*_left.jpgz*_right.jpgz*_left.disp.pngz*_right.disp.pngrR   )r   r   r   dirsr   rS   rT   rA   rU   rV   rW   r   r   r    r     s    zCREStereo.__init__r!   c                 C   s<   t jt|t jd}|d d d d d f d }d }||fS )NZdtypeg      @@r?   rs   r   r#   float32r[   r   r   r    r;     s    zCREStereo._read_disparityr<   c                    s   t tt |S )a  Return example at given index.

        Args:
            index(int): The index of the example to retrieve

        Returns:
            tuple: A 4-tuple with ``(img_left, img_right, disparity, valid_mask)``.
            The disparity is a numpy array of shape (1, H, W) and the images are PIL images.
            ``valid_mask`` is implicitly ``None`` if the ``transforms`` parameter does not
            generate a valid mask.
        r_   r`   r   r   r    rB     s    zCREStereo.__getitem__)Nru   r   r   r   r    r   x  s   ) r   c                       s`   e Zd ZdZdeeee dd fddZeee	j
df ddd	Zeed
 fddZ  ZS )FallingThingsStereoa  `FallingThings <https://research.nvidia.com/publication/2018-06_falling-things-synthetic-dataset-3d-object-detection-and-pose-estimation>`_ dataset.

    The dataset is expected to have the following structure: ::

        root
            FallingThings
                single
                    dir1
                        scene1
                            _object_settings.json
                            _camera_settings.json
                            image1.left.depth.png
                            image1.right.depth.png
                            image1.left.jpg
                            image1.right.jpg
                            image2.left.depth.png
                            image2.right.depth.png
                            image2.left.jpg
                            image2.right
                            ...
                        scene2
                    ...
                mixed
                    scene1
                        _object_settings.json
                        _camera_settings.json
                        image1.left.depth.png
                        image1.right.depth.png
                        image1.left.jpg
                        image1.right.jpg
                        image2.left.depth.png
                        image2.right.depth.png
                        image2.left.jpg
                        image2.right
                        ...
                    scene2
                    ...

    Args:
        root (string): Root directory where FallingThings is located.
        variant (string): Which variant to use. Either "single", "mixed", or "both".
        transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version.
    singleN)r   variantr   r   c                    s   t  || t|d }t|ddd dgdgddgd| }tdd tdd}|D ]}t|| ||  d	 }t|| ||  d
 }|  j| ||7  _t|| ||  d }	t|| ||  d }
|  j| |	|
7  _qZd S )NZFallingThingsr   )r   mixedr   rh   r   r   rN   )r   r   z
*.left.jpgz*.right.jpgz*.left.depth.pngz*.right.depth.pngr   r   r   r   rH   r   r:   r   )r   r   r   r   variantsZsplit_prefixr   rn   ro   rU   rV   r   r   r    r     s&    
zFallingThingsStereo.__init__r!   c                 C   s   t t|}t|jd }t|dp}t|}|d d d d }d\}}|| | |t j	 }	|	d d d d d f }	d }
|	|
fW  d    S 1 s0    Y  d S )Nz_camera_settings.jsonrZcamera_settingsr   Zintrinsic_settingsZfx)   d   )
r?   rs   r   r#   r   r   jsonloadastyper   )r   r"   depthZcamera_settings_pathfZ
intrinsicsZfocalZbaselineZpixel_constantr\   r]   r   r   r    r;     s    
z#FallingThingsStereo._read_disparityr<   c                    s   t tt |S r^   r_   r`   r   r   r    rB   )  s    zFallingThingsStereo.__getitem__)r   Nra   r   r   r   r    r     s   ,r   c                       sb   e Zd ZdZdeeeee dd fddZeee	j
df dd	d
Zeed fddZ  ZS )SceneFlowStereoa  Dataset interface for `Scene Flow <https://lmb.informatik.uni-freiburg.de/resources/datasets/SceneFlowDatasets.en.html>`_ datasets.
    This interface provides access to the `FlyingThings3D, `Monkaa` and `Driving` datasets.

    The dataset is expected to have the following structure: ::

        root
            SceneFlow
                Monkaa
                    frames_cleanpass
                        scene1
                            left
                                img1.png
                                img2.png
                            right
                                img1.png
                                img2.png
                        scene2
                            left
                                img1.png
                                img2.png
                            right
                                img1.png
                                img2.png
                    frames_finalpass
                        scene1
                            left
                                img1.png
                                img2.png
                            right
                                img1.png
                                img2.png
                        ...
                        ...
                    disparity
                        scene1
                            left
                                img1.pfm
                                img2.pfm
                            right
                                img1.pfm
                                img2.pfm
                FlyingThings3D
                    ...
                    ...

    Args:
        root (string): Root directory where SceneFlow is located.
        variant (string): Which dataset variant to user, "FlyingThings3D" (default), "Monkaa" or "Driving".
        pass_name (string): Which pass to use, "clean" (default), "final" or "both".
        transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version.

    FlyingThings3DcleanN)r   r   	pass_namer   r   c                    s(  t  || t|d }t|ddd t|ddd dgdgddgd| }|| }td	td	d	 d	 td	d	 d	 d
}|D ]}t|| ||  d d }t|| ||  d d }	|  j| ||	7  _t|d ||  d d }
t|d ||  d d }|  j| |
|7  _qd S )NZ	SceneFlowr   )r   DrivingMonkaarh   r   )r   finalr   Zframes_cleanpassZframes_finalpassrN   )r   r   r   r1   rj   r2   Z	disparityz*.pfmr   )r   r   r   r   r   ZpassesZprefix_directoriesr   rS   rT   rU   rV   r   r   r    r   n  s,    zSceneFlowStereo.__init__r!   c                 C   s   t |}t|}d }||fS r*   rX   r[   r   r   r    r;     s    
zSceneFlowStereo._read_disparityr<   c                    s   t tt |S r^   r_   r`   r   r   r    rB     s    zSceneFlowStereo.__getitem__)r   r   Nra   r   r   r   r    r   8  s   8   %r   c                       s   e Zd ZdZdZdeeee dd fddZee	eef dd	d
Z
eee	d e	ejejf f dddZeed fddZ  ZS )SintelStereoa  Sintel `Stereo Dataset <http://sintel.is.tue.mpg.de/stereo>`_.

    The dataset is expected to have the following structure: ::

        root
            Sintel
                training
                    final_left
                        scene1
                            img1.png
                            img2.png
                            ...
                        ...
                    final_right
                        scene2
                            img1.png
                            img2.png
                            ...
                        ...
                    disparities
                        scene1
                            img1.png
                            img2.png
                            ...
                        ...
                    occlusions
                        scene1
                            img1.png
                            img2.png
                            ...
                        ...
                    outofframe
                        scene1
                            img1.png
                            img2.png
                            ...
                        ...

    Args:
        root (string): Root directory where Sintel Stereo is located.
        pass_name (string): The name of the pass to use, either "final", "clean" or "both".
        transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version.
    Tr   N)r   r   r   r   c           	         s   t  || t|ddd t|d }dgdgddgd| }|D ]}t|d | d d	 d
 }t|d | d d	 d
 }|  j| ||7  _t|d d d	 d
 }|  j| |d 7  _qDd S )Nr   )r   r   r   rh   ZSintelr   r   Ztraining_leftrN   rj   _rightrW   )r   r   r   r   rH   r   r:   r   )	r   r   r   r   Z
pass_namesr   rn   ro   rp   r   r   r    r     s    zSintelStereo.__init__r!   c                 C   s   t |}|j}|j}|jj}t|d |j | }t|d |j | }tj|sdtd| dtj|std| d||fS )NZ
occlusionsZ
outofframezOcclusion mask z does not existzOut of frame mask )r   r   r   rH   r   r   r   r5   )r   r"   ZfpathbasenameZscenedirZ	sampledirZocclusion_pathZoutofframe_pathr   r   r    _get_occlussion_mask_paths  s    z'SintelStereo._get_occlussion_mask_pathsrl   c           
      C   s   |d u rdS t jt|t jd}t j|ddd\}}}|d |d  |d  }t |d	}| |\}}t t|d
k}t t|d
k}	t |	|}||fS )Nrl   r      )Zaxis   @   i @  )   r   r   r   )	r?   rs   r   r#   r   re   Z	transposer   logical_and)
r   r"   r\   r   gbZocclued_mask_pathZout_of_frame_mask_pathr]   Zoff_maskr   r   r    r;      s    zSintelStereo._read_disparityr<   c                    s   t tt |S )a  Return example at given index.

        Args:
            index(int): The index of the example to retrieve

        Returns:
            tuple: A 4-tuple with ``(img_left, img_right, disparity, valid_mask)`` is returned.
            The disparity is a numpy array of shape (1, H, W) and the images are PIL images whilst
            the valid_mask is a numpy array of shape (H, W).
        r   r`   r   r   r    rB     s    zSintelStereo.__getitem__)r   N)rD   rE   rF   rG   r>   rH   r	   r   r   r
   r   r   r?   r@   r;   rI   rK   rB   rL   r   r   r   r    r     s   ,(r   c                       s`   e Zd ZdZdeeee dd fddZeee	j
df ddd	Zeed
 fddZ  ZS )
InStereo2ka  `InStereo2k <https://github.com/YuhuaXu/StereoDataset>`_ dataset.

    The dataset is expected to have the following structure: ::

        root
            InStereo2k
                train
                    scene1
                        left.png
                        right.png
                        left_disp.png
                        right_disp.png
                        ...
                    scene2
                    ...
                test
                    scene1
                        left.png
                        right.png
                        left_disp.png
                        right_disp.png
                        ...
                    scene2
                    ...

    Args:
        root (string): Root directory where InStereo2k is located.
        split (string): Either "train" or "test".
        transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version.
    rc   Nrd   c                    s   t  || t|d | }t|ddd t|d d }t|d d }| ||| _t|d d }t|d d	 }| ||| _d S )
Nr   re   rf   rh   rN   zleft.pngz	right.pngzleft_disp.pngzright_disp.png)r   r   r   r   rH   r:   r   r   rw   r   r   r    r   B  s    zInStereo2k.__init__r!   c                 C   s<   t jt|t jd}|d d d d d f d }d }||fS )Nr   g      @r   r[   r   r   r    r;   Q  s    zInStereo2k._read_disparityr<   c                    s   t tt |S r^   r_   r`   r   r   r    rB   X  s    zInStereo2k.__getitem__)rc   Nra   r   r   r   r    r   "  s   r   c                       sr   e Zd ZdZdZdeeee dd fddZee	e
d e
ejejf f d	d
dZeed fddZ  ZS )ETH3DStereoaU  ETH3D `Low-Res Two-View <https://www.eth3d.net/datasets>`_ dataset.

    The dataset is expected to have the following structure: ::

        root
            ETH3D
                two_view_training
                    scene1
                        im1.png
                        im0.png
                        images.txt
                        cameras.txt
                        calib.txt
                    scene2
                        im1.png
                        im0.png
                        images.txt
                        cameras.txt
                        calib.txt
                    ...
                two_view_training_gt
                    scene1
                        disp0GT.pfm
                        mask0nocc.png
                    scene2
                        disp0GT.pfm
                        mask0nocc.png
                    ...
                two_view_testing
                    scene1
                        im1.png
                        im0.png
                        images.txt
                        cameras.txt
                        calib.txt
                    scene2
                        im1.png
                        im0.png
                        images.txt
                        cameras.txt
                        calib.txt
                    ...

    Args:
        root (string): Root directory of the ETH3D Dataset.
        split (string, optional): The dataset split of scenes, either "train" (default) or "test".
        transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version.
    Trc   Nrd   c           	         s   t  || t|ddd t|d }|dkr4dnd}d}t|| d	 d
 }t|| d	 d }| ||| _|dkrtdd | jD | _n"t|| d	 d }| |d | _d S )Nre   rf   rh   ZETH3Drc   Ztwo_view_trainingZtwo_view_testZtwo_view_training_gtrN   rO   rP   rg   c                 s   s   | ]
}d V  qdS rk   r   r+   r   r   r    r.     r/   z'ETH3DStereo.__init__.<locals>.<genexpr>rQ   )	r   r   r   r   rH   r:   r   r3   r   )	r   r   re   r   Zimg_dirZanot_dirrn   ro   rp   r   r   r    r     s    zETH3DStereo.__init__rl   r!   c                 C   sN   |d u rdS t |}t|}t|jd }t|}t|t	}||fS )Nrl   zmask0nocc.png)
rY   r?   rZ   r   r   r   r#   rs   r   r   )r   r"   r\   Z	mask_pathr]   r   r   r    r;     s    

zETH3DStereo._read_disparityr<   c                    s   t tt |S rt   r   r`   r   r   r    rB     s    zETH3DStereo.__getitem__)rc   N)rD   rE   rF   rG   r>   rH   r	   r   r   r   r
   r?   r@   r;   rI   rK   rB   rL   r   r   r   r    r   g  s
   1(r   )-	functoolsr   r   r   r   abcr   r   r   pathlibr   typingr   r   r   r	   r
   r   numpyr?   ZPILr   utilsr   r   r   Zvisionr   r@   rJ   rK   __all__partialrY   r   rM   rb   rv   rx   r   r   r   r   r   r   r   r   r   r    <module>   s:    q@UX  XhpzE