fix dataset
This commit is contained in:
parent
243bd9944d
commit
2365c70e4a
|
@ -72,6 +72,7 @@ class PGDataSet(Dataset):
|
||||||
def __getitem__(self, idx):
|
def __getitem__(self, idx):
|
||||||
file_idx = self.data_idx_order_list[idx]
|
file_idx = self.data_idx_order_list[idx]
|
||||||
data_line = self.data_lines[file_idx]
|
data_line = self.data_lines[file_idx]
|
||||||
|
img_id = 0
|
||||||
try:
|
try:
|
||||||
data_line = data_line.decode('utf-8')
|
data_line = data_line.decode('utf-8')
|
||||||
substr = data_line.strip("\n").split(self.delimiter)
|
substr = data_line.strip("\n").split(self.delimiter)
|
||||||
|
@ -79,9 +80,10 @@ class PGDataSet(Dataset):
|
||||||
label = substr[1]
|
label = substr[1]
|
||||||
img_path = os.path.join(self.data_dir, file_name)
|
img_path = os.path.join(self.data_dir, file_name)
|
||||||
if self.mode.lower() == 'eval':
|
if self.mode.lower() == 'eval':
|
||||||
img_id = int(data_line.split(".")[0][7:])
|
try:
|
||||||
else:
|
img_id = int(data_line.split(".")[0][7:])
|
||||||
img_id = 0
|
except:
|
||||||
|
img_id = 0
|
||||||
data = {'img_path': img_path, 'label': label, 'img_id': img_id}
|
data = {'img_path': img_path, 'label': label, 'img_id': img_id}
|
||||||
if not os.path.exists(img_path):
|
if not os.path.exists(img_path):
|
||||||
raise Exception("{} does not exist!".format(img_path))
|
raise Exception("{} does not exist!".format(img_path))
|
||||||
|
|
Loading…
Reference in New Issue