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}')
# print("scene: {} | sample: {}".format(scene_id, 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.img_transform(image))
image_tensor = torch.stack(images)
h, w = image_tensor.size(2), image_tensor.size(3)
image_tensor = image_tensor.view(18, h, w)
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)
boxes = []
num_objs = target['bounding_box'].shape[0]
for i in range(num_objs):
bb = target['bounding_box'][i]
xmin, xmax = min(bb[0]), max(bb[0])
ymin, ymax = min(bb[1]), max(bb[1])
boxes.append([xmin, xmax, ymin, ymax])
boxes = torch.as_tensor(boxes, dtype = torch.float32)
labels = torch.ones((num_objs,), dtype = torch.int64)
area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
target['category'] = torch.as_tensor(categories)
mask = Image.open(os.path.join(sample_path, "obj_map.png")).convert('L')
# print("Mask image shape: {}".format(mask.size))
to_tensor = torchvision.transforms.ToTensor()
# print("Mask tensor shape: {}".format(to_tensor(mask).shape))
mask = self.map_transform(mask)
target['mask'] = mask
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
return image_tensor, target, road_image, extra
else:
return image_tensor, target, road_image
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}')
sample_depth_path = os.path.join(self.depth_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[0](image))
image_tensor = torch.stack(images)
depths = []
for depth_image in image_names:
depth_path = os.path.join(sample_depth_path, depth_image)
image = Image.open(depth_path)
depths.append(self.transform[1](image))
depth_tensor = torch.stack(depths)
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
return image_tensor, target, road_image, extra, depth_tensor
else:
return image_tensor, target, road_image, depth_tensor
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.img_transform(image))
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)
random_rotate = random.randint(0, 1)
if random_rotate:
images.reverse()
ego_image = TF.rotate(ego_image, 180)
image_tensor = torch.cat(images)
# ego_image = torchvision.transforms.functional.to_tensor(ego_image)
# road_image = convert_map_to_road_map(ego_image)
ego_image = self.map_transform(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
# print(sample_path)
# print(type(image_tensor), type(target), type(ego_image), type(extra))
return image_tensor, target, ego_image, extra
else:
# print(type(image_tensor), type(target), type(ego_image))
return image_tensor, target, ego_image
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}')
front = [
Image.open(os.path.join(sample_path, image_names[0])),
Image.open(os.path.join(sample_path, image_names[1])),
Image.open(os.path.join(sample_path, image_names[2]))
]
back = [
Image.open(os.path.join(sample_path, image_names[3])),
Image.open(os.path.join(sample_path, image_names[4])),
Image.open(os.path.join(sample_path, image_names[5]))
]
combo_1 = append_images(front, direction='horizontal')
combo_2 = append_images(back, direction='horizontal')
combined = append_images([combo_1, combo_2], direction='vertical')
#image_tensor = self.transform(combined)
image_tensor = torch.stack([self.transform(combined)])
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
return image_tensor, target, road_image, extra
else:
return image_tensor, target, road_image
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)
image = self.transform(image)
image = image.unsqueeze(0)
images.append(image)
image_tensor = torch.cat(images, dim=0)
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().reshape(-1, 2, 4)
labels = self.build_labels(corners)
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)
labels[:, 1] = categories
# column order: category, x, y, width, height; Plan to transpose the label
# reserve first column for idx of each instance.
labels = torch.as_tensor(labels)
# build lane map label
road_image_temp = torch.zeros(road_image.shape)
road_image_temp[road_image] = 1
road_image = road_image_temp
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, labels, road_image, extra
else:
return image_tensor, labels, road_image
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)
image_tensor = image_tensor.reshape(18, 256, 306)
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 = {}
bboxes = torch.as_tensor(corners).view(-1, 2, 4)
#print("***")
#print(bboxes.shape)
target['boxes'] = self.reshapeTarget(bboxes)
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
return image_tensor, target, road_image, extra
else:
# print(type(image_tensor))
# print(type(target))
return image_tensor, target
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}')
#get camera and only get that image
image_path = os.path.join(sample_path, self.camera)
image = Image.open(image_path)
image = self.transform(image)
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)
road_image_rebinned = rebin(road_image, self.downsample_shape)
road_image_rebinned_mod = torch.Tensor(road_image_rebinned.numpy() *
self.mask)
road_image_rebinned_mod = road_image_rebinned[self.mask]
##Preprocess road image
#road_image = torch.Tensor(road_image.numpy()*self.mask)
## Get resample down
#road_image_mod = road_image[self.mask]
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
return image, target, road_image_rebinned, extra, road_image_rebinned_mod
else:
return image, target, road_image_rebinned
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
return image_tensor, target, road_image, extra
else:
return image_tensor, target, road_image
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)
# print("TBB shape: {}".format(target['bounding_box'].shape))
bbox_center = ((target['bounding_box'][:, :, 0] + target['bounding_box'][:, :, 3])/2) #.view(-1, 1)
bbox_width = (target['bounding_box'][:, :, 1][:, 0] - target['bounding_box'][:, :, 0][:, 0]).view(-1, 1)
bbox_length = (target['bounding_box'][:, :, 0][:, 1] - target['bounding_box'][:, :, 2][:, 1]).view(-1, 1)
target['bbox_alt'] = torch.cat([bbox_center, torch.cat([bbox_length, bbox_width], dim = 1)], dim = 1)
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
return image_tensor, target, road_image, extra #, scene_id, sample_id
else:
return image_tensor, target, road_image
def __getitem__(self, index):
#print(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}')
data_entries = self.annotation_dataframe[
(self.annotation_dataframe['scene'] == scene_id)
& (self.annotation_dataframe['sample'] == sample_id)]
#print(data_entries)
#corners = data_entries[['fl_x', 'fr_x', 'bl_x', 'br_x', 'fl_y', 'fr_y','bl_y', 'br_y']].to_numpy()
data_entries['min_x'] = data_entries[['fl_x', 'fr_x', 'bl_x',
'br_x']].min(axis=1)
data_entries['min_y'] = data_entries[['fl_y', 'fr_y', 'bl_y',
'br_y']].min(axis=1)
data_entries['max_x'] = data_entries[['fl_x', 'fr_x', 'bl_x',
'br_x']].max(axis=1)
data_entries['max_y'] = data_entries[['fl_y', 'fr_y', 'bl_y',
'br_y']].max(axis=1)
#print(data_entries)
corners = data_entries[['min_x', 'min_y', 'max_x', 'max_y']].to_numpy()
#print(corners)
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['bounding_box'] = torch.as_tensor(corners).view(-1, 2, 2) #<--------------Look into this later
target['bounding_box'] = torch.as_tensor(corners).view(-1, 4)
target['category'] = torch.as_tensor(categories)
images = []
for image_name in image_names:
image_path = os.path.join(sample_path, image_name)
image = Image.open(image_path)
new_image, new_box, new_label = transform(image,
target,
split=self.split)
images.append(new_image)
#<-------------------------Uncomment after we use all data!
image_tensor = torch.stack(images)
target['bounding_box'] = new_box
target['category'] = new_label
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
return image_tensor, target, road_image, extra
else:
image_tensor = image_tensor.squeeze(0)
boxes = target['bounding_box'].type(
torch.FloatTensor) # (n_objects, 4)
labels = target['category'].type(torch.LongTensor)
#return image_tensor,target['bounding_box'],target['category']
return image_tensor, boxes, labels
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}')
total_width = RESIZE_DIM * 3
total_height = RESIZE_DIM * 2
new_img = Image.new('RGB', (total_width, total_height))
x_offset = 0
for image_name in image_names[:3]:
image_path = os.path.join(sample_path, image_name)
image = Image.open(image_path)
image = image.resize((RESIZE_DIM, RESIZE_DIM))
new_img.paste(image, (x_offset, 0))
x_offset += RESIZE_DIM
x_offset = 0
y_offset = RESIZE_DIM
for image_name in image_names[3:]:
image_path = os.path.join(sample_path, image_name)
image = Image.open(image_path)
image = image.resize((RESIZE_DIM, RESIZE_DIM))
# keep spatial information
if 'CAM_BACK.jpeg' == image_name:
image = image.rotate(-180)
elif 'CAM_BACK_LEFT.jpeg' == image_name:
image = image.rotate(180)
else:
image = image.rotate(-180)
new_img.paste(image, (x_offset, y_offset))
x_offset += RESIZE_DIM
#new_img.save(f'../foo/{scene_id}_{sample_id}.jpg')
image_tensor = torch.as_tensor(self.transform(new_img))
# print(image_tensor.shape)
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().reshape(-1, 2, 4)
labels = np.zeros((corners.shape[0], 6))
for idx, corner in enumerate(corners):
point_squence = np.stack([
corner[:, 0], corner[:, 1], corner[:, 3], corner[:, 2],
corner[:, 0]
])
x = point_squence.T[0] * 10 + 400
y = -point_squence.T[1] * 10 + 400
xc = (x[0] + x[2]) / 2
yc = (y[0] + y[1]) / 2
w = np.abs(x[0] - x[2]) / 2
h = np.abs(y[0] - y[1]) / 2
# normalize to 0-1
labels[idx, 2] = xc / 800
labels[idx, 3] = yc / 800
labels[idx, 4] = w / 800
labels[idx, 5] = h / 800
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)
labels[:, 1] = categories
# column order: category, x, y, width, height; Plan to transpose the label
# reserve first column for idx of each instance.
labels = torch.as_tensor(labels)
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
return image_tensor, labels, road_image, extra
else:
return image_tensor, labels, road_image
def __getitem__(self, index):
#get 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]
sample_path = os.path.join(self.image_folder, f'scene_{scene_id}',
f'sample_{sample_id}')
image_path = os.path.join(self.image_folder, f'scene_{scene_id}',
f'sample_{sample_id}', image_name)
image = Image.open(image_path)
#get Labels
ego_path = os.path.join(sample_path,
f'ego_{index % NUM_IMAGE_PER_SAMPLE}.png')
ego_image = Image.open(ego_path)
ego_image = torchvision.transforms.functional.to_tensor(ego_image)
ego_image = ego_image.squeeze().long()
target_path = os.path.join(
sample_path, f'target_{index % NUM_IMAGE_PER_SAMPLE}.png')
target_image = Image.open(target_path)
target_image = torchvision.transforms.functional.to_tensor(
target_image)
target_image = target_image.squeeze().long()
if self.discr_run:
scene_id_discr = np.random.choice(self.scene_index)
sample_id_discr = np.random.choice(126)
sample_path_discr = os.path.join(self.image_folder,
f'scene_{scene_id_discr}',
f'sample_{sample_id_discr}')
ego_path_discr = os.path.join(
sample_path_discr, f'ego_{index % NUM_IMAGE_PER_SAMPLE}.png')
ego_image_discr = Image.open(ego_path_discr)
ego_image_discr = torchvision.transforms.functional.to_tensor(
ego_image_discr)
ego_image_discr = process_discr(ego_image_discr.squeeze())
target_path_discr = os.path.join(
sample_path_discr,
f'target_{index % NUM_IMAGE_PER_SAMPLE}.png')
target_image_discr = Image.open(target_path_discr)
target_image_discr = torchvision.transforms.functional.to_tensor(
target_image_discr)
target_image_discr = process_discr(target_image_discr.squeeze())
return self.transform(
image
), target_image, ego_image, target_image_discr, ego_image_discr
elif self.extra_info:
data_entries = self.annotation_dataframe[
(self.annotation_dataframe['scene'] == scene_id)
& (self.annotation_dataframe['sample'] == sample_id)]
categories = data_entries.category_id.to_numpy()
target['category'] = torch.as_tensor(categories)
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
return self.transform(image), target_image, ego_image, extra
else:
return self.transform(image), target_image, ego_image
def __getitem__(self, index):
#get 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]
scene_id_discr = np.random.choice(self.scene_index)
sample_id_discr = np.random.choice(126)
sample_path = os.path.join(self.image_folder, f'scene_{scene_id}',
f'sample_{sample_id}')
image_path = os.path.join(self.image_folder, f'scene_{scene_id}',
f'sample_{sample_id}', image_name)
sample_path_discr = os.path.join(self.image_folder,
f'scene_{scene_id_discr}',
f'sample_{sample_id_discr}')
image = Image.open(image_path)
#get Labels
data_entries = self.annotation_dataframe[
(self.annotation_dataframe['scene'] == scene_id)
& (self.annotation_dataframe['sample'] == sample_id)]
data_entries_discr = self.annotation_dataframe[
(self.annotation_dataframe['scene'] == scene_id_discr)
& (self.annotation_dataframe['sample'] == sample_id_discr)]
corners = data_entries[[
'fl_x', 'fr_x', 'bl_x', 'br_x', 'fl_y', 'fr_y', 'bl_y', 'br_y'
]].to_numpy()
corners_discr = data_entries_discr[[
'fl_x', 'fr_x', 'bl_x', 'br_x', 'fl_y', 'fr_y', 'bl_y', 'br_y'
]].to_numpy()
ego_path = os.path.join(sample_path, 'ego.png')
ego_path_discr = os.path.join(sample_path_discr, 'ego.png')
ego_image = Image.open(ego_path)
ego_image_discr = Image.open(ego_path_discr)
#ego_image = ego_image.resize((64,64)) #Trevor - Resize image so we can compute loss static input
ego_image = torchvision.transforms.functional.to_tensor(ego_image)
ego_image_discr = torchvision.transforms.functional.to_tensor(
ego_image_discr)
road_image = split_with_index(convert_map_to_road_map(ego_image),
index % NUM_IMAGE_PER_SAMPLE)
road_image_discr = split_with_index(
convert_map_to_road_map(ego_image_discr),
index % NUM_IMAGE_PER_SAMPLE)
#target = {}
bounding_box = bounding_box_to_image(
torch.as_tensor(corners).view(-1, 2, 4))
bounding_box_discr = bounding_box_to_image(
torch.as_tensor(corners_discr).view(-1, 2, 4))
target = split_with_index(bounding_box, index % NUM_IMAGE_PER_SAMPLE)
target_discr = split_with_index(bounding_box_discr,
index % NUM_IMAGE_PER_SAMPLE)
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 = {}
categories = data_entries.category_id.to_numpy()
extra['category'] = torch.as_tensor(categories)
extra['action'] = torch.as_tensor(actions)
extra['ego_image'] = ego_image
extra['lane_image'] = lane_image
return self.transform(
image
), target, road_image, target_discr, road_image_discr, extra
else:
return self.transform(
image), target, road_image, target_discr, road_image_discr
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()
#corners = (corners + 40)/0.1
# CUSTOM CODE
x_corners = (data_entries[['fl_x', 'fr_x', 'bl_x', 'br_x']].to_numpy()
+ 40) / 0.1
y_corners = (data_entries[['fl_y', 'fr_y', 'bl_y', 'br_y']].to_numpy()
+ 40) / 0.1
x_min_corn = np.amin(x_corners, axis=1)
y_min_corn = np.amin(y_corners, axis=1)
x_max_corn = np.amax(x_corners, axis=1)
y_max_corn = np.amax(y_corners, axis=1)
bboxes = np.stack((x_min_corn, y_min_corn, x_max_corn, y_max_corn),
axis=1)
image_id = torch.tensor([index, scene_id, sample_id])
# 0 needs to be the background class always.
categories = data_entries.category_id.to_numpy() + 1
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['boxes'] = torch.as_tensor(bboxes).float()
target['labels'] = torch.as_tensor(categories)
target['image_id'] = image_id
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
return image_tensor, target
else:
return image_tensor, target, road_image
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, 'scene_' + str(scene_id),
'sample_' + str(sample_id))
images = []
for image_name in image_names:
image_path = os.path.join(sample_path, image_name)
image = Image.open(image_path)
image.load()
images.append(self.transform["image"](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.load()
ego_image = torchvision.transforms.functional.to_tensor(ego_image)
road_image = convert_map_to_road_map(ego_image)
road_image = self.transform["road"](road_image.type(torch.FloatTensor))
# print(torch.as_tensor(corners).view(-1, 2, 4).transpose(1,2).flatten(1,2))
bounding_box = torch.as_tensor(corners).view(
-1, 2, 4) #.transpose(1,2)#.flatten(1,2)
bounding_box[:, 0] = (bounding_box[:, 0] * 10) + 400
bounding_box[:, 1] = (-bounding_box[:, 1] * 10) + 400
bounding_box = (bounding_box * 256) / 800
bounding_box = bounding_box.transpose(1, 2)
# print(bounding_box[:, :, 0].shape)
bbox = torch.zeros(bounding_box.shape[0], 4)
# print(bbox.shape, bounding_box.shape)
# bbox = (bbox * 256)/800
bbox_new = torch.zeros(bounding_box.shape[0], 5)
# print(bbox.shape, bounding_box.shape)
# bbox[:, 0] = bounding_box[:, :, 0].min(dim=1)[0]
# bbox[:, 1] = bounding_box[:, :, 1].min(dim=1)[0]
# bbox[:, 2] = bounding_box[:, :, 0].max(dim=1)[0]
# bbox[:, 3] = bounding_box[:, :, 1].max(dim=1)[0]
# Computre rotate angle from center point
for i, box in enumerate(bounding_box):
if box[0][0] <= box[2][0] and box[0][1] >= box[1][1]:
br = box[0]
bl = box[1]
fr = box[2]
fl = box[3]
else:
fl = box[0]
fr = box[1]
bl = box[2]
br = box[3]
# print("before:",box)
centerpoint = (fl + br) / 2
if fl[0] > fr[0]: # negative angle
if fr[0] != centerpoint[0]:
theta = torch.atan(
(fr[1] - centerpoint[1]) / abs(fr[0] - centerpoint[0]))
else:
theta = (np.pi / 2)
a = bl - centerpoint
b = fl - centerpoint
tempangle = torch.acos(
torch.dot(a, b) / (torch.norm(a, 2) * torch.norm(b, 2)))
beta = (np.pi - tempangle) / 2
if fr[0] > centerpoint[0]:
gamma = -(theta - beta)
else:
gamma = -(np.pi - theta - beta)
# print ("-----test----")
# print (torch.norm(a, 2))
# print (torch.norm(b, 2))
# print (theta)
# print (beta)
# print (gamma)
elif fl[0] < fr[0]: # positive angle
if centerpoint[0] != br[0]:
theta = torch.atan(
(br[1] - centerpoint[1]) / abs(centerpoint[0] - br[0]))
else:
theta = np.pi / 2
a = fl - centerpoint
b = bl - centerpoint
tempangle = torch.acos(
torch.dot(a, b) / (torch.norm(a, 2) * torch.norm(b, 2)))
beta = (np.pi - tempangle) / 2
if br[0] > centerpoint[0]:
gamma = (theta - beta)
else:
gamma = (np.pi - theta - beta)
else:
gamma = 0
# print((gamma*180)/np.pi)
#theta = np.arctan((fr[1] - br[1])/(fr[0]-br[0]))
bbox_new[i, 4] = gamma
translation_matrix = torch.tensor([[1, 0, centerpoint[0]],
[0, 1, centerpoint[1]],
[0, 0, 1]])
reverse_translation_matrix = torch.tensor([[1, 0, -centerpoint[0]],
[0, 1, -centerpoint[1]],
[0, 0, 1]])
rotation_matrix = torch.tensor(
[[torch.cos(-gamma), -torch.sin(-gamma), 0],
[torch.sin(-gamma), torch.cos(-gamma), 0], [0, 0, 1]])
# print(translation_matrix,reverse_translation_matrix,rotation_matrix)
# print(box.shape)
box = torch.cat([
box.transpose(0, 1),
torch.ones(box.shape[0]).type(torch.DoubleTensor).unsqueeze(0)
],
dim=0)
# print(box)
bbox_rotated = torch.matmul(
translation_matrix,
torch.matmul(rotation_matrix,
torch.matmul(reverse_translation_matrix,
box)))[:2]
# print(bbox_rotated)
# print("\nrotation matrix shape:",rotation_matrix.shape)
# rotation_matrix = torch.from_numpy(rotation_matrix)
# bbox_rotated = torch.matmul(rotation_matrix, torch.transpose(box, 0, 1))
# print("\nbbox_rotated shape:",bbox_rotated.shape)
# print("\nrotated_bbox:", bbox_rotated)
# print("\nbbox new shape:",bbox_new.shape)
if box[0][0] <= box[2][0] and box[0][1] >= box[1][1]:
bbox_new[i, 0] = bbox_rotated[0, 1]
bbox_new[i, 1] = bbox_rotated[1, 1]
bbox_new[i, 2] = bbox_rotated[0, 2]
bbox_new[i, 3] = bbox_rotated[1, 2]
else:
bbox_new[i, 0] = bbox_rotated[0, 0]
bbox_new[i, 1] = bbox_rotated[1, 0]
bbox_new[i, 2] = bbox_rotated[0, 3]
bbox_new[i, 3] = bbox_rotated[1, 3]
# print("\nafter:",bbox_new[i])
# if len(bbox_rotated[bbox_rotated<0])>0:
# print(bbox[0])
# print(scene_id, sample_id, bounding_box.shape)
classes = torch.as_tensor(categories).view(-1, 1)
# print(bbox_new.shape,classes.shape)
if self.args.gen_semantic_map:
semantic_map_path = os.path.join(sample_path, "semantic_map.npy")
semantic_map = np.load(semantic_map_path)
semantic_map = F.one_hot(
torch.tensor(semantic_map).to(torch.int64), 11)
else: # self.args.gen_object_map:
semantic_map_path = os.path.join(sample_path, "object_map.npy")
semantic_map = np.load(semantic_map_path)
semantic_map = F.one_hot(
torch.tensor(semantic_map).to(torch.int64), 3)
semantic_map = semantic_map.transpose(1, 2).transpose(0, 1)
# plt.imshow(semantic_map)
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)
action = torch.as_tensor(actions)
ego = self.transform["road"](ego_image)
road = lane_image
# print(scene_id, sample_id, bounding_box[0])
# print(bounding_box.shape,classes.shape)
# print(classes)
# exit(0)
return index, image_tensor, bbox_new, classes, action, ego, road_image, semantic_map
else:
return index, image_tensor, bbox_new, classes
