/
data_list.py
112 lines (92 loc) · 3.38 KB
/
data_list.py
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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
import torch
import numpy as np
import random
from PIL import Image
from torch.utils.data import Dataset
import os
import os.path
import cv2
import torchvision
def make_dataset(image_list, labels):
if labels:
len_ = len(image_list)
images = [(image_list[i].strip(), labels[i, :]) for i in range(len_)]
else:
if len(image_list[0].split()) > 2:
images = [(val.split()[0], np.array([int(la) for la in val.split()[1:]])) for val in image_list]
else:
images = [(val.split()[0], int(val.split()[1])) for val in image_list]
return images
def rgb_loader(path):
with open(path, 'rb') as f:
with Image.open(f) as img:
return img.convert('RGB')
def l_loader(path):
with open(path, 'rb') as f:
with Image.open(f) as img:
return img.convert('L')
class ImageList(Dataset):
def __init__(self, image_list, labels=None, transform=None, target_transform=None, mode='RGB'):
imgs = make_dataset(image_list, labels)
if len(imgs) == 0:
raise(RuntimeError("Found 0 images in subfolders of: " + root + "\n"
"Supported image extensions are: " + ",".join(IMG_EXTENSIONS)))
self.imgs = imgs
self.transform = transform
self.target_transform = target_transform
if mode == 'RGB':
self.loader = rgb_loader
elif mode == 'L':
self.loader = l_loader
def __getitem__(self, index):
path, target = self.imgs[index]
img = self.loader(path)
if self.transform is not None:
img = self.transform(img)
if self.target_transform is not None:
target = self.target_transform(target)
return img, target
def __len__(self):
return len(self.imgs)
class ImageList_idx(Dataset):
def __init__(self, image_list, labels=None, transform=None, target_transform=None, transform1=None,mode='RGB'):
imgs = make_dataset(image_list, labels)
if len(imgs) == 0:
raise(RuntimeError("Found 0 images in subfolders of: " + root + "\n"
"Supported image extensions are: " + ",".join(IMG_EXTENSIONS)))
self.imgs = imgs
self.transform = transform
self.target_transform = target_transform
self.transform1=transform1
if mode == 'RGB':
self.loader = rgb_loader
elif mode == 'L':
self.loader = l_loader
def __getitem__(self, index):
path, target = self.imgs[index]
img = self.loader(path)
if self.transform is not None:
img1 = self.transform(img)
if self.transform1 is not None:
img2=self.transform1(img)
if self.target_transform is not None:
target = self.target_transform(target)
if self.transform1 is not None:
return [img1,img2], target, index
else:
return img1, target, index,path
def __len__(self):
return len(self.imgs)
class Listset(Dataset):
def __init__(self,image_list,lists,uns=False):
self.image_list = image_list
self.lists = lists
self.uns=uns
def __getitem__(self, index):
ids=self.lists[index]
if self.uns:
return self.image_list.__getitem__(ids),index
else:
return self.image_list.__getitem__(ids)
def __len__(self):
return len(self.lists)