data_utils.py

import numpy as np
import torch
import os
from pathlib import Path

def load_warcraft_data(folder_path,k):
    """
    k (int): hyperparameter, resolution of problem
    # https://www.google.com/url?q=https%3A%2F%2Fedmond.mpdl.mpg.de%2Fdataset.xhtml%3FpersistentId%3Ddoi%3A10.17617%2F3.YJCQ5S
    """
    
    data = {}
    split_types = ['train','val','test']
    data_types = ['maps','vertex_weights','shortest_paths']
    
    for s in split_types:
        data[s] = {}
        for d in data_types:
            # For files split into parts
            data_dir = Path(f"{folder_path}/{k}x{k}")
            search_path = os.path.join(f"{s}_{d}*.npy")
            files = [file for file in data_dir.glob(search_path)]
            container = []
            for file in files:
                container.append(np.load(file))
                data[s][d] = np.concatenate(container, axis = 0)

    return data, split_types, data_types

class WarcraftDataset(torch.utils.data.Dataset):
    def __init__(self, data_dct, data_types):
        """
        data_dct should be one of the nested dictionaries eg data['train']
        """
        self.data = data_dct
        self.data_types = data_types

    def __len__(self):
        return self.data[self.data_types[0]].shape[0]

    def __getitem__(self, idx):
        out_dct = {}
        for d in self.data_types:
            if d == "maps":
                # normalise and transform from [H,W,C] to [C,H,W]
                out_dct[d] = torch.FloatTensor(self.data[d][idx].transpose(2, 0, 1)/255).detach()
            else:
                # flatten
                out_dct[d] = torch.FloatTensor(self.data[d][idx]).reshape(-1)
            
        out_dct['objective'] = (out_dct['vertex_weights'] * out_dct['shortest_paths']).sum()
            
        return out_dct