File size: 2,125 Bytes
803ef9e |
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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import os
import copy
import json
import operator
import numpy as np
from PIL import Image
from glob import glob
from os.path import join
from itertools import chain
from scipy.io import loadmat
from collections import defaultdict
import torch
import torch.utils.data as data
from torchvision import transforms
DATA_ROOTS = 'data/TrafficSign'
# wget https://sid.erda.dk/public/archives/ff17dc924eba88d5d01a807357d6614c/FullIJCNN2013.zip
# unzip FullIJCNN2013.zip
class TrafficSign(data.Dataset):
NUM_CLASSES = 43
def __init__(self, root=DATA_ROOTS, train=True, image_transforms=None):
super().__init__()
self.root = root
self.train = train
self.image_transforms = image_transforms
paths, labels = self.load_images()
self.paths, self.labels = paths, labels
def load_images(self):
split = 'Final_Training'
rs = np.random.RandomState(42)
all_filepaths, all_labels = [], []
for class_i in range(self.NUM_CLASSES):
# class_dir_i = join(self.root, split, 'Images', '{:05d}'.format(class_i))
class_dir_i = join(self.root, '{:02d}'.format(class_i))
image_paths = glob(join(class_dir_i, "*.ppm"))
# train test splitting
image_paths = np.array(image_paths)
num = len(image_paths)
indexer = np.arange(num)
rs.shuffle(indexer)
image_paths = image_paths[indexer].tolist()
if self.train:
image_paths = image_paths[:int(0.8 * num)]
else:
image_paths = image_paths[int(0.8 * num):]
labels = [class_i] * len(image_paths)
all_filepaths.extend(image_paths)
all_labels.extend(labels)
return all_filepaths, all_labels
def __len__(self):
return len(self.paths)
def __getitem__(self, index):
path = self.paths[index]
label = self.labels[index]
image = Image.open(path).convert(mode='RGB')
if self.image_transforms:
image = self.image_transforms(image)
return image, label |