-
Notifications
You must be signed in to change notification settings - Fork 0
/
data_helper.py
146 lines (118 loc) · 5.62 KB
/
data_helper.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
import os
from PIL import Image
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from helper import convert_map_to_lane_map, convert_map_to_road_map
NUM_SAMPLE_PER_SCENE = 126
NUM_IMAGE_PER_SAMPLE = 6
image_names = [
'CAM_FRONT_LEFT.jpeg',
'CAM_FRONT.jpeg',
'CAM_FRONT_RIGHT.jpeg',
'CAM_BACK_LEFT.jpeg',
'CAM_BACK.jpeg',
'CAM_BACK_RIGHT.jpeg',
]
# The dataset class for unlabeled data.
class UnlabeledDataset(torch.utils.data.Dataset):
def __init__(self, image_folder, scene_index, first_dim, transform):
"""
Args:
image_folder (string): the location of the image folder
scene_index (list): a list of scene indices for the unlabeled data
first_dim ({'sample', 'image'}):
'sample' will return [batch_size, NUM_IMAGE_PER_SAMPLE, 3, H, W]
'image' will return [batch_size, 3, H, W] and the index of the camera [0 - 5]
CAM_FRONT_LEFT: 0
CAM_FRONT: 1
CAM_FRONT_RIGHT: 2
CAM_BACK_LEFT: 3
CAM_BACK.jpeg: 4
CAM_BACK_RIGHT: 5
transform (Transform): The function to process the image
"""
self.image_folder = image_folder
self.scene_index = scene_index
self.transform = transform
assert first_dim in ['sample', 'image']
self.first_dim = first_dim
def __len__(self):
if self.first_dim == 'sample':
return self.scene_index.size * NUM_SAMPLE_PER_SCENE
elif self.first_dim == 'image':
return self.scene_index.size * NUM_SAMPLE_PER_SCENE * NUM_IMAGE_PER_SAMPLE
def __getitem__(self, index):
if self.first_dim == 'sample':
scene_id = self.scene_index[index // NUM_SAMPLE_PER_SCENE]
sample_id = index % NUM_SAMPLE_PER_SCENE
sample_path = os.path.join(self.image_folder, f'scene_{scene_id}', f'sample_{sample_id}')
images = []
for image_name in image_names:
image_path = os.path.join(sample_path, image_name)
image = Image.open(image_path)
images.append(self.transform(image))
image_tensor = torch.stack(images)
return image_tensor
elif self.first_dim == 'image':
scene_id = self.scene_index[index // (NUM_SAMPLE_PER_SCENE * NUM_IMAGE_PER_SAMPLE)]
sample_id = (index % (NUM_SAMPLE_PER_SCENE * NUM_IMAGE_PER_SAMPLE)) // NUM_IMAGE_PER_SAMPLE
image_name = image_names[index % NUM_IMAGE_PER_SAMPLE]
image_path = os.path.join(self.image_folder, f'scene_{scene_id}', f'sample_{sample_id}', image_name)
image = Image.open(image_path)
return self.transform(image), index % NUM_IMAGE_PER_SAMPLE
# The dataset class for labeled data.
class LabeledDataset(torch.utils.data.Dataset):
def __init__(self, image_folder, annotation_file, scene_index, transform, extra_info=True):
"""
Args:
image_folder (string): the location of the image folder
annotation_file (string): the location of the annotations
scene_index (list): a list of scene indices for the unlabeled data
transform (Transform): The function to process the image
extra_info (Boolean): whether you want the extra information
"""
self.image_folder = image_folder
self.annotation_dataframe = pd.read_csv(annotation_file)
self.scene_index = scene_index
self.transform = transform
self.extra_info = extra_info
def __len__(self):
return self.scene_index.size * NUM_SAMPLE_PER_SCENE
def __getitem__(self, index):
scene_id = self.scene_index[index // NUM_SAMPLE_PER_SCENE]
sample_id = index % NUM_SAMPLE_PER_SCENE
sample_path = os.path.join(self.image_folder, f'scene_{scene_id}', f'sample_{sample_id}')
images = []
for image_name in image_names:
image_path = os.path.join(sample_path, image_name)
image = Image.open(image_path)
images.append(self.transform(image))
image_tensor = torch.stack(images)
data_entries = self.annotation_dataframe[(self.annotation_dataframe['scene'] == scene_id) & (self.annotation_dataframe['sample'] == sample_id)]
corners = data_entries[['fl_x', 'fr_x', 'bl_x', 'br_x', 'fl_y', 'fr_y','bl_y', 'br_y']].to_numpy()
categories = data_entries.category_id.to_numpy()
ego_path = os.path.join(sample_path, 'ego.png')
ego_image = Image.open(ego_path)
ego_image = torchvision.transforms.functional.to_tensor(ego_image)
road_image = convert_map_to_road_map(ego_image)
target = {}
target['bounding_box'] = torch.as_tensor(corners).view(-1, 2, 4)
target['category'] = torch.as_tensor(categories)
if self.extra_info:
actions = data_entries.action_id.to_numpy()
# You can change the binary_lane to False to get a lane with
lane_image = convert_map_to_lane_map(ego_image, binary_lane=True)
extra = {}
extra['action'] = torch.as_tensor(actions)
extra['ego_image'] = ego_image
extra['lane_image'] = lane_image
extra['file_path'] = sample_path
extra['scene_id'] = scene_id
extra['sample_id'] = sample_id
return image_tensor, target, road_image, extra
else:
return image_tensor, target, road_image