-
Notifications
You must be signed in to change notification settings - Fork 2
/
dataset.py
59 lines (49 loc) · 2.26 KB
/
dataset.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
import torch
import os
import pandas as pd
from PIL import Image
class VOCDataset(torch.utils.data.Dataset):
def __init__(self, csv_file, img_dir, label_dir, S=7, B=2, C=20, transform=None):
self.annotations = pd.read_csv(csv_file)
self.img_dir = img_dir
self.label_dir = label_dir
self.transform = transform
self.S = S
self.B = B
self.C = C
def __len__(self):
return len(self.annotations)
def __getitem__(self, index):
label_path = os.path.join(self.label_dir, self.annotations.iloc[index, 1])
boxes = []
with open(label_path) as f:
for label in f.readlines():
class_label, x, y, width, height = [
float(x) if float(x) != int(float(x)) else int(x)
for x in label.replace("\n", "").split()
]
boxes.append([class_label, x, y, width, height])
img_path = os.path.join(self.img_dir, self.annotations.iloc[index, 0])
image = Image.open(img_path)
boxes = torch.tensor(boxes) # transform을 하게 되면 augmentation에서 boxex가 필요, 안할 경우 필요없는 작업
if self.transform:
# data augmentation하게 되면 좌표도 같이 수정해야되서 boxes를 입력 받음
image, boxes = self.transform(image, boxes)
label_matrix = torch.zeros((self.S, self.S, self.C + 5 * self.B))
for box in boxes:
class_label, x, y, width, height = box.tolist()
class_label = int(class_label)
i, j = int(self.S * y), int(self.S * x)
x_cell, y_cell = self.S * x - j, self.S * y - i
width_cell, height_cell = (
width * self.S,
height * self.S
)
if label_matrix[i, j, 20] == 0:
label_matrix[i, j, 20] = 1 # obj면 1
box_coordinates = torch.tensor(
[x_cell, y_cell, width_cell, height_cell]
)
label_matrix[i, j, 21:25] = box_coordinates # 좌표 입력
label_matrix[i, j, class_label] = 1 # class 나타내는 20길이 벡터에 해당 클래스만 1 할당 -> one-hot encoding
return image, label_matrix