-
Notifications
You must be signed in to change notification settings - Fork 1
/
get_data.py
161 lines (131 loc) · 5.88 KB
/
get_data.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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import time
import os
import h5py
import random
import numpy as np
from PIL import Image
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import Dataset
from torchvision import models, transforms
from superpixel.slic import SLIC
from lib.make_index import make_data_index, default_loader
from lib.image_process import RandomCropPatches, NonOverlappingCropPatches
from lib.utils import mos_rescale
class IQADataset(Dataset):
"""
IQA Dataset
"""
def __init__(self, args, status='train', loader=default_loader):
"""
:param args: arguments of the model
:param status: train/val/test
:param loader: image loader
"""
self.args = args
self.status = status
self.loader = loader
self.train_database = args.train_database
self.test_database = args.test_database
self.image_n_nodes = args.image_n_nodes
self.patch_n_nodes = args.patch_n_nodes
self.region_size = args.region_size
self.ruler = args.ruler
self.iterate = args.iterate
self.patch_size = args.patch_size
self.n_patches_train = args.n_patches_train
# Train
train_Info = h5py.File(args.train_info, 'r')
train_index = train_Info['index']
train_index = train_index[:, 0 % train_index.shape[1]]
train_ref_ids = train_Info['ref_ids'][0, :]
# Test
test_Info = h5py.File(args.test_info, 'r')
test_index = test_Info['index']
test_index = test_index[:, 0 % test_index.shape[1]]
test_ref_ids = test_Info['ref_ids'][0, :]
# Get dataset index
train_index_, test_index_ = [], []
if 'train' in status:
print('The Training Set Index is ', train_index, ' and num of Training index is ', len(train_index))
for i in range(len(train_ref_ids)):
train_index_.append(i)
self.index = train_index_
print("Number of Training Images: {}\n".format(len(self.index)))
self.mos = train_Info['subjective_scores'][0, self.index]
# self.mos_std = train_Info['subjective_scoresSTD'][0, self.index]
im_names = [train_Info[train_Info['im_names'][0, :][i]][()].tobytes()[::2].decode() for i in self.index]
self.label = []
self.im_names = []
self.dis_type = []
for idx in range(len(self.index)):
self.im_names.append(os.path.join(args.train_im_dir, im_names[idx]))
self.label.append(self.mos[idx])
if self.train_database == 'TID2008' or self.train_database == 'TID2013':
self.dis_type.append(int(im_names[idx][4:6]) - 1)
elif self.train_database == 'KADID':
self.dis_type.append(int(im_names[idx][4:6]) - 1)
elif self.train_database == 'CSIQ':
# Distortion Type
if 'AWGN' in im_names[idx]:
self.dis_type.append(0)
elif 'BLUR' in im_names[idx]:
self.dis_type.append(1)
elif 'contrast' in im_names[idx]:
self.dis_type.append(2)
elif 'fnoise' in im_names[idx]:
self.dis_type.append(3)
elif 'JPEG' in im_names[idx]:
self.dis_type.append(4)
elif 'jpeg2000' in im_names[idx]:
self.dis_type.append(5)
elif self.train_database == 'LIVE':
# Distortion Type
if 'jp2k' in im_names[idx]:
self.dis_type.append(0)
elif 'jpeg' in im_names[idx]:
self.dis_type.append(1)
elif 'wn' in im_names[idx]:
self.dis_type.append(2)
elif 'gblur' in im_names[idx]:
self.dis_type.append(3)
elif 'fastfading' in im_names[idx]:
self.dis_type.append(4)
elif self.train_database == 'SIQAD':
loc = im_names[idx].find('_')
self.dis_type.append(int(im_names[idx][loc + 1]) - 1)
elif self.train_database == 'SCID':
self.dis_type.append(int(im_names[idx][6]) - 1)
else:
print('The Testing Set Index is ', test_index, ' and num of test index is ', len(test_index))
for i in range(len(test_ref_ids)):
test_index_.append(i)
self.index = test_index_
print("Number of Testing Images: {}".format(len(self.index)), '\n')
self.mos = test_Info['subjective_scores'][0, self.index]
im_names = [test_Info[test_Info['im_names'][0, :][i]][()].tobytes()[::2].decode() for i in self.index]
self.label = []
self.im_names = []
self.dis_type = []
for idx in range(len(self.index)):
self.im_names.append(os.path.join(args.test_im_dir, im_names[idx]))
self.label.append(self.mos[idx])
self.dis_type.append(0)
def __len__(self):
return len(self.index)
def __getitem__(self, idx):
im = self.loader(self.im_names[idx])
# Get CNN input
if self.status == 'train':
cnn_input, patch, patch_graph = RandomCropPatches(im, args=self.args, transforms=True)
else:
cnn_input, patch, patch_graph = NonOverlappingCropPatches(im, args=self.args, transforms=True)
# Get labels
label = torch.as_tensor([self.label[idx], ])
# Choose whether to use distortion type or distortion level
dis_type = torch.as_tensor([self.dis_type[idx], ])
return patch, patch_graph, cnn_input, label, dis_type