<filename>ET/alfred/data/semantic_map.py
import os
import json
import pickle
import torch
from torch.functional import split
from torch.utils.data.dataset import Dataset

from io import BytesIO
from tqdm.auto import tqdm


import random
import numpy as np

class AlfredDataset(Dataset):
    def __init__(self, root_folder:str, split="train", grid_size=37) -> None:
        super().__init__()
        self.root_folder = root_folder
        self.name = self.root_folder
        self.split = split
        self.grid_size = grid_size

        self.json_folder = os.path.join(self.root_folder, "jsons")
        self.feat_folder = os.path.join(self.root_folder, "feats")
        
        self.json_paths = []
        self.feat_paths = []

        self.load_json_mask_feat_paths()

        with open('/home/ubuntu/bak/ET/data/neural_slam_map_data.pkl', 'rb') as f:
            self.data = pickle.load(f)
    
    def __len__(self):
        return len(self.json_paths)
    
    def load_json_mask_feat_paths(self):
        for json_file in os.listdir(self.json_folder):
            if json_file.startswith(self.split):
                feat_file = json_file.split(".")[0] + ".pt"
                if not os.path.exists(os.path.join(self.feat_folder, feat_file)): continue
                self.json_paths.append(os.path.join(self.json_folder,json_file))
                self.feat_paths.append(os.path.join(self.feat_folder,feat_file))

    def __getitem__(self, index):
        task_json = random.choice(pickle.load(open(self.json_paths[index], "rb")))
        task_json["dataset_name"] = self.name
        #feat_dict = self.load_features(task_json)
        feat_dict = dict()
        traj_feat = torch.load(self.feat_paths[index])
        traj_feat = traj_feat.view(-1, 4, 512, 7, 7)

        key = self.json_paths[index].split('/')[-1][len(self.split)+1:-4]
        semantic_map, poses, offset = self.data[key]
        t, r = zip(*poses)
        t, r = np.array(t), np.array(r)
        t = (t * 4 - offset) / self.grid_size * 2 # So that t is in [-1, 1]
        poses = np.concatenate([t, r[:,None]], 1)
        assert r.shape[0] == traj_feat.shape[0] + 1
        
        return traj_feat, poses, semantic_map

