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 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):

111
src/perception/datasets.py
View File

@ -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)

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