scalelsd/ssl/datasets/yorkurban_dataset.py

from torch.utils.data import Dataset
import os
import math
from tqdm import tqdm
from skimage.io import imread
from skimage import color
import PIL
import numpy as np
import h5py
import cv2
import pickle
from .synthetic_util import get_line_heatmap
from torchvision import transforms
import torch
import torch.utils.data.dataloader as torch_loader
# Augmentation libs
from ..config.project_config import Config as cfg
from .transforms import photometric_transforms as photoaug
from .transforms import homographic_transforms as homoaug
# Some visualization tools
from ..misc.visualize_util import plot_junctions, plot_line_segments
# Some data parsing tools
from ..misc.train_utils import parse_h5_data
# Inherit from private dataset
# from dataset.private_dataset import PrivateDataset


# Implements the customized collate_fn for yorkurban dataset
def yorkurban_collate_fn(batch):
    batch_keys = ["image", "junction_map", "valid_mask", "heatmap",
                  "heatmap_pos", "heatmap_neg", "homography"]
    list_keys = ["junctions", "line_map", "line_map_pos", "line_map_neg", "file_key",
                 "aux_junctions", "aux_line_map"]

    outputs = {}
    for data_key in batch[0].keys():
        batch_match = sum([_ in data_key for _ in batch_keys])
        list_match = sum([_ in data_key for _ in list_keys])
        # print(batch_match, list_match)
        if batch_match > 0 and list_match == 0:
            outputs[data_key] = torch_loader.default_collate([b[data_key] for b in batch])
        elif batch_match == 0 and list_match > 0:
            outputs[data_key] = [b[data_key] for b in batch]
        elif batch_match == 0 and list_match == 0:
            continue
        else:
            raise ValueError("[Error] A key matches batch keys and list keys simultaneously.")

    return outputs

# The processed wireframe.
class YorkUrbanDataset(Dataset):
    # Initialize the dataset
    def __init__(self, mode="test", config=None):
        super(YorkUrbanDataset, self).__init__()
        # Check mode => "train", "val", "test
        if not mode in ["test"]:
            raise ValueError("[Error] Unknown mode for york urban dataset. Only 'test' mode is available.")
        self.mode = mode

        self.config = config

        # Get cache setting
        self.dataset_name = self.get_dataset_name()
        self.cache_name = self.get_cache_name()
        self.cache_path = cfg.yorkurban_cache_path

        # Get the filename dataset
        print("[Info] Initializing york urban dataset...")
        self.filename_dataset, self.datapoints = self.get_filename_dataset()
        # Get dataset length
        self.dataset_length = len(self.datapoints)
        
        # Get repeatability evaluation set
        if self.mode == "test" and self.config.get("evaluation", None) is not None:
            # Get the cache name
            tmp = self.cache_name.split(self.mode)
            self.rep_i_cache_name = tmp[0] + self.mode + "_rep_i" + tmp[1]
            self.rep_v_cache_name = tmp[0] + self.mode + "_rep_v" + tmp[1]

            # Get the repeatability config
            self.rep_config = self.config["evaluation"]["repeatability"]

            self.rep_eval_dataset = self.construct_rep_eval_dataset()
            self.rep_eval_datapoints = self.get_rep_eval_datapoints()

        # Print some info
        print("[Info] Successfully initialized dataset")
        print("\t Name: yorkurban")
        print("\t Mode: %s" %(self.mode))
        print("\t Gt: %s" %(self.config.get("gt_source_%s"%(self.mode), "official")))
        print("\t Counts: %d" %(self.dataset_length))
        print("----------------------------------------")
    
    def get_filename_dataset(self):
        # Get the path to the dataset
        if self.mode == "train":
            raise NotImplementedError
        elif self.mode == "test":
            dataset_path = os.path.join(cfg.yorkurban_dataroot)
        # Get paths to all image files
        folder_lst = sorted([os.path.join(dataset_path, _) for _ in os.listdir(dataset_path) \
            if os.path.isdir(os.path.join(dataset_path, _))])
        folder_lst = folder_lst[:-1]
        #folder_lst = [f for f in folder_lst if f.startswith('P')]
        image_paths = []
        for folder in folder_lst:
            image_path = [os.path.join(folder, _) for _ in os.listdir(folder) \
                if os.path.splitext(_)[-1] == ".jpg" or os.path.splitext(_)[-1] == ".png"]
            image_paths += image_path

        # Verify all the images and labels exist
        for idx in range(len(image_paths)):
            image_path = image_paths[idx]
            if not os.path.exists(image_path):
                raise ValueError("[Error] The image does not exist. %s"%(image_path))
        
        # Construct the filename dataset
        num_pad = int(math.ceil(math.log10(len(image_paths))) + 1)
        filename_dataset = {}
        for idx in range(len(image_paths)):
            # Get the file key
            key = self.get_padded_filename(num_pad, idx)

            filename_dataset[key] = {
                "image": image_paths[idx]
            }
        
        # Get the datapoints
        datapoints = list(sorted(filename_dataset.keys()))

        return filename_dataset, datapoints
    
    # Get the padded filename using adaptive padding
    @staticmethod
    def get_padded_filename(num_pad, idx):
        file_len = len("%d" % (idx))
        filename = "0" * (num_pad - file_len) + "%d" % (idx)

        return filename
    
    # Get dataset name from dataset config / default config
    def get_dataset_name(self):
        if self.config["dataset_name"] is None:
            dataset_name = "yorkurban_dataset" + f"_{self.mode}"
        else:
            dataset_name = self.config["dataset_name"] + f"_{self.mode}" 

        return dataset_name
    
    # Get cache name from dataset config / default config
    def get_cache_name(self):
        if self.config["dataset_name"] is None:
            dataset_name = "yorkurban_dataset" + f"_{self.mode}"
        else:
            dataset_name = self.config["dataset_name"] + f"_{self.mode}" 
        # Compose cache name
        cache_name = dataset_name + "_cache.pkl"

        return cache_name
    
    ###########################################
    ## Repeatability evaluation related APIs ##
    ###########################################
    # Construct repeatability evaluation dataset (from scratch or from cache)
    def construct_rep_eval_dataset(self):
        rep_eval_dataset = {}
        # Check if viewpoint and illumination cache exists
        if self.rep_config["photometric"]["enable"]:
            if self._check_rep_eval_dataset_cache(split="i"):
                print("\t Found repeatability illumination cache %s at %s"%(self.rep_i_cache_name, self.cache_path))
                print("\t Load repeatability illumination cache...")
                rep_i_keymap, rep_i_dataset_name = self.get_rep_eval_dataset_from_cache(split="i")
            else:
                print("\t Can't find repeatability illumination cache ...")
                print("\t Create repeatability illumination dataset from scratch...")
                rep_i_keymap, rep_i_dataset_name = self.get_rep_eval_dataset(split="i")
                print("\t Create filename dataset cache...")
                self.create_rep_eval_dataset_cache("i", rep_i_keymap, rep_i_dataset_name)
        else:
            rep_i_keymap = None
            rep_i_dataset_name = None
        
        rep_eval_dataset["illumination"] = {
            "keymap": rep_i_keymap,
            "dataset_name": rep_i_dataset_name
        }
        
        if self.rep_config["homographic"]["enable"]:
            if self._check_rep_eval_dataset_cache(split="v"):
                print("\t Found repeatability viewpoint cache %s at %s"%(self.rep_v_cache_name, self.cache_path))
                print("\t Load repeatability viewpoint cache...")
                rep_v_keymap, rep_v_dataset_name = self.get_rep_eval_dataset_from_cache(split="v")
            else:
                print("\t Can't find repeatability viewpoint cache ...")
                print("\t Create repeatability viewpoint dataset from scratch...")
                rep_v_keymap, rep_v_dataset_name = self.get_rep_eval_dataset(split="v")
                print("\t Create filename dataset cache...")
                self.create_rep_eval_dataset_cache("v", rep_v_keymap, rep_v_dataset_name)
        else:
            rep_v_keymap = None
            rep_v_dataset_name = None
        
        rep_eval_dataset["viewpoint"] = {
            "keymap": rep_v_keymap,
            "dataset_name": rep_v_dataset_name
        }

        return rep_eval_dataset
    
    # Create filename dataset cache for faster initialization
    def create_rep_eval_dataset_cache(self, split, keymap, dataset_name):
        # Check cache path exists
        if not os.path.exists(self.cache_path):
            os.makedirs(self.cache_path)

        if split == "i":
            cache_file_path = os.path.join(self.cache_path, self.rep_i_cache_name)
        elif split == "v":
            cache_file_path = os.path.join(self.cache_path, self.rep_v_cache_name)
        else:
            raise ValueError("[Error] Unknown split for repeatability evaluation.")

        data = {
            "keymap": keymap,
            "dataset_name": dataset_name
        }
        with open(cache_file_path, "wb") as f:
            pickle.dump(data, f, pickle.HIGHEST_PROTOCOL)
    
    # Get filename dataset from cache
    def get_rep_eval_dataset_from_cache(self, split):
        # Load from pkl cache
        if split == "i":
            cache_file_path = os.path.join(self.cache_path, self.rep_i_cache_name)
        elif split == "v":
            cache_file_path = os.path.join(self.cache_path, self.rep_v_cache_name)
        else:
            raise ValueError("[Error] Unknown split for repeatability evaluation.")

        with open(cache_file_path, "rb") as f:
            data = pickle.load(f)
        
        return data["keymap"], data["dataset_name"]
    
    # Initialize the repeatability evaluation dataset from scratch
    def get_rep_eval_dataset(self, split):
        image_shape = self.config["preprocessing"]["resize"]

        # Initialize the illumination set
        if split == "i":
            # Set the random seed before continuing
            seed = self.rep_config["seed"]
            np.random.seed(seed)
            torch.manual_seed(seed)

            raise NotImplementedError
        
        # Initialize the viewpoint set
        elif split == "v":
            # Set the random seed before continuing
            seed = self.rep_config["seed"]
            np.random.seed(seed)
            torch.manual_seed(seed)

            v_keymap = {}
            # Get the name for the output h5 dataset
            v_dataset_name = self.rep_v_cache_name.split(".pkl")[0] + ".h5"
            v_dataset_path = os.path.join(self.cache_path, v_dataset_name)
            with h5py.File(v_dataset_path, "w") as f:
                # Iterate through all the file_key in test set
                for idx, key in enumerate(tqdm(list(self.filename_dataset.keys()), ascii=True)):
                    # Sample N random homography
                    file_key_lst = []
                    for i in range(self.rep_config["homographic"]["num_samples"]):
                        file_key = key + "_" + str(i)

                        # Sample a random homography
                        homo_mat, _ = homoaug.sample_homography(image_shape, 
                                **self.rep_config["homographic"]["params"])
                        
                        file_key_lst.append(file_key)
                        f.create_dataset(file_key, data=homo_mat, compression="gzip")
                    
                    v_keymap[key] = file_key_lst
            
            return v_keymap, v_dataset_name
        
        else:
            raise ValueError("[Error] Unknow split for repeatability evaluation.")
    
    # Convert ref image and warped images to list of evaluation pairs
    def get_rep_eval_datapoints(self):
        datapoints = {
            "illumination": [],
            "viewpoint": []
        }

        # Iterate through all the ref image
        if self.rep_eval_dataset["illumination"]["keymap"] is not None:
            for ref_key in sorted(self.rep_eval_dataset["illumination"]["keymap"].keys()):
                pair_lst = [[ref_key, _] for _ in self.rep_eval_dataset["illumination"]["keymap"][ref_key]]
                datapoints["illumination"] += pair_lst
        
        if self.rep_eval_dataset["viewpoint"]["keymap"] is not None:
            for ref_key in sorted(self.rep_eval_dataset["viewpoint"]["keymap"].keys()):
                pair_lst = [[ref_key, _] for _ in self.rep_eval_dataset["viewpoint"]["keymap"][ref_key]]
                datapoints["viewpoint"] += pair_lst
        
        return datapoints
    
    # Check if the repeatability cache dataset exists
    def _check_rep_eval_dataset_cache(self, split):
        if split == "i":
            cache_file_path = os.path.join(self.cache_path, self.rep_i_cache_name)
        else:
            cache_file_path = os.path.join(self.cache_path, self.rep_v_cache_name)
        
        return os.path.exists(cache_file_path)
    
    ###########################################
    ## Repeatability evaluation related APIs ##
    ###########################################
    # Get the corresponding data according to the "index in rep_eval_datapoints".
    def get_rep_eval_data(self, split, idx):
        assert split in ["viewpoint", "illumination"]
        datapoint = self.rep_eval_datapoints[split][idx]

        # Get reference image
        ref_key = datapoint[0]
        # Get the data paths
        data_path = self.filename_dataset[ref_key]
        # Read in the image and npz labels (but haven't applied any transform)
        image = imread(data_path["image"])

        # Resize the image before photometric and homographical augmentations
        image_size = image.shape[:2]
        if not(list(image_size) == self.config["preprocessing"]["resize"]):
            # Resize the image and the point location.
            size_old = list(image.shape)[:2] # Only H and W dimensions

            image = cv2.resize(image, tuple(self.config['preprocessing']['resize'][::-1]),
                               interpolation=cv2.INTER_LINEAR)
            image = np.array(image, dtype=np.uint8)
        
        # Optionally convert the image to grayscale
        if self.config["gray_scale"]:
            image = (color.rgb2gray(image) *255.).astype(np.uint8)
        
        image_transform = photoaug.normalize_image()
        image = image_transform(image)

        # Get target image
        if split == "viewpoint":
            target_key = datapoint[1]
            dataset_path = os.path.join(self.cache_path, self.rep_eval_dataset[split]["dataset_name"])
            with h5py.File(dataset_path, "r") as f:
                homo_mat = np.array(f[target_key])
            
            # Warp the image
            target_size = (image.shape[1], image.shape[0])
            target_image = cv2.warpPerspective(image, homo_mat, target_size,
                                            flags=cv2.INTER_LINEAR)

        else:
            raise NotImplementedError

        # Convert to tensor and return the results
        to_tensor = transforms.ToTensor()

        return {
            "ref_image": to_tensor(image),
            "ref_key": ref_key,
            "target_image": to_tensor(target_image),
            "target_key": target_key,
            "homo_mat": homo_mat
        }
    
    ############################################
    ## Pytorch and preprocessing related APIs ##
    ############################################
    # Get the length of the dataset
    def __len__(self):
        return self.dataset_length

    # Get data from the information from filename dataset
    @staticmethod
    def get_data_from_path(data_path):
        output = {}

        # Get image data
        image_path = data_path["image"]
        image = imread(image_path)
        output["image"] = image
        
        return output
    
    # The test preprocessing
    def test_preprocessing(self, data, numpy=False):
        # Fetch the corresponding entries
        image = data["image"]
        image_size = image.shape[:2]

        # Resize the image before photometric and homographical augmentations
        if not(list(image_size) == self.config["preprocessing"]["resize"]):
            # Resize the image and the point location.
            size_old = list(image.shape)[:2] # Only H and W dimensions

            image = cv2.resize(image, tuple(self.config['preprocessing']['resize'][::-1]),
                               interpolation=cv2.INTER_LINEAR)
            image = np.array(image, dtype=np.uint8)

        # Optionally convert the image to grayscale
        if self.config["gray_scale"]:
            image = (color.rgb2gray(image) *255.).astype(np.uint8)

        # Still need to normalize image
        image_transform = photoaug.normalize_image()
        image = image_transform(image)

        # Update image size
        image_size = image.shape[:2]
        valid_mask = np.ones(image_size)

        # Convert to tensor and return the results
        to_tensor = transforms.ToTensor()
        if not numpy:
            return {
                "image": to_tensor(image),
                "valid_mask": to_tensor(valid_mask).to(torch.int32)
            }
        else:
            return {
                "image": image,
                "valid_mask": valid_mask.astype(np.int32)
            }
    
    # Define the getitem method
    def __getitem__(self, idx):
        """Return data
        file_key: str, keys used to retrieve certain data from the filename dataset.
        image: torch.float, C*H*W range 0~1,
        valid_mask: torch.int32, 1*H*W range 0 or 1
        """
        # Get the corresponding datapoint and get contents from filename dataset
        file_key = self.datapoints[idx]
        data_path = self.filename_dataset[file_key]
        # Read in the image and npz labels (but haven't applied any transform)
        data = self.get_data_from_path(data_path)

        # Perform transform and augmentation
        if self.mode == "train" or self.config["add_augmentation_to_all_splits"]:
            raise NotImplementedError
        else:
            data = self.test_preprocessing(data)
        
        # Add file key to the output
        data["file_key"] = file_key
        
        return data
    
if __name__ == "__main__":
    import sys
    import yaml
    import matplotlib
    import matplotlib.pyplot as plt
    plt.switch_backend("TkAgg")
    from torch.utils.data import DataLoader
    sys.path.append("../")

    # Load configuration file
    with open("./config/yorkurban_dataset_config.yaml", "r") as f:
        config = yaml.safe_load(f)
    
    config["add_augmentation_to_all_splits"] = False

    # Initialize the dataset
    test_dataset = YorkUrbanDataset(mode="test", config=config)
    import ipdb; ipdb.set_trace()