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