a
    h                     @   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 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Optional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f 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   S/var/www/auris/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   r9   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(   )r4   r   )r   r   r   r   __len__   s    zStereoMatchingDataset.__len__)N)N)__name__
__module____qualname____doc__r<   r	   strr   r   r   r   r
   r%   r1   tupler8   r   r=   r>   r9   intT1T2r@   rA   __classcell__r   r   r   r   r      s   $	 &#r   c                       sf   e Zd ZdZdeeef 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 )CarlaStereoaz  
    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 (str or ``pathlib.Path``): 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   rF   r8   r   r   )	r   r   r   left_image_patternright_image_patternr?   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   r9      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   rI   r   r@   r   r;   r   r   r   r@      s    zCarlaStereo.__getitem__)NrB   rC   rD   rE   r	   rF   r   r   r   r   rG   r=   r>   r9   rH   rI   r@   rK   r   r   r   r   rL      s   $rL   c                       sp   e Zd ZdZdZdeeef 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 (str or ``pathlib.Path``): 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 )Nrd   rb   testZvalid_valuesZ	Kitti2012ingZ	colored_0z*_10.pngZ	colored_1rb   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   rF   r8   r   r   r1   )r   r   rd   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 Nrk   g      p@r=   asarrayr
   r!   rZ   r   r   r   r9     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   r@     s    zKitti2012Stereo.__getitem__)rb   NrB   rC   rD   rE   r<   r	   rF   r   r   r   r   rG   r=   r>   r9   rH   rI   r@   rK   r   r   r   r   ra      s
   &&ra   c                       sp   e Zd ZdZdZdeeef 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 )Kitti2015StereoaM  
    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 (str or ``pathlib.Path``): 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.
    Trb   Nrc   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 )Nrd   re   rg   Z	Kitti2015rh   Zimage_2ri   Zimage_3rb   Z
disp_occ_0Z
disp_occ_1c                 s   s   | ]
}d V  qdS rj   r   r)   r   r   r   r,   Z  r-   z+Kitti2015Stereo.__init__.<locals>.<genexpr>rl   r   r   rd   r   rm   rn   rT   rU   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 rp   rq   rZ   r   r   r   r9   \  s    zKitti2015Stereo._read_disparityr:   c                    s   t tt |S rs   r^   r_   r   r   r   r@   g  s    zKitti2015Stereo.__getitem__)rb   Nrt   r   r   r   r   ru     s
   )&ru   c                	       s   e Zd ZdZg dg dg ddZdZdeeef 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eef d
dddZeed fddZ  ZS )Middlebury2014StereoaZ	  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 (str or ``pathlib.Path``): 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)rb   
additionalrf   Trb   perfectFN)r   rd   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 )Nrd   )rb   rf   rx   rg   rz   )ry   	imperfectbothNrf   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)Nry   r}   r~   rM   rN   rO   c                 s   s   | ]
}d V  qdS rj   r   r)   r   r   r   r,     r-   z	disp0.pfmz	disp1.pfm)r   r   r   rd   r5   _download_datasetr   ospathexistsr3   splitsanylistdirrF   r   r8   r1   r   r{   )r   r   rd   rz   r{   r   r|   Zcalibrartion_suffixesZcalibration_suffixZscene_patternrm   rn   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.
        rO   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{   parentr1   filterappendrandomchoicer   r%   )r   r    Zambient_file_pathsr   r   r   r%   "  s    


zMiddlebury2014Stereo._read_imgrk   c                 C   sB   |d u rdS t |}t|}d||tjk< |dkd}||fS )Nrk   r   )rX   r=   rY   infZsqueezerZ   r   r   r   r9   8  s    
z$Middlebury2014Stereo._read_disparity)r   r   c                    s<  d}t  d  | j}|dkr| j| D ]Z} | }dD ]H}| d| }| d| d}||  s8t|| dt|dd	 q8q(nt d  t fd
d| jd D r8d}	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   rf   )ry   r}   -/z.zipT)urlfilenamedownload_rootremove_finishedc                 3   s    | ]}|t  d  vV  qdS )rf   N)r   r   )r*   r   r   r   r   r,   Y  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   rd   r   r   r   rF   r   makedirsr   walkshutilmovermtree)r   r   base_urlZ
split_nameZsplit_sceneZ
split_rootrz   Z
scene_nameZ	scene_urlZtest_set_urlZ	scene_dirZscene_namesr+   r   Zscene_dst_dirZscene_src_dirr   r   r   r   C  s6    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   rJ   r   r@   r_   r   r   r   r@   i  s    z Middlebury2014Stereo.__getitem__)rb   ry   FNF)rB   rC   rD   rE   r   r<   r	   rF   r   r   boolr   r   r
   r%   rG   r=   r>   r9   r   rH   rJ   r@   rK   r   r   r   r   rw   w  s0   9/     
A(&rw   c                       sj   e Zd ZdZdZdeeef 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.pngrQ   )r   r   r   dirsr   rR   rS   r?   rT   rU   rV   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=   rr   r
   r!   float32rZ   r   r   r   r9     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   r@     s    zCREStereo.__getitem__)Nrt   r   r   r   r   r   w  s   ) 
r   c                       sh   e Zd ZdZdeeef 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 (str or ``pathlib.Path``): 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~   rg   r   r   rM   )r   r   z
*.left.jpgz*.right.jpgz*.left.depth.pngz*.right.depth.pngr   r   r   r   rF   r   r8   r   )r   r   r   r   variantsZsplit_prefixr   rm   rn   rT   rU   r   r   r   r     s&    
zFallingThingsStereo.__init__r   c                 C   s   t t|}t|jd }t|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.jsonZcamera_settingsr   Zintrinsic_settingsZfx)   d   )
r=   rr   r
   r!   r   r   jsonloadastyper   )r   r    depthZcamera_settings_pathfZ
intrinsicsZfocalZbaselineZpixel_constantr[   r\   r   r   r   r9     s    

z#FallingThingsStereo._read_disparityr:   c                    s   t tt |S r]   r^   r_   r   r   r   r@   (  s    zFallingThingsStereo.__getitem__)r   Nr`   r   r   r   r   r     s   ,&r   c                       sj   e Zd ZdZdeeef 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 (str or ``pathlib.Path``): 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Monkaarg   r   )r   finalr~   Zframes_cleanpassZframes_finalpassrM   )r   r   r   r/   ri   r0   Z	disparityz*.pfmr   )r   r   r   r   r   ZpassesZprefix_directoriesr   rR   rS   rT   rU   r   r   r   r   m  s,    zSceneFlowStereo.__init__r   c                 C   s   t |}t|}d }||fS r(   rW   rZ   r   r   r   r9     s    
zSceneFlowStereo._read_disparityr:   c                    s   t tt |S r]   r^   r_   r   r   r   r@     s    zSceneFlowStereo.__getitem__)r   r   Nr`   r   r   r   r   r   7  s   8   
%r   c                       s   e Zd ZdZdZdeeef 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 (str or ``pathlib.Path``): 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~   rg   ZSintelr   r   Ztraining_leftrM   ri   _rightrV   )r   r   r   r   rF   r   r8   r   )	r   r   r   r   Z
pass_namesr   rm   rn   ro   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   rF   r   r   r   r3   )r   r    ZfpathbasenameZscenedirZ	sampledirZocclusion_pathZoutofframe_pathr   r   r   _get_occlussion_mask_paths  s    z'SintelStereo._get_occlussion_mask_pathsrk   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 )Nrk   r      )Zaxis   @   i @  )   r   r   r   )	r=   rr   r
   r!   r   rd   Z	transposer   logical_and)
r   r    r[   rgbZocclued_mask_pathZout_of_frame_mask_pathr\   Zoff_maskr   r   r   r9     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   r@     s    zSintelStereo.__getitem__)r   N)rB   rC   rD   rE   r<   r	   rF   r   r   r   r   rG   r   r=   r>   r9   rH   rJ   r@   rK   r   r   r   r   r     s   ,&(r   c                       sh   e Zd ZdZdeeef 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 (str or ``pathlib.Path``): 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.
    rb   Nrc   c                    s   t  || t|d | }t|ddd t|d d }t|d d }| ||| _t|d d }t|d d	 }| ||| _d S )
Nr   rd   re   rg   rM   zleft.pngz	right.pngzleft_disp.pngzright_disp.png)r   r   r   r   rF   r8   r   r   rv   r   r   r   r   A  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   rZ   r   r   r   r9   P  s    zInStereo2k._read_disparityr:   c                    s   t tt |S r]   r^   r_   r   r   r   r@   W  s    zInStereo2k.__getitem__)rb   Nr`   r   r   r   r   r   !  s   &r   c                       sz   e Zd ZdZdZdeeef 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 )ETH3DStereoaf  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 (str or ``pathlib.Path``): 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.
    Trb   Nrc   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 )Nrd   re   rg   ZETH3Drb   Ztwo_view_trainingZtwo_view_testZtwo_view_training_gtrM   rN   rO   rf   c                 s   s   | ]
}d V  qdS rj   r   r)   r   r   r   r,     r-   z'ETH3DStereo.__init__.<locals>.<genexpr>rP   )	r   r   r   r   rF   r8   r   r1   r   )	r   r   rd   r   Zimg_dirZanot_dirrm   rn   ro   r   r   r   r     s    zETH3DStereo.__init__rk   r   c                 C   sN   |d u rdS t |}t|}t|jd }t|}t|t	}||fS )Nrk   zmask0nocc.png)
rX   r=   rY   r   r   r
   r!   rr   r   r   )r   r    r[   Z	mask_pathr\   r   r   r   r9     s    

zETH3DStereo._read_disparityr:   c                    s   t tt |S rs   r   r_   r   r   r   r@     s    zETH3DStereo.__getitem__)rb   N)rB   rC   rD   rE   r<   r	   rF   r   r   r   r   rG   r=   r>   r9   rH   rJ   r@   rK   r   r   r   r   r   f  s
   1&(r   ),	functoolsr   r   r   r   abcr   r   r   pathlibr   typingr   r   r   r	   numpyr=   ZPILr
   utilsr   r   r   Zvisionr   rG   r>   rI   rJ   __all__partialrX   r   rL   ra   ru   rw   r   r   r   r   r   r   r   r   r   r   <module>   s:   q@UX  XhpzE