| import os
|
|
|
| import h5py
|
| import torch
|
| from torch.utils.data import Dataset
|
|
|
|
|
| class ImagenetResults(Dataset):
|
| def __init__(self, path):
|
| super(ImagenetResults, self).__init__()
|
|
|
| self.path = os.path.join(path, "results.hdf5")
|
| self.data = None
|
|
|
| print("Reading dataset length...")
|
| with h5py.File(self.path, "r") as f:
|
|
|
| self.data_length = len(f["/image"])
|
|
|
| def __len__(self):
|
| return self.data_length
|
|
|
| def __getitem__(self, item):
|
| if self.data is None:
|
| self.data = h5py.File(self.path, "r")
|
|
|
| image = torch.tensor(self.data["image"][item])
|
| vis = torch.tensor(self.data["vis"][item])
|
| target = torch.tensor(self.data["target"][item]).long()
|
|
|
| return image, vis, target
|
|
|
|
|
| if __name__ == "__main__":
|
| import imageio
|
| import numpy as np
|
| from utils import render
|
|
|
| ds = ImagenetResults("../visualizations/fullgrad")
|
| sample_loader = torch.utils.data.DataLoader(ds, batch_size=5, shuffle=False)
|
|
|
| iterator = iter(sample_loader)
|
| image, vis, target = next(iterator)
|
|
|
| maps = (
|
| render.hm_to_rgb(vis[0].data.cpu().numpy(), scaling=3, sigma=1, cmap="seismic")
|
| * 255
|
| ).astype(np.uint8)
|
|
|
|
|
|
|
| print(len(ds))
|
|
|
