sequence_folders2.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import configparser 
from builtins import open
from builtins import str

import torch.utils.data as data
import numpy as np
from imageio import imread
from path import Path
import random



def load_as_float(path):
    return imread(path).astype(np.float32)


class SequenceFolder(data.Dataset):
    """A sequence data loader where the files are arranged in this way:
        root/scene_1/0000000.jpg
        root/scene_1/0000001.jpg
        ..
        root/scene_1/cam.txt
        root/scene_2/0000000.jpg
        .

        transform functions must take in a list a images and a numpy array (usually intrinsics matrix)
    """

    def __init__(self, root, seed=None, train=True, sequence_length=3, transform=None, target_transform=None):
        np.random.seed(seed)
        random.seed(seed)
        self.root = Path(root)
        scene_list_path = self.root/'train_left.txt' if train else self.root/'val_left.txt'
        self.scenes = [self.root/folder[:-1] for folder in open(scene_list_path)]
        self.transform = transform
        self.crawl_folders(sequence_length)

    def crawl_folders(self, sequence_length):
        sequence_set = []
        demi_length = (sequence_length-1)//2       ##demi=(3-1)//2=1
        shifts = list(range(-demi_length, demi_length + 1))  #list(range(-1,2))=[-1,0,1]
        shifts.pop(demi_length)
        for scene in self.scenes:
            print(scene+'/cam.txt')
            intrinsics = np.genfromtxt(scene +'/cam.txt').astype(np.float32).reshape((3, 3))   ###load camera intrinsics K in every scene!!
            imgs = sorted(scene.files('*.jpg'))
            if len(imgs) < sequence_length:
                continue
            for i in range(demi_length, len(imgs)-demi_length):   ##if this sequence has 100 images. range(0,99)
            #for i in range(0, len(imgs)):
                sample = {'intrinsics': intrinsics, 'tgt': imgs[i], 'ref_imgs': []}     #i start from 1!! if target image is 1, ref_image is[0,1,2]. Correct!
                for j in shifts:
                    sample['ref_imgs'].append(imgs[i+j])
                sequence_set.append(sample)
        ###random.shuffle(sequence_set)
        self.samples = sequence_set
        ###print("samples num:",len(self.samples))
        
        

    def __getitem__(self, index):
        sample = self.samples[index]
        
        tgt_img = load_as_float(sample['tgt'])
        ref_imgs = [load_as_float(ref_img) for ref_img in sample['ref_imgs']]
        if self.transform is not None:
            imgs, intrinsics = self.transform([tgt_img] + ref_imgs, np.copy(sample['intrinsics']))
            tgt_img = imgs[0]
            ref_imgs = imgs[1:]
        else:
            intrinsics = np.copy(sample['intrinsics'])
        return tgt_img, ref_imgs, intrinsics, np.linalg.inv(intrinsics)

    def __len__(self):
        return len(self.samples)