File size: 1,457 Bytes
5f0437a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# Copyright (c) 2023 Image Processing Research Group of University Federico II of Naples ('GRIP-UNINA').
#
# All rights reserved.
# This work should only be used for nonprofit purposes.
#
# By downloading and/or using any of these files, you implicitly agree to all the
# terms of the license, as specified in the document LICENSE.txt
# (included in this package) and online at
# http://www.grip.unina.it/download/LICENSE_OPEN.txt
"""
Created in September 2022
@author: fabrizio.guillaro
"""
from torch.utils.data import Dataset
import random
import numpy as np
import torch
from PIL import Image
class TestDataset(Dataset):
def __init__(self, list_img=None):
self.img_list = list_img
def shuffle(self):
random.shuffle(self.img_list)
def __len__(self):
return len(self.img_list)
def __getitem__(self, index):
assert self.img_list
assert 0 <= index < len(self.img_list), f"Index {index} is not available!"
rgb_path = self.img_list[index]
img_RGB = np.array(Image.open(rgb_path).convert("RGB"))
return torch.tensor(img_RGB.transpose(2, 0, 1), dtype=torch.float) / 256.0, rgb_path
def get_filename(self, index):
item = self.img_list[index]
if isinstance(item, list):
return item[0]
else:
return item |