update
This commit is contained in:
parent
e83298d09f
commit
0b73c6585c
461
data.ipynb
461
data.ipynb
File diff suppressed because one or more lines are too long
1250
perception.ipynb
1250
perception.ipynb
File diff suppressed because one or more lines are too long
BIN
sample.png
BIN
sample.png
Binary file not shown.
Before Width: | Height: | Size: 11 KiB After Width: | Height: | Size: 20 KiB |
|
@ -11,6 +11,7 @@ import os
|
||||||
import cv2
|
import cv2
|
||||||
import math
|
import math
|
||||||
import timm
|
import timm
|
||||||
|
import json
|
||||||
import random
|
import random
|
||||||
import gitinfo
|
import gitinfo
|
||||||
import argparse
|
import argparse
|
||||||
|
@ -29,7 +30,7 @@ from albumentations.pytorch import ToTensorV2
|
||||||
OBJ_LABELS = {
|
OBJ_LABELS = {
|
||||||
"truck": 0, "bicycle": 1, "car": 2, "motorcycle": 3,
|
"truck": 0, "bicycle": 1, "car": 2, "motorcycle": 3,
|
||||||
"train": 4, "bus": 5, "traffic sign": 6, "rider": 7, "person": 8,
|
"train": 4, "bus": 5, "traffic sign": 6, "rider": 7, "person": 8,
|
||||||
"traffic light NA": 9, "traffic light R": 10, "traffic light G": 11, "traffic light Y": 12
|
"traffic light NA": 9, "traffic light red": 10, "traffic light green": 11, "traffic light yellow": 12
|
||||||
}
|
}
|
||||||
|
|
||||||
def use_device(GPU):
|
def use_device(GPU):
|
||||||
|
|
|
@ -1,98 +1,57 @@
|
||||||
from src.perception.base import *
|
from src.perception.base import *
|
||||||
|
|
||||||
class AutoDriveDataset(Dataset):
|
class AutoDriveDataset(Dataset):
|
||||||
def __init__(self, csv_file, image_dir, lane_dir, da_dir, transform=True):
|
def __init__(self, image_dir, lane_dir, drivable_dir, transform=True):
|
||||||
self.data_frame = pd.read_json(csv_file)
|
self.lane_dir = [str(i) for i in list(sorted(Path(f"{lane_dir}").glob("*")))]
|
||||||
self.image_dir, self.lane_dir, self.da_dir = image_dir, lane_dir, da_dir
|
self.image_dir = [str(i) for i in list(sorted(Path(f"{image_dir}").glob("*")))]
|
||||||
|
self.drivable_dir = [str(i) for i in list(sorted(Path(f"{drivable_dir}").glob("*")))]
|
||||||
self.transform = transform
|
self.transform = transform
|
||||||
|
|
||||||
def __len__(self):
|
def __len__(self):
|
||||||
return len(self.data_frame)
|
return len(self.image_dir)
|
||||||
|
|
||||||
def __readdata__(self, idx):
|
def __readdata__(self, idx):
|
||||||
image_name = os.path.join(self.image_dir, self.data_frame.iloc[idx, 0])
|
name = self.image_dir[idx].split("/")[-1].split(".")[0]
|
||||||
label_name = os.path.join(self.da_dir, self.data_frame.iloc[idx, 0]).replace("jpg", "png")
|
lane = cv2.cvtColor(cv2.imread(f"{[i for i in self.lane_dir if name in i][0]}"), cv2.COLOR_BGR2RGB)
|
||||||
image = cv2.cvtColor(cv2.imread(f"{image_name}"), cv2.COLOR_BGR2RGB)
|
image = cv2.cvtColor(cv2.imread(f"{[i for i in self.image_dir if name in i][0]}"), cv2.COLOR_BGR2RGB)
|
||||||
drivable = cv2.cvtColor(cv2.imread("{}".format(label_name)), cv2.COLOR_BGR2RGB)
|
drivable = cv2.cvtColor(cv2.imread(f"{[i for i in self.drivable_dir if name in i][0]}"), cv2.COLOR_BGR2RGB)
|
||||||
lane = cv2.cvtColor(cv2.imread("{}".format(label_name.replace("drivable", "lane"))), cv2.COLOR_BGR2RGB)
|
drivable[np.all(drivable == [219, 94, 86], axis=-1)] = [255, 0, 0]
|
||||||
label_data = self.data_frame.iloc[idx, 3]
|
drivable[np.all(drivable == [86, 211, 219], axis=-1)] = [0, 0, 255]
|
||||||
boxes = []
|
return image, drivable
|
||||||
labels = []
|
|
||||||
for item in label_data:
|
|
||||||
xmin, ymin, xmax, ymax = item["box2d"]["x1"], item["box2d"]["y1"], item["box2d"]["x2"], item["box2d"]["y2"]
|
|
||||||
boxes.append([xmin, ymin, xmax, ymax])
|
|
||||||
if item["category"] == "traffic light":
|
|
||||||
item["category"] = item["category"] + " " + item["attributes"]["trafficLightColor"]
|
|
||||||
elif item["category"] == "other vehicle" or "trailer":
|
|
||||||
item["category"] = "car"
|
|
||||||
elif item["category"] == "other person" or "pedestrain":
|
|
||||||
item["category"] = "person"
|
|
||||||
else:
|
|
||||||
pass
|
|
||||||
labels.append(OBJ_LABELS[item['category']])
|
|
||||||
|
|
||||||
boxes = torch.as_tensor(boxes, dtype=torch.float32)
|
|
||||||
labels = torch.as_tensor(labels, dtype=torch.int64)
|
|
||||||
object_det = {"boxes": boxes, "labels": labels}
|
|
||||||
return image, lane, drivable, object_det
|
|
||||||
|
|
||||||
def __getitem__(self, idx):
|
def __getitem__(self, idx):
|
||||||
image, lane, drivable, object_det = self.__readdata__(idx=idx)
|
image, drivable = self.__readdata__(idx=idx)
|
||||||
if self.transform:
|
if self.transform:
|
||||||
image, lane, drivable, object_det = self.__augmentation__(image, lane, drivable, object_det)
|
image_t, drivable_t = self.__augmentation__(image, drivable)
|
||||||
|
|
||||||
image = torch.from_numpy(image).float() / 255
|
image = torch.from_numpy(image).float() / 255
|
||||||
lane = torch.from_numpy(lane).float() / 255
|
image_t = torch.from_numpy(image_t).float() / 255
|
||||||
drivable = torch.from_numpy(drivable).float() / 255
|
drivable_t = torch.from_numpy(drivable_t).float() / 255
|
||||||
object_det["boxes"] = torch.from_numpy(np.array(object_det["boxes"]))
|
return image, image_t, drivable_t
|
||||||
return image, lane, drivable, object_det
|
|
||||||
|
|
||||||
def __augmentation__(self, image, lane, drivable, object_det):
|
def __augmentation__(self, image, drivable):
|
||||||
transform = Aug.Compose([
|
transform = Aug.Compose([
|
||||||
Aug.Resize (360, 640, p=1), Aug.HorizontalFlip(p=0.5), Aug.RandomBrightnessContrast(p=0.5)],
|
Aug.Normalize(mean=(0.485, 0.56, 0.406), std=(0.229, 0.224, 0.225)),
|
||||||
bbox_params=Aug.BboxParams(format='pascal_voc', label_fields=['labels']))
|
Aug.Resize (360, 640, p=1), Aug.HorizontalFlip(p=0.5), Aug.RandomBrightnessContrast(p=0.5)])
|
||||||
transformed = transform(image=image, masks=[lane, drivable],
|
transformed = transform(image=image, mask=drivable)
|
||||||
bboxes=object_det["boxes"], labels=object_det["labels"])
|
|
||||||
image = transformed["image"].transpose(2, 0, 1)
|
image = transformed["image"].transpose(2, 0, 1)
|
||||||
lane = transformed["masks"][0].transpose(2, 0, 1)
|
drivable = transformed["mask"].transpose(2, 0, 1)
|
||||||
drivable = transformed["masks"][1].transpose(2, 0, 1)
|
return image, drivable
|
||||||
object_det = {"boxes": transformed["bboxes"], "labels": transformed["labels"]}
|
|
||||||
return image, lane, drivable, object_det
|
|
||||||
|
|
||||||
def collate_fn(self, batch):
|
def collate_fn(self, batch):
|
||||||
images, lane, drivable, object_det = zip(*batch)
|
image, image_t, drivable_t = zip(*batch)
|
||||||
images = torch.stack(images, dim=0)
|
image = torch.stack(image, dim=0)
|
||||||
lane = torch.stack(lane, dim=0)
|
image_t = torch.stack(image_t, dim=0)
|
||||||
drivable = torch.stack(drivable, dim=0)
|
drivable_t = torch.stack(drivable_t, dim=0)
|
||||||
object_det = [{k: v for k, v in t.items()} for t in object_det]
|
return image, image_t, drivable_t
|
||||||
return images, lane, drivable, object_det
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
dataset = AutoDriveDataset(csv_file="/home/bayes/data/ETS2/bdd100k/label/det_20/det_train.json",
|
dataset = AutoDriveDataset(image_dir="/home/bayes/data/ETS2/bdd100k/images/100k/train/",
|
||||||
image_dir="/home/bayes/data/ETS2/bdd100k/image/100k/train/",
|
drivable_dir="/home/bayes/data/ETS2/bdd100k/labels/drivable/colormaps/train/", transform=True)
|
||||||
lane_dir="/home/bayes/data/ETS2/bdd100k/label/lane/colormaps/train/",
|
|
||||||
da_dir="/home/bayes/data/ETS2/bdd100k/label/drivable/colormaps/train/", transform=True)
|
|
||||||
A, B, C, D = dataset.__getitem__(idx=0)
|
|
||||||
dataloader = DataLoader(dataset, batch_size=1, shuffle=False, collate_fn=dataset.collate_fn)
|
dataloader = DataLoader(dataset, batch_size=1, shuffle=False, collate_fn=dataset.collate_fn)
|
||||||
images, lane, drivable, object_det = next(iter(dataloader))
|
image, image_t, drivable_t = next(iter(dataloader))
|
||||||
fig, ax = plt.subplots(2)
|
fig, ax = plt.subplots(1)
|
||||||
ax[0].imshow(A.permute(1,2,0).numpy())
|
ax.imshow(image[0].permute(1,2,0).numpy())
|
||||||
ax[0].imshow(B.permute(1,2,0).numpy(), alpha=0.5)
|
ax.imshow(drivable_t[0].permute(1,2,0).numpy(), alpha=0.2)
|
||||||
ax[0].imshow(C.permute(1,2,0).numpy(), alpha=0.2)
|
ax.axis("off")
|
||||||
for i in range(0, len(D["boxes"])):
|
|
||||||
rect = patches.Rectangle((D["boxes"][i][0], D["boxes"][i][1]),
|
|
||||||
D["boxes"][i][2] - D["boxes"][i][0], D["boxes"][i][3] - D["boxes"][i][1], linewidth=1, edgecolor='r', facecolor='none')
|
|
||||||
ax[0].add_patch(rect) # import matplotlib.patches as patches
|
|
||||||
|
|
||||||
ax[0].axis("off")
|
|
||||||
|
|
||||||
ax[1].imshow(images[0].permute(1,2,0).numpy())
|
|
||||||
ax[1].imshow(lane[0].permute(1,2,0).numpy(), alpha=0.5)
|
|
||||||
ax[1].imshow(drivable[0].permute(1,2,0).numpy(), alpha=0.2)
|
|
||||||
for i in range(0, len(object_det[0]["boxes"])):
|
|
||||||
rect = patches.Rectangle((object_det[0]["boxes"][i][0], object_det[0]["boxes"][i][1]),
|
|
||||||
object_det[0]["boxes"][i][2] - object_det[0]["boxes"][i][0], object_det[0]["boxes"][i][3] - object_det[0]["boxes"][i][1], linewidth=1, edgecolor='r', facecolor='none')
|
|
||||||
ax[1].add_patch(rect) # import matplotlib.patches as patches
|
|
||||||
|
|
||||||
ax[1].axis("off")
|
|
||||||
plt.savefig("sample.png", dpi=250)
|
plt.savefig("sample.png", dpi=250)
|
|
@ -0,0 +1,105 @@
|
||||||
|
from src.perception.base import *
|
||||||
|
|
||||||
|
class AutoDriveDataset(Dataset):
|
||||||
|
def __init__(self, image_dir, lane_dir, da_dir, object_dir, transform=True):
|
||||||
|
self.image_dir, self.lane_dir, self.da_dir = image_dir, lane_dir, da_dir
|
||||||
|
self.da_dir = [str(i) for i in list(sorted(Path(f"{da_dir}").glob("*")))]
|
||||||
|
self.lane_dir = [str(i) for i in list(sorted(Path(f"{lane_dir}").glob("*")))]
|
||||||
|
self.image_dir = [str(i) for i in list(sorted(Path(f"{image_dir}").glob("*")))]
|
||||||
|
self.object_dir = [str(i) for i in list(sorted(Path(f"{object_dir}").glob("*")))]
|
||||||
|
self.transform = transform
|
||||||
|
|
||||||
|
def __len__(self):
|
||||||
|
return len(self.image_dir)
|
||||||
|
|
||||||
|
def __readdata__(self, idx):
|
||||||
|
name = self.object_dir[idx].split("/")[-1].split(".")[0]
|
||||||
|
image = cv2.cvtColor(cv2.imread(f"{[i for i in self.image_dir if name in i][0]}"), cv2.COLOR_BGR2RGB)
|
||||||
|
drivable = cv2.cvtColor(cv2.imread(f"{[i for i in self.da_dir if name in i][0]}"), cv2.COLOR_BGR2RGB)
|
||||||
|
lane = cv2.cvtColor(cv2.imread(f"{[i for i in self.lane_dir if name in i][0]}"), cv2.COLOR_BGR2RGB)
|
||||||
|
label_data = json.load(open(f"{[i for i in self.object_dir if name in i][0]}"))["frames"]
|
||||||
|
boxes = []
|
||||||
|
labels = []
|
||||||
|
for item in label_data[0]["objects"]:
|
||||||
|
if item["category"] not in list(OBJ_LABELS.keys()):
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
xmin, ymin, xmax, ymax = item["box2d"]["x1"], item["box2d"]["y1"], item["box2d"]["x2"], item["box2d"]["y2"]
|
||||||
|
boxes.append([xmin, ymin, xmax, ymax])
|
||||||
|
if item["category"] == "traffic light":
|
||||||
|
item["category"] = item["category"] + " " + item["attributes"]["trafficLightColor"]
|
||||||
|
elif item["category"] == "other vehicle" or "trailer":
|
||||||
|
item["category"] = "car"
|
||||||
|
elif item["category"] == "other person" or "pedestrain":
|
||||||
|
item["category"] = "person"
|
||||||
|
else:
|
||||||
|
pass
|
||||||
|
labels.append(OBJ_LABELS[item['category']])
|
||||||
|
|
||||||
|
boxes = torch.as_tensor(boxes, dtype=torch.float32)
|
||||||
|
labels = torch.as_tensor(labels, dtype=torch.int64)
|
||||||
|
object_det = {"boxes": boxes, "labels": labels}
|
||||||
|
return image, lane, drivable, object_det
|
||||||
|
|
||||||
|
def __getitem__(self, idx):
|
||||||
|
image, lane, drivable, object_det = self.__readdata__(idx=idx)
|
||||||
|
if self.transform:
|
||||||
|
image, lane, drivable, object_det = self.__augmentation__(image, lane, drivable, object_det)
|
||||||
|
|
||||||
|
image = torch.from_numpy(image).float() / 255
|
||||||
|
lane = torch.from_numpy(lane).float() / 255
|
||||||
|
drivable = torch.from_numpy(drivable).float() / 255
|
||||||
|
object_det["boxes"] = torch.from_numpy(np.array(object_det["boxes"]))
|
||||||
|
return image, lane, drivable, object_det
|
||||||
|
|
||||||
|
def __augmentation__(self, image, lane, drivable, object_det):
|
||||||
|
transform = Aug.Compose([
|
||||||
|
Aug.Normlize(mean=(0.485, 0.56, 0.406), std=(0.229, 0.224, 0.225)),
|
||||||
|
Aug.Resize (360, 640, p=1), Aug.HorizontalFlip(p=0.5), Aug.RandomBrightnessContrast(p=0.5)],
|
||||||
|
bbox_params=Aug.BboxParams(format='pascal_voc', label_fields=['labels']))
|
||||||
|
transformed = transform(image=image, masks=[lane, drivable],
|
||||||
|
bboxes=object_det["boxes"], labels=object_det["labels"])
|
||||||
|
image = transformed["image"].transpose(2, 0, 1)
|
||||||
|
lane = transformed["masks"][0].transpose(2, 0, 1)
|
||||||
|
drivable = transformed["masks"][1].transpose(2, 0, 1)
|
||||||
|
object_det = {"boxes": transformed["bboxes"], "labels": transformed["labels"]}
|
||||||
|
return image, lane, drivable, object_det
|
||||||
|
|
||||||
|
def collate_fn(self, batch):
|
||||||
|
images, lane, drivable, object_det = zip(*batch)
|
||||||
|
images = torch.stack(images, dim=0)
|
||||||
|
lane = torch.stack(lane, dim=0)
|
||||||
|
drivable = torch.stack(drivable, dim=0)
|
||||||
|
object_det = [{k: v for k, v in t.items()} for t in object_det]
|
||||||
|
return images, lane, drivable, object_det
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
dataset = AutoDriveDataset(image_dir="/home/bayes/data/ETS2/bdd100k/images/100k/train/",
|
||||||
|
lane_dir="/home/bayes/data/ETS2/bdd100k/labels/lane_gt/train/",
|
||||||
|
da_dir="/home/bayes/data/ETS2/bdd100k/labels/drivable_gt/train/",
|
||||||
|
object_dir="/home/bayes/data/ETS2/bdd100k/labels/det_gt/train/", transform=True)
|
||||||
|
A, B, C, D = dataset.__getitem__(idx=10)
|
||||||
|
dataloader = DataLoader(dataset, batch_size=1, shuffle=False, collate_fn=dataset.collate_fn)
|
||||||
|
images, lane, drivable, object_det = next(iter(dataloader))
|
||||||
|
fig, ax = plt.subplots(2)
|
||||||
|
ax[0].imshow(A.permute(1,2,0).numpy())
|
||||||
|
ax[0].imshow(B.permute(1,2,0).numpy(), alpha=0.5)
|
||||||
|
ax[0].imshow(C.permute(1,2,0).numpy(), alpha=0.2)
|
||||||
|
for i in range(0, len(D["boxes"])):
|
||||||
|
rect = patches.Rectangle((D["boxes"][i][0], D["boxes"][i][1]),
|
||||||
|
D["boxes"][i][2] - D["boxes"][i][0], D["boxes"][i][3] - D["boxes"][i][1], linewidth=1, edgecolor='r', facecolor='none')
|
||||||
|
ax[0].add_patch(rect) # import matplotlib.patches as patches
|
||||||
|
|
||||||
|
ax[0].axis("off")
|
||||||
|
|
||||||
|
ax[1].imshow(images[0].permute(1,2,0).numpy())
|
||||||
|
ax[1].imshow(lane[0].permute(1,2,0).numpy(), alpha=0.5)
|
||||||
|
ax[1].imshow(drivable[0].permute(1,2,0).numpy(), alpha=0.2)
|
||||||
|
for i in range(0, len(object_det[0]["boxes"])):
|
||||||
|
rect = patches.Rectangle((object_det[0]["boxes"][i][0], object_det[0]["boxes"][i][1]),
|
||||||
|
object_det[0]["boxes"][i][2] - object_det[0]["boxes"][i][0], object_det[0]["boxes"][i][3] - object_det[0]["boxes"][i][1], linewidth=1, edgecolor='r', facecolor='none')
|
||||||
|
ax[1].add_patch(rect) # import matplotlib.patches as patches
|
||||||
|
|
||||||
|
ax[1].axis("off")
|
||||||
|
plt.savefig("sample.png", dpi=250)
|
||||||
|
[str(i) for i in list(sorted(Path("/home/bayes/data/ETS2/bdd100k/images/100k/train/").glob("*")))]
|
Loading…
Reference in New Issue