Acyclic
/
AutoETS2
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update

This commit is contained in:
TerenceLiu98 2024-05-09 13:26:29 +00:00
parent e83298d09f
commit 0b73c6585c
7 changed files with 1357 additions and 569 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

461
data.ipynb
View File

File diff suppressed because one or more lines are too long

1246
perception.ipynb
View File

File diff suppressed because one or more lines are too long

BIN
sample.png
View File

Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 11 KiB

After

Width:  |  Height:  |  Size: 20 KiB

3
src/perception/base.py
View File

@ -11,6 +11,7 @@ import os
import cv2
import math
import timm
import json
import random
import gitinfo
import argparse
@ -29,7 +30,7 @@ from albumentations.pytorch import ToTensorV2
OBJ_LABELS = {
"truck": 0, "bicycle": 1, "car": 2, "motorcycle": 3,
"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):

111
src/perception/datasets.py
View File

@ -1,98 +1,57 @@
from src.perception.base import *
class AutoDriveDataset(Dataset):
def __init__(self, csv_file, image_dir, lane_dir, da_dir, transform=True):
self.data_frame = pd.read_json(csv_file)
self.image_dir, self.lane_dir, self.da_dir = image_dir, lane_dir, da_dir
def __init__(self, image_dir, lane_dir, drivable_dir, transform=True):
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.drivable_dir = [str(i) for i in list(sorted(Path(f"{drivable_dir}").glob("*")))]
self.transform = transform
def __len__(self):
return len(self.data_frame)
return len(self.image_dir)
def __readdata__(self, idx):
image_name = os.path.join(self.image_dir, self.data_frame.iloc[idx, 0])
label_name = os.path.join(self.da_dir, self.data_frame.iloc[idx, 0]).replace("jpg", "png")
image = cv2.cvtColor(cv2.imread(f"{image_name}"), cv2.COLOR_BGR2RGB)
drivable = cv2.cvtColor(cv2.imread("{}".format(label_name)), cv2.COLOR_BGR2RGB)
lane = cv2.cvtColor(cv2.imread("{}".format(label_name.replace("drivable", "lane"))), cv2.COLOR_BGR2RGB)
label_data = self.data_frame.iloc[idx, 3]
boxes = []
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
name = self.image_dir[idx].split("/")[-1].split(".")[0]
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"{[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.drivable_dir if name in i][0]}"), cv2.COLOR_BGR2RGB)
drivable[np.all(drivable == [219, 94, 86], axis=-1)] = [255, 0, 0]
drivable[np.all(drivable == [86, 211, 219], axis=-1)] = [0, 0, 255]
return image, drivable
def __getitem__(self, idx):
image, lane, drivable, object_det = self.__readdata__(idx=idx)
image, drivable = self.__readdata__(idx=idx)
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
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
image_t = torch.from_numpy(image_t).float() / 255
drivable_t = torch.from_numpy(drivable_t).float() / 255
return image, image_t, drivable_t
def __augmentation__(self, image, lane, drivable, object_det):
def __augmentation__(self, image, drivable):
transform = Aug.Compose([
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"])
Aug.Normalize(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)])
transformed = transform(image=image, mask=drivable)
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
drivable = transformed["mask"].transpose(2, 0, 1)
return image, drivable
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
image, image_t, drivable_t = zip(*batch)
image = torch.stack(image, dim=0)
image_t = torch.stack(image_t, dim=0)
drivable_t = torch.stack(drivable_t, dim=0)
return image, image_t, drivable_t
if __name__ == "__main__":
dataset = AutoDriveDataset(csv_file="/home/bayes/data/ETS2/bdd100k/label/det_20/det_train.json",
image_dir="/home/bayes/data/ETS2/bdd100k/image/100k/train/",
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)
dataset = AutoDriveDataset(image_dir="/home/bayes/data/ETS2/bdd100k/images/100k/train/",
drivable_dir="/home/bayes/data/ETS2/bdd100k/labels/drivable/colormaps/train/", transform=True)
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")
image, image_t, drivable_t = next(iter(dataloader))
fig, ax = plt.subplots(1)
ax.imshow(image[0].permute(1,2,0).numpy())
ax.imshow(drivable_t[0].permute(1,2,0).numpy(), alpha=0.2)
ax.axis("off")
plt.savefig("sample.png", dpi=250)

105
src/perception/datasets_old.py Normal file
View File

@ -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("*")))]

BIN
test.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 66 KiB