-
Notifications
You must be signed in to change notification settings - Fork 20
/
car.py
107 lines (94 loc) · 3.82 KB
/
car.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
import torch
import numpy as np
import os
from PIL import Image, TarIO
import pickle
import tarfile
class car(torch.utils.data.Dataset):
def __init__(self, root, train=True, transform=None):
super(car, self).__init__()
self.root = root
self.train = train
self.transform = transform
if self._check_processed():
print('Train file has been extracted' if self.train else 'Test file has been extracted')
else:
self._extract()
if self.train:
self.train_data, self.train_label = pickle.load(
open(os.path.join(self.root, 'processed/train.pkl'), 'rb')
)
else:
self.test_data, self.test_label = pickle.load(
open(os.path.join(self.root, 'processed/test.pkl'), 'rb')
)
def __len__(self):
return len(self.train_data) if self.train else len(self.test_data)
def __getitem__(self, idx):
if self.train:
img, label = self.train_data[idx], self.train_label[idx]
else:
img, label = self.test_data[idx], self.test_label[idx]
img = Image.fromarray(img)
if self.transform is not None:
img = self.transform(img)
return img, label
def _check_processed(self):
assert os.path.isdir(self.root)
assert os.path.isdir(os.path.join(self.root, 'car_ims'))
assert os.path.isfile(os.path.join(self.root, 'car_nori.list'))
return (os.path.isfile(os.path.join(self.root, 'processed/train.pkl')) and
os.path.isfile(os.path.join(self.root, 'processed/test.pkl')))
def _extract(self):
processed_data_path = os.path.join(self.root, 'processed')
if not os.path.isdir(processed_data_path):
os.mkdir(processed_data_path)
car_nori_path = os.path.join(self.root, 'car_nori.list')
imgs_path = os.path.join(self.root, 'car_ims')
car_nori = open(car_nori_path)
print('Finish loading images.txt and train_test_split.txt')
train_data = []
train_labels = []
test_data = []
test_labels = []
print('Start extract images..')
cnt = 0
train_cnt = 0
test_cnt = 0
for annos in car_nori:
cnt += 1
if cnt < 3: continue
annos_list = annos.split()
image = Image.open(os.path.join(imgs_path, annos_list[1]))
if image.getbands()[0] == 'L':
image = image.convert('RGB')
image_np = np.array(image)
image.close()
if len(annos_list) == 8:
label = int(annos_list[6]) - 1
if annos_list[7] == '1':
train_cnt += 1
train_data.append(image_np)
train_labels.append(label)
else:
test_cnt += 1
test_data.append(image_np)
test_labels.append(label)
else:
tmp = annos_list[5].split(')')
label = int(tmp[1]) - 1
if annos_list[6] == '1':
train_cnt += 1
train_data.append(image_np)
train_labels.append(label)
else:
test_cnt += 1
test_data.append(image_np)
test_labels.append(label)
if (cnt-2)%1000 == 0:
print('{} images have been extracted'.format(cnt-2))
print('Total images: {}, training images: {}. testing images: {}'.format(cnt-2, train_cnt, test_cnt))
pickle.dump((train_data, train_labels),
open(os.path.join(self._root, 'processed/train.pkl'), 'wb'))
pickle.dump((test_data, test_labels),
open(os.path.join(self._root, 'processed/test.pkl'), 'wb'))