# AUTOGENERATED! DO NOT EDIT! File to edit: ../dataset.ipynb.

# %% auto 0
__all__ = ['WxBSDataset', 'EVDDataset']

# %% ../dataset.ipynb 3
import os
import torch
import torchvision as tv
import numpy as np
from PIL import Image

class WxBSDataset(torch.utils.data.Dataset):
    """Wide multiple baselines stereo dataset."""
    urls = {
        'v1.1': [
            'http://cmp.felk.cvut.cz/wbs/datasets/WxBS_v1.1.zip',
            'WxBS_v1.1.zip',
            '66da037e7d20487e86db147307615870'
        ]}
    validation_pairs = ['kyiv_dolltheater2', 'petrzin']
                        
    def __init__(self, root_dir:str, subset: str = 'test',  version='v1.1', download: bool = True):
        """
        Args:
            root_dir (string): Directory with all the images.
            subset (string): can be "val" or "test".
            download (bool): download if not found, default: False
        """
        assert version in self.urls.keys()
        self.root_dir = root_dir
        self.subset = subset
        self.validation_mode = self.subset == 'val'
        self.version = version
        self.data_dir = os.path.join(self.root_dir, version)
        self.data_down = os.path.join(self.root_dir, '{}.zip'.format(version))
        if not self._check_unpacked():
            if download:
                self.download()
            else:
                raise ValueError(f'Data is not in {self.data_dir}, consider set download=True')
        self.index_dataset()
        return
        
    def __len__(self):
        return len(self.pairs)

    def _check_downloaded(self):
        if os.path.exists(self.data_down):
            import hashlib
            md5 = hashlib.md5(self.data_down).hexdigest()
            if md5 == self.urls[self.version][2]:
                return True
        return False

    def _check_unpacked(self):
        return os.path.exists(self.data_dir)

    def download(self) -> None:
        if not self._check_downloaded():
            # download files
            url = self.urls[self.version][0]
            filename = self.urls[self.version][1]
            md5 = self.urls[self.version][2]
            fpath = os.path.join(self.root_dir, filename)
            tv.datasets.phototour.download_url(url, self.root_dir, filename, md5)
            print('# Extracting data {}\n'.format(self.data_down))
            import zipfile
            with zipfile.ZipFile(fpath, 'r') as z:
                z.extractall(self.data_dir)
            os.unlink(fpath)
        return

    def index_dataset(self):
        sets = sorted([x for x in os.listdir(self.data_dir) if os.path.isdir(os.path.join(self.data_dir, x))])
        img_pairs_list = []
        for s in sets:
            if s == '.DS_Store':
                continue
            ss = os.path.join(self.data_dir, s)
            pairs = os.listdir(ss)
            for p in sorted(pairs):
                if p == '.DS_Store':
                    continue
                cur_dir = os.path.join(ss, p)
                if self.validation_mode:
                    if p not in self.validation_pairs:
                        continue
                else:
                    if p in self.validation_pairs:
                        continue
                if os.path.isfile(os.path.join(cur_dir, '01.png')):
                    img_pairs_list.append((os.path.join(cur_dir, '01.png'),
                                           os.path.join(cur_dir, '02.png'),
                                           os.path.join(cur_dir, 'corrs.txt'),
                                           os.path.join(cur_dir, 'crossval_errors.txt')))
                elif os.path.isfile(os.path.join(cur_dir, '01.jpg')):
                    img_pairs_list.append((os.path.join(cur_dir, '01.jpg'),
                                           os.path.join(cur_dir, '02.jpg'),
                                           os.path.join(cur_dir, 'corrs.txt'),
                                           os.path.join(cur_dir, 'crossval_errors.txt')))
                else:
                    continue
        self.pairs = img_pairs_list
        return

    def __getitem__(self, idx):
        imgfname1, imgfname2, pts_fname, err_fname = self.pairs[idx]
        img1 = np.array(Image.open(imgfname1))
        img2 = np.array(Image.open(imgfname2))
        pts = np.loadtxt(pts_fname)
        crossval_errors = np.loadtxt(err_fname)
        pair_name = '/'.join(pts_fname.split('/')[-3:-1])
        out = {'img1': img1,
               'img2': img2,
               'pts': pts,
               'errors': crossval_errors,
               'name': pair_name}
        return out

# %% ../dataset.ipynb 5
class EVDDataset(torch.utils.data.Dataset):
    """Extreme View Dataset."""
    urls = {
        'v1.0': [
            'http://cmp.felk.cvut.cz/wbs/datasets/EVD.zip',
            'EVD.zip',
            '7ce52151abdafa71a609424d09e43075'
        ]}
                        
    def __init__(self, root_dir:str, subset: str = 'test',  version='v1.0', download: bool = True):
        """
        Args:
            root_dir (string): Directory with all the images.
            subset (string): can be "val" or "test".
            download (bool): download if not found, default: False
        """
        assert version in self.urls.keys()
        self.root_dir = root_dir
        self.subset = subset
        self.version = version
        self.data_dir = os.path.join(root_dir, 'EVD')
        self.data_down = self.root_dir
        if not self._check_unpacked():
            if download:
                self.download()
            else:
                raise ValueError(f'Data is not in {self.data_dir}, consider set download=True')
        self.index_dataset()
        return
        
    def __len__(self):
        return len(self.pairs)

    def _check_downloaded(self):
        if os.path.exists(self.data_down):
            import hashlib
            md5 = hashlib.md5(self.data_down).hexdigest()
            if md5 == self.urls[self.version][2]:
                return True
        return False

    def _check_unpacked(self):
        return os.path.exists(self.data_dir)

    def download(self) -> None:
        if not self._check_downloaded():
            # download files
            url = self.urls[self.version][0]
            filename = self.urls[self.version][1]
            md5 = self.urls[self.version][2]
            fpath = os.path.join(self.root_dir, filename)
            tv.datasets.phototour.download_url(url, self.root_dir, filename, md5)
            print('# Extracting data {}\n'.format(self.data_down))
            import zipfile
            with zipfile.ZipFile(fpath, 'r') as z:
                z.extractall(self.data_dir)
            os.unlink(fpath)
        return

    def index_dataset(self):
        sets = sorted([x for x in os.listdir(os.path.join(self.data_dir, '1'))])
        img_pairs_list = []
        for s in sets:
            if s == '.DS_Store':
                continue
            img_pairs_list.append(((os.path.join(self.data_dir, '1', s)),
                                  (os.path.join(self.data_dir, '2', s)),
                                  (os.path.join(self.data_dir, 'h', s.replace('png', 'txt')))))
        self.pairs = img_pairs_list
        return

    def __getitem__(self, idx):
        imgfname1, imgfname2, h_fname = self.pairs[idx]
        img1 = np.array(Image.open(imgfname1))
        img2 = np.array(Image.open(imgfname2))
        with open(h_fname, 'r') as hf:
            lines = hf.readlines()
            H = []
            for l in lines:
                H.append([float(x) for x in l.replace('\t',' ').strip().split(' ') if len(x) > 0])
            H = np.array(H)
            H = H / H[2, 2]
        img1_shape = img1.shape[:2]
        img2_shape = img2.shape[:2]
        
        pair_name = imgfname1.split('/')[-1].split('.')[0]
        out = {'img1': img1,
               'img2': img2,
               'img1_shape': img1_shape,
               'img2_shape': img2_shape,
               'H': H,
               'name': pair_name}
        return out
