feat: initial HSAP platform
Huaxu Sentinel Active Safety Platform with embedded algorithm code, Docker Compose setup, and vendored dataset scaffolds for clone-and-run. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -0,0 +1,7 @@
|
||||
from .registry import build_dataset, build_dataloader
|
||||
|
||||
from .tusimple import TuSimple
|
||||
from .culane import CULane
|
||||
from .llamas import LLAMAS
|
||||
from .mufld import MufldLane
|
||||
from .process import *
|
||||
@@ -0,0 +1,67 @@
|
||||
import os.path as osp
|
||||
import os
|
||||
import numpy as np
|
||||
import cv2
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
import torchvision
|
||||
import logging
|
||||
from .registry import DATASETS
|
||||
from .process import Process
|
||||
from clrnet.utils.visualization import imshow_lanes
|
||||
from mmcv.parallel import DataContainer as DC
|
||||
|
||||
|
||||
@DATASETS.register_module
|
||||
class BaseDataset(Dataset):
|
||||
def __init__(self, data_root, split, processes=None, cfg=None):
|
||||
self.cfg = cfg
|
||||
self.logger = logging.getLogger(__name__)
|
||||
self.data_root = data_root
|
||||
self.training = 'train' in split
|
||||
self.processes = Process(processes, cfg)
|
||||
|
||||
def view(self, predictions, img_metas):
|
||||
img_metas = [item for img_meta in img_metas.data for item in img_meta]
|
||||
for lanes, img_meta in zip(predictions, img_metas):
|
||||
img_name = img_meta['img_name']
|
||||
img = cv2.imread(osp.join(self.data_root, img_name))
|
||||
out_file = osp.join(self.cfg.work_dir, 'visualization',
|
||||
img_name.replace('/', '_'))
|
||||
lanes = [lane.to_array(self.cfg) for lane in lanes]
|
||||
imshow_lanes(img, lanes, out_file=out_file)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_infos)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
data_info = self.data_infos[idx]
|
||||
img = cv2.imread(data_info['img_path'])
|
||||
img = img[self.cfg.cut_height:, :, :]
|
||||
sample = data_info.copy()
|
||||
sample.update({'img': img})
|
||||
|
||||
if self.training:
|
||||
label = cv2.imread(sample['mask_path'], cv2.IMREAD_UNCHANGED)
|
||||
if len(label.shape) > 2:
|
||||
label = label[:, :, 0]
|
||||
label = label.squeeze()
|
||||
label = label[self.cfg.cut_height:, :]
|
||||
sample.update({'mask': label})
|
||||
|
||||
if self.cfg.cut_height != 0:
|
||||
new_lanes = []
|
||||
for i in sample['lanes']:
|
||||
lanes = []
|
||||
for p in i:
|
||||
lanes.append((p[0], p[1] - self.cfg.cut_height))
|
||||
new_lanes.append(lanes)
|
||||
sample.update({'lanes': new_lanes})
|
||||
|
||||
sample = self.processes(sample)
|
||||
meta = {'full_img_path': data_info['img_path'],
|
||||
'img_name': data_info['img_name']}
|
||||
meta = DC(meta, cpu_only=True)
|
||||
sample.update({'meta': meta})
|
||||
|
||||
return sample
|
||||
143
algorithms/lane_ufld/code/CLRNet-main/clrnet/datasets/culane.py
Normal file
143
algorithms/lane_ufld/code/CLRNet-main/clrnet/datasets/culane.py
Normal file
@@ -0,0 +1,143 @@
|
||||
import os
|
||||
import os.path as osp
|
||||
import numpy as np
|
||||
from .base_dataset import BaseDataset
|
||||
from .registry import DATASETS
|
||||
import clrnet.utils.culane_metric as culane_metric
|
||||
import cv2
|
||||
from tqdm import tqdm
|
||||
import logging
|
||||
import pickle as pkl
|
||||
|
||||
LIST_FILE = {
|
||||
'train': 'list/train_gt.txt',
|
||||
'val': 'list/val.txt',
|
||||
'test': 'list/test.txt',
|
||||
}
|
||||
|
||||
CATEGORYS = {
|
||||
'normal': 'list/test_split/test0_normal.txt',
|
||||
'crowd': 'list/test_split/test1_crowd.txt',
|
||||
'hlight': 'list/test_split/test2_hlight.txt',
|
||||
'shadow': 'list/test_split/test3_shadow.txt',
|
||||
'noline': 'list/test_split/test4_noline.txt',
|
||||
'arrow': 'list/test_split/test5_arrow.txt',
|
||||
'curve': 'list/test_split/test6_curve.txt',
|
||||
'cross': 'list/test_split/test7_cross.txt',
|
||||
'night': 'list/test_split/test8_night.txt',
|
||||
}
|
||||
|
||||
|
||||
@DATASETS.register_module
|
||||
class CULane(BaseDataset):
|
||||
def __init__(self, data_root, split, processes=None, cfg=None):
|
||||
super().__init__(data_root, split, processes=processes, cfg=cfg)
|
||||
self.list_path = osp.join(data_root, LIST_FILE[split])
|
||||
self.split = split
|
||||
self.load_annotations()
|
||||
|
||||
def load_annotations(self):
|
||||
self.logger.info('Loading CULane annotations...')
|
||||
# Waiting for the dataset to load is tedious, let's cache it
|
||||
os.makedirs('cache', exist_ok=True)
|
||||
cache_path = 'cache/culane_{}.pkl'.format(self.split)
|
||||
if os.path.exists(cache_path):
|
||||
with open(cache_path, 'rb') as cache_file:
|
||||
self.data_infos = pkl.load(cache_file)
|
||||
self.max_lanes = max(
|
||||
len(anno['lanes']) for anno in self.data_infos)
|
||||
return
|
||||
|
||||
self.data_infos = []
|
||||
with open(self.list_path) as list_file:
|
||||
for line in list_file:
|
||||
infos = self.load_annotation(line.split())
|
||||
self.data_infos.append(infos)
|
||||
|
||||
# cache data infos to file
|
||||
with open(cache_path, 'wb') as cache_file:
|
||||
pkl.dump(self.data_infos, cache_file)
|
||||
|
||||
def load_annotation(self, line):
|
||||
infos = {}
|
||||
img_line = line[0]
|
||||
img_line = img_line[1 if img_line[0] == '/' else 0::]
|
||||
img_path = os.path.join(self.data_root, img_line)
|
||||
infos['img_name'] = img_line
|
||||
infos['img_path'] = img_path
|
||||
if len(line) > 1:
|
||||
mask_line = line[1]
|
||||
mask_line = mask_line[1 if mask_line[0] == '/' else 0::]
|
||||
mask_path = os.path.join(self.data_root, mask_line)
|
||||
infos['mask_path'] = mask_path
|
||||
|
||||
if len(line) > 2:
|
||||
exist_list = [int(l) for l in line[2:]]
|
||||
infos['lane_exist'] = np.array(exist_list)
|
||||
|
||||
anno_path = img_path[:-3] + 'lines.txt' # remove sufix jpg and add lines.txt
|
||||
with open(anno_path, 'r') as anno_file:
|
||||
data = [
|
||||
list(map(float, line.split()))
|
||||
for line in anno_file.readlines()
|
||||
]
|
||||
lanes = [[(lane[i], lane[i + 1]) for i in range(0, len(lane), 2)
|
||||
if lane[i] >= 0 and lane[i + 1] >= 0] for lane in data]
|
||||
lanes = [list(set(lane)) for lane in lanes] # remove duplicated points
|
||||
lanes = [lane for lane in lanes
|
||||
if len(lane) > 2] # remove lanes with less than 2 points
|
||||
|
||||
lanes = [sorted(lane, key=lambda x: x[1])
|
||||
for lane in lanes] # sort by y
|
||||
infos['lanes'] = lanes
|
||||
|
||||
return infos
|
||||
|
||||
def get_prediction_string(self, pred):
|
||||
ys = np.arange(270, 590, 8) / self.cfg.ori_img_h
|
||||
out = []
|
||||
for lane in pred:
|
||||
xs = lane(ys)
|
||||
valid_mask = (xs >= 0) & (xs < 1)
|
||||
xs = xs * self.cfg.ori_img_w
|
||||
lane_xs = xs[valid_mask]
|
||||
lane_ys = ys[valid_mask] * self.cfg.ori_img_h
|
||||
lane_xs, lane_ys = lane_xs[::-1], lane_ys[::-1]
|
||||
lane_str = ' '.join([
|
||||
'{:.5f} {:.5f}'.format(x, y) for x, y in zip(lane_xs, lane_ys)
|
||||
])
|
||||
if lane_str != '':
|
||||
out.append(lane_str)
|
||||
|
||||
return '\n'.join(out)
|
||||
|
||||
def evaluate(self, predictions, output_basedir):
|
||||
loss_lines = [[], [], [], []]
|
||||
print('Generating prediction output...')
|
||||
for idx, pred in enumerate(predictions):
|
||||
output_dir = os.path.join(
|
||||
output_basedir,
|
||||
os.path.dirname(self.data_infos[idx]['img_name']))
|
||||
output_filename = os.path.basename(
|
||||
self.data_infos[idx]['img_name'])[:-3] + 'lines.txt'
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
output = self.get_prediction_string(pred)
|
||||
|
||||
with open(os.path.join(output_dir, output_filename),
|
||||
'w') as out_file:
|
||||
out_file.write(output)
|
||||
|
||||
for cate, cate_file in CATEGORYS.items():
|
||||
result = culane_metric.eval_predictions(output_basedir,
|
||||
self.data_root,
|
||||
os.path.join(self.data_root, cate_file),
|
||||
iou_thresholds=[0.5],
|
||||
official=True)
|
||||
|
||||
result = culane_metric.eval_predictions(output_basedir,
|
||||
self.data_root,
|
||||
self.list_path,
|
||||
iou_thresholds=np.linspace(0.5, 0.95, 10),
|
||||
official=True)
|
||||
|
||||
return result[0.5]['F1']
|
||||
180
algorithms/lane_ufld/code/CLRNet-main/clrnet/datasets/llamas.py
Normal file
180
algorithms/lane_ufld/code/CLRNet-main/clrnet/datasets/llamas.py
Normal file
@@ -0,0 +1,180 @@
|
||||
import os
|
||||
import pickle as pkl
|
||||
import cv2
|
||||
|
||||
from .registry import DATASETS
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
from .base_dataset import BaseDataset
|
||||
|
||||
TRAIN_LABELS_DIR = 'labels/train'
|
||||
TEST_LABELS_DIR = 'labels/valid'
|
||||
TEST_IMGS_DIR = 'color_images/test'
|
||||
SPLIT_DIRECTORIES = {'train': 'labels/train', 'val': 'labels/valid'}
|
||||
from clrnet.utils.llamas_utils import get_horizontal_values_for_four_lanes
|
||||
import clrnet.utils.llamas_metric as llamas_metric
|
||||
|
||||
|
||||
@DATASETS.register_module
|
||||
class LLAMAS(BaseDataset):
|
||||
def __init__(self, data_root, split='train', processes=None, cfg=None):
|
||||
self.split = split
|
||||
self.data_root = data_root
|
||||
super().__init__(data_root, split, processes, cfg)
|
||||
if split != 'test' and split not in SPLIT_DIRECTORIES.keys():
|
||||
raise Exception('Split `{}` does not exist.'.format(split))
|
||||
if split != 'test':
|
||||
self.labels_dir = os.path.join(self.data_root,
|
||||
SPLIT_DIRECTORIES[split])
|
||||
|
||||
self.data_infos = []
|
||||
self.load_annotations()
|
||||
|
||||
def get_img_heigth(self, _):
|
||||
return self.cfg.ori_img_h
|
||||
|
||||
def get_img_width(self, _):
|
||||
return self.cfg.ori_img_w
|
||||
|
||||
def get_metrics(self, lanes, _):
|
||||
# Placeholders
|
||||
return [0] * len(lanes), [0] * len(lanes), [1] * len(lanes), [
|
||||
1
|
||||
] * len(lanes)
|
||||
|
||||
def get_img_path(self, json_path):
|
||||
# /foo/bar/test/folder/image_label.ext --> test/folder/image_label.ext
|
||||
base_name = '/'.join(json_path.split('/')[-3:])
|
||||
image_path = os.path.join(
|
||||
'color_images', base_name.replace('.json', '_color_rect.png'))
|
||||
return image_path
|
||||
|
||||
def get_img_name(self, json_path):
|
||||
base_name = (json_path.split('/')[-1]).replace('.json',
|
||||
'_color_rect.png')
|
||||
return base_name
|
||||
|
||||
def get_json_paths(self):
|
||||
json_paths = []
|
||||
for root, _, files in os.walk(self.labels_dir):
|
||||
for file in files:
|
||||
if file.endswith(".json"):
|
||||
json_paths.append(os.path.join(root, file))
|
||||
return json_paths
|
||||
|
||||
def load_annotations(self):
|
||||
# the labels are not public for the test set yet
|
||||
if self.split == 'test':
|
||||
imgs_dir = os.path.join(self.data_root, TEST_IMGS_DIR)
|
||||
self.data_infos = [{
|
||||
'img_path':
|
||||
os.path.join(root, file),
|
||||
'img_name':
|
||||
os.path.join(TEST_IMGS_DIR,
|
||||
root.split('/')[-1], file),
|
||||
'lanes': [],
|
||||
'relative_path':
|
||||
os.path.join(root.split('/')[-1], file)
|
||||
} for root, _, files in os.walk(imgs_dir) for file in files
|
||||
if file.endswith('.png')]
|
||||
self.data_infos = sorted(self.data_infos,
|
||||
key=lambda x: x['img_path'])
|
||||
return
|
||||
|
||||
# Waiting for the dataset to load is tedious, let's cache it
|
||||
os.makedirs('cache', exist_ok=True)
|
||||
cache_path = 'cache/llamas_{}.pkl'.format(self.split)
|
||||
if os.path.exists(cache_path):
|
||||
with open(cache_path, 'rb') as cache_file:
|
||||
self.data_infos = pkl.load(cache_file)
|
||||
self.max_lanes = max(
|
||||
len(anno['lanes']) for anno in self.data_infos)
|
||||
return
|
||||
|
||||
self.max_lanes = 0
|
||||
print("Searching annotation files...")
|
||||
json_paths = self.get_json_paths()
|
||||
print('{} annotations found.'.format(len(json_paths)))
|
||||
|
||||
for json_path in tqdm(json_paths):
|
||||
lanes = get_horizontal_values_for_four_lanes(json_path)
|
||||
lanes = [[(x, y) for x, y in zip(lane, range(self.cfg.ori_img_h))
|
||||
if x >= 0] for lane in lanes]
|
||||
lanes = [lane for lane in lanes if len(lane) > 0]
|
||||
lanes = [list(set(lane))
|
||||
for lane in lanes] # remove duplicated points
|
||||
lanes = [lane for lane in lanes
|
||||
if len(lane) > 2] # remove lanes with less than 2 points
|
||||
|
||||
lanes = [sorted(lane, key=lambda x: x[1])
|
||||
for lane in lanes] # sort by y
|
||||
lanes.sort(key=lambda lane: lane[0][0])
|
||||
mask_path = json_path.replace('.json', '.png')
|
||||
|
||||
# generate seg labels
|
||||
seg = np.zeros((717, 1276, 3))
|
||||
for i, lane in enumerate(lanes):
|
||||
for j in range(0, len(lane) - 1):
|
||||
cv2.line(seg, (round(lane[j][0]), lane[j][1]),
|
||||
(round(lane[j + 1][0]), lane[j + 1][1]),
|
||||
(i + 1, i + 1, i + 1),
|
||||
thickness=15)
|
||||
|
||||
cv2.imwrite(mask_path, seg)
|
||||
|
||||
relative_path = self.get_img_path(json_path)
|
||||
img_path = os.path.join(self.data_root, relative_path)
|
||||
self.max_lanes = max(self.max_lanes, len(lanes))
|
||||
self.data_infos.append({
|
||||
'img_path': img_path,
|
||||
'img_name': relative_path,
|
||||
'mask_path': mask_path,
|
||||
'lanes': lanes,
|
||||
'relative_path': relative_path
|
||||
})
|
||||
|
||||
with open(cache_path, 'wb') as cache_file:
|
||||
pkl.dump(self.data_infos, cache_file)
|
||||
|
||||
def assign_class_to_lanes(self, lanes):
|
||||
return {
|
||||
label: value
|
||||
for label, value in zip(['l0', 'l1', 'r0', 'r1'], lanes)
|
||||
}
|
||||
|
||||
def get_prediction_string(self, pred):
|
||||
ys = np.arange(300, 717, 1) / (self.cfg.ori_img_h - 1)
|
||||
out = []
|
||||
for lane in pred:
|
||||
xs = lane(ys)
|
||||
valid_mask = (xs >= 0) & (xs < 1)
|
||||
xs = xs * (self.cfg.ori_img_w - 1)
|
||||
lane_xs = xs[valid_mask]
|
||||
lane_ys = ys[valid_mask] * (self.cfg.ori_img_h - 1)
|
||||
lane_xs, lane_ys = lane_xs[::-1], lane_ys[::-1]
|
||||
lane_str = ' '.join([
|
||||
'{:.5f} {:.5f}'.format(x, y) for x, y in zip(lane_xs, lane_ys)
|
||||
])
|
||||
if lane_str != '':
|
||||
out.append(lane_str)
|
||||
|
||||
return '\n'.join(out)
|
||||
|
||||
def evaluate(self, predictions, output_basedir):
|
||||
print('Generating prediction output...')
|
||||
for idx, pred in enumerate(predictions):
|
||||
relative_path = self.data_infos[idx]['relative_path']
|
||||
output_filename = '/'.join(relative_path.split('/')[-2:]).replace(
|
||||
'_color_rect.png', '.lines.txt')
|
||||
output_filepath = os.path.join(output_basedir, output_filename)
|
||||
os.makedirs(os.path.dirname(output_filepath), exist_ok=True)
|
||||
output = self.get_prediction_string(pred)
|
||||
with open(output_filepath, 'w') as out_file:
|
||||
out_file.write(output)
|
||||
if self.split == 'test':
|
||||
return None
|
||||
result = llamas_metric.eval_predictions(output_basedir,
|
||||
self.labels_dir,
|
||||
iou_thresholds=np.linspace(0.5, 0.95, 10),
|
||||
unofficial=False)
|
||||
return result[0.5]['F1']
|
||||
188
algorithms/lane_ufld/code/CLRNet-main/clrnet/datasets/mufld.py
Normal file
188
algorithms/lane_ufld/code/CLRNet-main/clrnet/datasets/mufld.py
Normal file
@@ -0,0 +1,188 @@
|
||||
import os
|
||||
import os.path as osp
|
||||
import pickle as pkl
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from .base_dataset import BaseDataset
|
||||
from .registry import DATASETS
|
||||
from clrnet.utils.mask_to_lanes import lanes_from_mask, normalize_mask_labels
|
||||
from clrnet.utils.dataset_packs import resolve_list_file
|
||||
|
||||
DEFAULT_LIST = {
|
||||
"train": "list/train_gt.txt",
|
||||
"val": "list/val_gt.txt",
|
||||
"test": "list/test_gt.txt",
|
||||
}
|
||||
|
||||
|
||||
@DATASETS.register_module
|
||||
class MufldLane(BaseDataset):
|
||||
"""
|
||||
MUFLD / lane0_copy DATASET packs.
|
||||
list: <img_rel> <mask_rel> per line; lanes from mask or cached .lines.txt
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data_root,
|
||||
split,
|
||||
processes=None,
|
||||
cfg=None,
|
||||
list_file=None,
|
||||
):
|
||||
super().__init__(data_root, split, processes=processes, cfg=cfg)
|
||||
self.split = split
|
||||
if list_file is None:
|
||||
list_file = getattr(cfg, f"{split}_list_file", None)
|
||||
if list_file is None:
|
||||
list_file = resolve_list_file(cfg, split)
|
||||
if list_file is None:
|
||||
rel = DEFAULT_LIST.get(split, DEFAULT_LIST["train"])
|
||||
packs = getattr(cfg, "train_packs" if split == "train" else "val_packs", None)
|
||||
if packs:
|
||||
pack = packs[0] if isinstance(packs, (list, tuple)) else packs
|
||||
list_file = f"{pack}/{rel}"
|
||||
else:
|
||||
list_file = rel
|
||||
if osp.isabs(list_file):
|
||||
self.list_path = list_file
|
||||
else:
|
||||
self.list_path = osp.join(data_root, list_file)
|
||||
self.sample_ys = list(getattr(cfg, "sample_y", range(710, 150, -10)))
|
||||
self.num_lanes = getattr(cfg, "max_lanes", 4)
|
||||
self.lines_cache = getattr(cfg, "lines_cache_dir", "cache/mufld_lines")
|
||||
self.load_annotations()
|
||||
|
||||
def load_annotations(self):
|
||||
self.logger.info("Loading MufldLane annotations from %s", self.list_path)
|
||||
os.makedirs("cache", exist_ok=True)
|
||||
cache_key = self.list_path.replace("/", "_")
|
||||
cache_path = osp.join("cache", f"mufld_{self.split}_{cache_key}.pkl")
|
||||
if osp.exists(cache_path):
|
||||
with open(cache_path, "rb") as f:
|
||||
self.data_infos = pkl.load(f)
|
||||
self.max_lanes = max(len(a["lanes"]) for a in self.data_infos) if self.data_infos else self.num_lanes
|
||||
return
|
||||
|
||||
self.data_infos = []
|
||||
with open(self.list_path) as f:
|
||||
for line in f:
|
||||
parts = line.strip().split()
|
||||
if len(parts) < 2:
|
||||
continue
|
||||
info = self.load_annotation(parts)
|
||||
if info and len(info.get("lanes", [])) > 0:
|
||||
self.data_infos.append(info)
|
||||
|
||||
with open(cache_path, "wb") as f:
|
||||
pkl.dump(self.data_infos, f)
|
||||
self.max_lanes = max(len(a["lanes"]) for a in self.data_infos) if self.data_infos else self.num_lanes
|
||||
self.logger.info("Loaded %d samples, max_lanes=%d", len(self.data_infos), self.max_lanes)
|
||||
|
||||
def _lines_path(self, img_path: str) -> str:
|
||||
base = img_path[:-4] if img_path.lower().endswith((".jpg", ".png")) else img_path
|
||||
cache_root = osp.join(self.data_root, self.lines_cache)
|
||||
rel = osp.relpath(base, self.data_root)
|
||||
return osp.join(cache_root, rel + ".lines.txt")
|
||||
|
||||
def load_annotation(self, line):
|
||||
img_line = line[0].lstrip("/")
|
||||
mask_line = line[1].lstrip("/")
|
||||
img_path = osp.join(self.data_root, img_line)
|
||||
mask_path = osp.join(self.data_root, mask_line)
|
||||
infos = {
|
||||
"img_name": img_line,
|
||||
"img_path": img_path,
|
||||
"mask_path": mask_path,
|
||||
}
|
||||
if len(line) > 2:
|
||||
infos["lane_exist"] = np.array([int(x) for x in line[2:]])
|
||||
|
||||
lines_path = self._lines_path(img_path)
|
||||
if osp.isfile(lines_path):
|
||||
with open(lines_path) as f:
|
||||
data = [list(map(float, ln.split())) for ln in f.readlines() if ln.strip()]
|
||||
lanes = [
|
||||
[(lane[i], lane[i + 1]) for i in range(0, len(lane), 2) if lane[i] >= 0 and lane[i + 1] >= 0]
|
||||
for lane in data
|
||||
]
|
||||
elif osp.isfile(mask_path):
|
||||
mask = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
|
||||
if mask is None:
|
||||
return None
|
||||
if mask.ndim > 2:
|
||||
mask = mask[:, :, 0]
|
||||
lanes = lanes_from_mask(mask, self.sample_ys, self.num_lanes)
|
||||
if getattr(self.cfg, "write_lines_cache", False):
|
||||
os.makedirs(osp.dirname(lines_path), exist_ok=True)
|
||||
with open(lines_path, "w") as out:
|
||||
for lane in lanes:
|
||||
out.write(" ".join(f"{x:.5f} {y:.5f}" for x, y in lane) + "\n")
|
||||
else:
|
||||
return None
|
||||
|
||||
lanes = [lane for lane in lanes if len(lane) > 2]
|
||||
lanes = [sorted(lane, key=lambda x: x[1]) for lane in lanes]
|
||||
infos["lanes"] = lanes
|
||||
return infos
|
||||
|
||||
def __getitem__(self, idx):
|
||||
data_info = self.data_infos[idx]
|
||||
img = cv2.imread(data_info["img_path"])
|
||||
if img is None:
|
||||
raise FileNotFoundError(data_info["img_path"])
|
||||
img = img[self.cfg.cut_height :, :, :]
|
||||
sample = data_info.copy()
|
||||
sample.update({"img": img})
|
||||
|
||||
if self.training:
|
||||
label = cv2.imread(sample["mask_path"], cv2.IMREAD_UNCHANGED)
|
||||
if label is None:
|
||||
raise FileNotFoundError(sample["mask_path"])
|
||||
if label.ndim > 2:
|
||||
label = label[:, :, 0]
|
||||
label = normalize_mask_labels(label.squeeze(), self.num_lanes)
|
||||
label = label[self.cfg.cut_height :, :]
|
||||
sample.update({"mask": label})
|
||||
|
||||
if self.cfg.cut_height != 0:
|
||||
new_lanes = []
|
||||
for lane in sample["lanes"]:
|
||||
new_lanes.append([(p[0], p[1] - self.cfg.cut_height) for p in lane])
|
||||
sample.update({"lanes": new_lanes})
|
||||
|
||||
from mmcv.parallel import DataContainer as DC
|
||||
from clrnet.datasets.process import Process
|
||||
|
||||
sample = self.processes(sample)
|
||||
meta = {"full_img_path": data_info["img_path"], "img_name": data_info["img_name"]}
|
||||
sample.update({"meta": DC(meta, cpu_only=True)})
|
||||
return sample
|
||||
|
||||
def get_prediction_string(self, pred):
|
||||
ys = np.array(self.sample_ys) / self.cfg.ori_img_h
|
||||
out = []
|
||||
for lane in pred:
|
||||
xs = lane(ys)
|
||||
valid = (xs >= 0) & (xs < 1)
|
||||
xs = xs[valid] * self.cfg.ori_img_w
|
||||
lane_ys = ys[valid] * self.cfg.ori_img_h
|
||||
xs, lane_ys = xs[::-1], lane_ys[::-1]
|
||||
s = " ".join(f"{x:.5f} {y:.5f}" for x, y in zip(xs, lane_ys))
|
||||
if s:
|
||||
out.append(s)
|
||||
return "\n".join(out)
|
||||
|
||||
def evaluate(self, predictions, output_basedir):
|
||||
os.makedirs(output_basedir, exist_ok=True)
|
||||
for idx, pred in enumerate(predictions):
|
||||
rel = self.data_infos[idx]["img_name"]
|
||||
out_dir = osp.join(output_basedir, osp.dirname(rel))
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
out_file = osp.join(out_dir, osp.basename(rel)[:-4] + ".lines.txt")
|
||||
with open(out_file, "w") as f:
|
||||
f.write(self.get_prediction_string(pred))
|
||||
self.logger.info("Wrote predictions under %s (MUFLD: no CULane official eval)", output_basedir)
|
||||
return 0.0
|
||||
@@ -0,0 +1,21 @@
|
||||
from .transforms import (RandomLROffsetLABEL, RandomUDoffsetLABEL, Resize,
|
||||
RandomCrop, CenterCrop, RandomRotation, RandomBlur,
|
||||
RandomHorizontalFlip, Normalize, ToTensor)
|
||||
|
||||
from .generate_lane_line import GenerateLaneLine
|
||||
from .process import Process
|
||||
|
||||
__all__ = [
|
||||
'Process',
|
||||
'RandomLROffsetLABEL',
|
||||
'RandomUDoffsetLABEL',
|
||||
'Resize',
|
||||
'RandomCrop',
|
||||
'CenterCrop',
|
||||
'RandomRotation',
|
||||
'RandomBlur',
|
||||
'RandomHorizontalFlip',
|
||||
'Normalize',
|
||||
'ToTensor',
|
||||
'GenerateLaneLine',
|
||||
]
|
||||
@@ -0,0 +1,218 @@
|
||||
import math
|
||||
import numpy as np
|
||||
import cv2
|
||||
import imgaug.augmenters as iaa
|
||||
from imgaug.augmentables.lines import LineString, LineStringsOnImage
|
||||
from imgaug.augmentables.segmaps import SegmentationMapsOnImage
|
||||
from scipy.interpolate import InterpolatedUnivariateSpline
|
||||
from clrnet.datasets.process.transforms import CLRTransforms
|
||||
|
||||
from ..registry import PROCESS
|
||||
|
||||
|
||||
@PROCESS.register_module
|
||||
class GenerateLaneLine(object):
|
||||
def __init__(self, transforms=None, cfg=None, training=True):
|
||||
self.transforms = transforms
|
||||
self.img_w, self.img_h = cfg.img_w, cfg.img_h
|
||||
self.num_points = cfg.num_points
|
||||
self.n_offsets = cfg.num_points
|
||||
self.n_strips = cfg.num_points - 1
|
||||
self.strip_size = self.img_h / self.n_strips
|
||||
self.max_lanes = cfg.max_lanes
|
||||
self.offsets_ys = np.arange(self.img_h, -1, -self.strip_size)
|
||||
self.training = training
|
||||
|
||||
if transforms is None:
|
||||
transforms = CLRTransforms(self.img_h, self.img_w)
|
||||
|
||||
if transforms is not None:
|
||||
img_transforms = []
|
||||
for aug in transforms:
|
||||
p = aug['p']
|
||||
if aug['name'] != 'OneOf':
|
||||
img_transforms.append(
|
||||
iaa.Sometimes(p=p,
|
||||
then_list=getattr(
|
||||
iaa,
|
||||
aug['name'])(**aug['parameters'])))
|
||||
else:
|
||||
img_transforms.append(
|
||||
iaa.Sometimes(
|
||||
p=p,
|
||||
then_list=iaa.OneOf([
|
||||
getattr(iaa,
|
||||
aug_['name'])(**aug_['parameters'])
|
||||
for aug_ in aug['transforms']
|
||||
])))
|
||||
else:
|
||||
img_transforms = []
|
||||
self.transform = iaa.Sequential(img_transforms)
|
||||
|
||||
def lane_to_linestrings(self, lanes):
|
||||
lines = []
|
||||
for lane in lanes:
|
||||
lines.append(LineString(lane))
|
||||
|
||||
return lines
|
||||
|
||||
def sample_lane(self, points, sample_ys):
|
||||
# this function expects the points to be sorted
|
||||
points = np.array(points)
|
||||
if not np.all(points[1:, 1] < points[:-1, 1]):
|
||||
raise Exception('Annotaion points have to be sorted')
|
||||
x, y = points[:, 0], points[:, 1]
|
||||
|
||||
# interpolate points inside domain
|
||||
assert len(points) > 1
|
||||
interp = InterpolatedUnivariateSpline(y[::-1],
|
||||
x[::-1],
|
||||
k=min(3,
|
||||
len(points) - 1))
|
||||
domain_min_y = y.min()
|
||||
domain_max_y = y.max()
|
||||
sample_ys_inside_domain = sample_ys[(sample_ys >= domain_min_y)
|
||||
& (sample_ys <= domain_max_y)]
|
||||
assert len(sample_ys_inside_domain) > 0
|
||||
interp_xs = interp(sample_ys_inside_domain)
|
||||
|
||||
# extrapolate lane to the bottom of the image with a straight line using the 2 points closest to the bottom
|
||||
two_closest_points = points[:2]
|
||||
extrap = np.polyfit(two_closest_points[:, 1],
|
||||
two_closest_points[:, 0],
|
||||
deg=1)
|
||||
extrap_ys = sample_ys[sample_ys > domain_max_y]
|
||||
extrap_xs = np.polyval(extrap, extrap_ys)
|
||||
all_xs = np.hstack((extrap_xs, interp_xs))
|
||||
|
||||
# separate between inside and outside points
|
||||
inside_mask = (all_xs >= 0) & (all_xs < self.img_w)
|
||||
xs_inside_image = all_xs[inside_mask]
|
||||
xs_outside_image = all_xs[~inside_mask]
|
||||
|
||||
return xs_outside_image, xs_inside_image
|
||||
|
||||
def filter_lane(self, lane):
|
||||
assert lane[-1][1] <= lane[0][1]
|
||||
filtered_lane = []
|
||||
used = set()
|
||||
for p in lane:
|
||||
if p[1] not in used:
|
||||
filtered_lane.append(p)
|
||||
used.add(p[1])
|
||||
|
||||
return filtered_lane
|
||||
|
||||
def transform_annotation(self, anno, img_wh=None):
|
||||
img_w, img_h = self.img_w, self.img_h
|
||||
|
||||
old_lanes = anno['lanes']
|
||||
|
||||
# removing lanes with less than 2 points
|
||||
old_lanes = filter(lambda x: len(x) > 1, old_lanes)
|
||||
# sort lane points by Y (bottom to top of the image)
|
||||
old_lanes = [sorted(lane, key=lambda x: -x[1]) for lane in old_lanes]
|
||||
# remove points with same Y (keep first occurrence)
|
||||
old_lanes = [self.filter_lane(lane) for lane in old_lanes]
|
||||
# normalize the annotation coordinates
|
||||
old_lanes = [[[
|
||||
x * self.img_w / float(img_w), y * self.img_h / float(img_h)
|
||||
] for x, y in lane] for lane in old_lanes]
|
||||
# create tranformed annotations
|
||||
lanes = np.ones(
|
||||
(self.max_lanes, 2 + 1 + 1 + 2 + self.n_offsets), dtype=np.float32
|
||||
) * -1e5 # 2 scores, 1 start_y, 1 start_x, 1 theta, 1 length, S+1 coordinates
|
||||
lanes_endpoints = np.ones((self.max_lanes, 2))
|
||||
# lanes are invalid by default
|
||||
lanes[:, 0] = 1
|
||||
lanes[:, 1] = 0
|
||||
for lane_idx, lane in enumerate(old_lanes):
|
||||
if lane_idx >= self.max_lanes:
|
||||
break
|
||||
|
||||
try:
|
||||
xs_outside_image, xs_inside_image = self.sample_lane(
|
||||
lane, self.offsets_ys)
|
||||
except AssertionError:
|
||||
continue
|
||||
if len(xs_inside_image) <= 1:
|
||||
continue
|
||||
all_xs = np.hstack((xs_outside_image, xs_inside_image))
|
||||
lanes[lane_idx, 0] = 0
|
||||
lanes[lane_idx, 1] = 1
|
||||
lanes[lane_idx, 2] = len(xs_outside_image) / self.n_strips
|
||||
lanes[lane_idx, 3] = xs_inside_image[0]
|
||||
|
||||
thetas = []
|
||||
for i in range(1, len(xs_inside_image)):
|
||||
theta = math.atan(
|
||||
i * self.strip_size /
|
||||
(xs_inside_image[i] - xs_inside_image[0] + 1e-5)) / math.pi
|
||||
theta = theta if theta > 0 else 1 - abs(theta)
|
||||
thetas.append(theta)
|
||||
|
||||
theta_far = sum(thetas) / len(thetas)
|
||||
|
||||
# lanes[lane_idx,
|
||||
# 4] = (theta_closest + theta_far) / 2 # averaged angle
|
||||
lanes[lane_idx, 4] = theta_far
|
||||
lanes[lane_idx, 5] = len(xs_inside_image)
|
||||
lanes[lane_idx, 6:6 + len(all_xs)] = all_xs
|
||||
lanes_endpoints[lane_idx, 0] = (len(all_xs) - 1) / self.n_strips
|
||||
lanes_endpoints[lane_idx, 1] = xs_inside_image[-1]
|
||||
|
||||
new_anno = {
|
||||
'label': lanes,
|
||||
'old_anno': anno,
|
||||
'lane_endpoints': lanes_endpoints
|
||||
}
|
||||
return new_anno
|
||||
|
||||
def linestrings_to_lanes(self, lines):
|
||||
lanes = []
|
||||
for line in lines:
|
||||
lanes.append(line.coords)
|
||||
|
||||
return lanes
|
||||
|
||||
def __call__(self, sample):
|
||||
img_org = sample['img']
|
||||
line_strings_org = self.lane_to_linestrings(sample['lanes'])
|
||||
line_strings_org = LineStringsOnImage(line_strings_org,
|
||||
shape=img_org.shape)
|
||||
|
||||
for i in range(30):
|
||||
if self.training:
|
||||
mask_org = SegmentationMapsOnImage(sample['mask'],
|
||||
shape=img_org.shape)
|
||||
img, line_strings, seg = self.transform(
|
||||
image=img_org.copy().astype(np.uint8),
|
||||
line_strings=line_strings_org,
|
||||
segmentation_maps=mask_org)
|
||||
else:
|
||||
img, line_strings = self.transform(
|
||||
image=img_org.copy().astype(np.uint8),
|
||||
line_strings=line_strings_org)
|
||||
line_strings.clip_out_of_image_()
|
||||
new_anno = {'lanes': self.linestrings_to_lanes(line_strings)}
|
||||
try:
|
||||
annos = self.transform_annotation(new_anno,
|
||||
img_wh=(self.img_w,
|
||||
self.img_h))
|
||||
label = annos['label']
|
||||
lane_endpoints = annos['lane_endpoints']
|
||||
break
|
||||
except:
|
||||
if (i + 1) == 30:
|
||||
self.logger.critical(
|
||||
'Transform annotation failed 30 times :(')
|
||||
exit()
|
||||
|
||||
sample['img'] = img.astype(np.float32) / 255.
|
||||
sample['lane_line'] = label
|
||||
sample['lanes_endpoints'] = lane_endpoints
|
||||
sample['gt_points'] = new_anno['lanes']
|
||||
sample['seg'] = seg.get_arr() if self.training else np.zeros(
|
||||
img_org.shape)
|
||||
|
||||
return sample
|
||||
@@ -0,0 +1,48 @@
|
||||
import collections
|
||||
|
||||
from clrnet.utils import build_from_cfg
|
||||
|
||||
from ..registry import PROCESS
|
||||
|
||||
|
||||
class Process(object):
|
||||
"""Compose multiple process sequentially.
|
||||
Args:
|
||||
process (Sequence[dict | callable]): Sequence of process object or
|
||||
config dict to be composed.
|
||||
"""
|
||||
def __init__(self, processes, cfg):
|
||||
assert isinstance(processes, collections.abc.Sequence)
|
||||
self.processes = []
|
||||
for process in processes:
|
||||
if isinstance(process, dict):
|
||||
process = build_from_cfg(process,
|
||||
PROCESS,
|
||||
default_args=dict(cfg=cfg))
|
||||
self.processes.append(process)
|
||||
elif callable(process):
|
||||
self.processes.append(process)
|
||||
else:
|
||||
raise TypeError('process must be callable or a dict')
|
||||
|
||||
def __call__(self, data):
|
||||
"""Call function to apply processes sequentially.
|
||||
Args:
|
||||
data (dict): A result dict contains the data to process.
|
||||
Returns:
|
||||
dict: Processed data.
|
||||
"""
|
||||
|
||||
for t in self.processes:
|
||||
data = t(data)
|
||||
if data is None:
|
||||
return None
|
||||
return data
|
||||
|
||||
def __repr__(self):
|
||||
format_string = self.__class__.__name__ + '('
|
||||
for t in self.processes:
|
||||
format_string += '\n'
|
||||
format_string += f' {t}'
|
||||
format_string += '\n)'
|
||||
return format_string
|
||||
@@ -0,0 +1,311 @@
|
||||
import random
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import numbers
|
||||
import collections
|
||||
from PIL import Image
|
||||
|
||||
from ..registry import PROCESS
|
||||
|
||||
|
||||
def to_tensor(data):
|
||||
"""Convert objects of various python types to :obj:`torch.Tensor`.
|
||||
|
||||
Supported types are: :class:`numpy.ndarray`, :class:`torch.Tensor`,
|
||||
:class:`Sequence`, :class:`int` and :class:`float`.
|
||||
|
||||
Args:
|
||||
data (torch.Tensor | numpy.ndarray | Sequence | int | float): Data to
|
||||
be converted.
|
||||
"""
|
||||
|
||||
if isinstance(data, torch.Tensor):
|
||||
return data
|
||||
elif isinstance(data, np.ndarray):
|
||||
return torch.from_numpy(data)
|
||||
elif isinstance(data, int):
|
||||
return torch.LongTensor([data])
|
||||
elif isinstance(data, float):
|
||||
return torch.FloatTensor([data])
|
||||
else:
|
||||
raise TypeError(f'type {type(data)} cannot be converted to tensor.')
|
||||
|
||||
|
||||
@PROCESS.register_module
|
||||
class ToTensor(object):
|
||||
"""Convert some results to :obj:`torch.Tensor` by given keys.
|
||||
|
||||
Args:
|
||||
keys (Sequence[str]): Keys that need to be converted to Tensor.
|
||||
"""
|
||||
def __init__(self, keys=['img', 'mask'], cfg=None):
|
||||
self.keys = keys
|
||||
|
||||
def __call__(self, sample):
|
||||
data = {}
|
||||
if len(sample['img'].shape) < 3:
|
||||
sample['img'] = np.expand_dims(img, -1)
|
||||
for key in self.keys:
|
||||
if key == 'img_metas' or key == 'gt_masks' or key == 'lane_line':
|
||||
data[key] = sample[key]
|
||||
continue
|
||||
data[key] = to_tensor(sample[key])
|
||||
data['img'] = data['img'].permute(2, 0, 1)
|
||||
return data
|
||||
|
||||
def __repr__(self):
|
||||
return self.__class__.__name__ + f'(keys={self.keys})'
|
||||
|
||||
|
||||
@PROCESS.register_module
|
||||
class RandomLROffsetLABEL(object):
|
||||
def __init__(self, max_offset, cfg=None):
|
||||
self.max_offset = max_offset
|
||||
|
||||
def __call__(self, sample):
|
||||
img = sample['img']
|
||||
label = sample['mask']
|
||||
offset = np.random.randint(-self.max_offset, self.max_offset)
|
||||
h, w = img.shape[:2]
|
||||
|
||||
img = np.array(img)
|
||||
if offset > 0:
|
||||
img[:, offset:, :] = img[:, 0:w - offset, :]
|
||||
img[:, :offset, :] = 0
|
||||
if offset < 0:
|
||||
real_offset = -offset
|
||||
img[:, 0:w - real_offset, :] = img[:, real_offset:, :]
|
||||
img[:, w - real_offset:, :] = 0
|
||||
|
||||
label = np.array(label)
|
||||
if offset > 0:
|
||||
label[:, offset:] = label[:, 0:w - offset]
|
||||
label[:, :offset] = 0
|
||||
if offset < 0:
|
||||
offset = -offset
|
||||
label[:, 0:w - offset] = label[:, offset:]
|
||||
label[:, w - offset:] = 0
|
||||
sample['img'] = img
|
||||
sample['mask'] = label
|
||||
|
||||
return sample
|
||||
|
||||
|
||||
@PROCESS.register_module
|
||||
class RandomUDoffsetLABEL(object):
|
||||
def __init__(self, max_offset, cfg=None):
|
||||
self.max_offset = max_offset
|
||||
|
||||
def __call__(self, sample):
|
||||
img = sample['img']
|
||||
label = sample['mask']
|
||||
offset = np.random.randint(-self.max_offset, self.max_offset)
|
||||
h, w = img.shape[:2]
|
||||
|
||||
img = np.array(img)
|
||||
if offset > 0:
|
||||
img[offset:, :, :] = img[0:h - offset, :, :]
|
||||
img[:offset, :, :] = 0
|
||||
if offset < 0:
|
||||
real_offset = -offset
|
||||
img[0:h - real_offset, :, :] = img[real_offset:, :, :]
|
||||
img[h - real_offset:, :, :] = 0
|
||||
|
||||
label = np.array(label)
|
||||
if offset > 0:
|
||||
label[offset:, :] = label[0:h - offset, :]
|
||||
label[:offset, :] = 0
|
||||
if offset < 0:
|
||||
offset = -offset
|
||||
label[0:h - offset, :] = label[offset:, :]
|
||||
label[h - offset:, :] = 0
|
||||
sample['img'] = img
|
||||
sample['mask'] = label
|
||||
return sample
|
||||
|
||||
|
||||
@PROCESS.register_module
|
||||
class Resize(object):
|
||||
def __init__(self, size, cfg=None):
|
||||
assert (isinstance(size, collections.Iterable) and len(size) == 2)
|
||||
self.size = size
|
||||
|
||||
def __call__(self, sample):
|
||||
out = list()
|
||||
sample['img'] = cv2.resize(sample['img'],
|
||||
self.size,
|
||||
interpolation=cv2.INTER_CUBIC)
|
||||
if 'mask' in sample:
|
||||
sample['mask'] = cv2.resize(sample['mask'],
|
||||
self.size,
|
||||
interpolation=cv2.INTER_NEAREST)
|
||||
return sample
|
||||
|
||||
|
||||
@PROCESS.register_module
|
||||
class RandomCrop(object):
|
||||
def __init__(self, size, cfg=None):
|
||||
if isinstance(size, numbers.Number):
|
||||
self.size = (int(size), int(size))
|
||||
else:
|
||||
self.size = size
|
||||
|
||||
def __call__(self, img_group):
|
||||
h, w = img_group[0].shape[0:2]
|
||||
th, tw = self.size
|
||||
|
||||
out_images = list()
|
||||
h1 = random.randint(0, max(0, h - th))
|
||||
w1 = random.randint(0, max(0, w - tw))
|
||||
h2 = min(h1 + th, h)
|
||||
w2 = min(w1 + tw, w)
|
||||
|
||||
for img in img_group:
|
||||
assert (img.shape[0] == h and img.shape[1] == w)
|
||||
out_images.append(img[h1:h2, w1:w2, ...])
|
||||
return out_images
|
||||
|
||||
|
||||
@PROCESS.register_module
|
||||
class CenterCrop(object):
|
||||
def __init__(self, size, cfg=None):
|
||||
if isinstance(size, numbers.Number):
|
||||
self.size = (int(size), int(size))
|
||||
else:
|
||||
self.size = size
|
||||
|
||||
def __call__(self, img_group):
|
||||
h, w = img_group[0].shape[0:2]
|
||||
th, tw = self.size
|
||||
|
||||
out_images = list()
|
||||
h1 = max(0, int((h - th) / 2))
|
||||
w1 = max(0, int((w - tw) / 2))
|
||||
h2 = min(h1 + th, h)
|
||||
w2 = min(w1 + tw, w)
|
||||
|
||||
for img in img_group:
|
||||
assert (img.shape[0] == h and img.shape[1] == w)
|
||||
out_images.append(img[h1:h2, w1:w2, ...])
|
||||
return out_images
|
||||
|
||||
|
||||
@PROCESS.register_module
|
||||
class RandomRotation(object):
|
||||
def __init__(self,
|
||||
degree=(-10, 10),
|
||||
interpolation=(cv2.INTER_LINEAR, cv2.INTER_NEAREST),
|
||||
padding=None,
|
||||
cfg=None):
|
||||
self.degree = degree
|
||||
self.interpolation = interpolation
|
||||
self.padding = padding
|
||||
if self.padding is None:
|
||||
self.padding = [0, 0]
|
||||
|
||||
def _rotate_img(self, sample, map_matrix):
|
||||
h, w = sample['img'].shape[0:2]
|
||||
sample['img'] = cv2.warpAffine(sample['img'],
|
||||
map_matrix, (w, h),
|
||||
flags=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=self.padding)
|
||||
|
||||
def _rotate_mask(self, sample, map_matrix):
|
||||
if 'mask' not in sample:
|
||||
return
|
||||
h, w = sample['mask'].shape[0:2]
|
||||
sample['mask'] = cv2.warpAffine(sample['mask'],
|
||||
map_matrix, (w, h),
|
||||
flags=cv2.INTER_NEAREST,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=self.padding)
|
||||
|
||||
def __call__(self, sample):
|
||||
v = random.random()
|
||||
if v < 0.5:
|
||||
degree = random.uniform(self.degree[0], self.degree[1])
|
||||
h, w = sample['img'].shape[0:2]
|
||||
center = (w / 2, h / 2)
|
||||
map_matrix = cv2.getRotationMatrix2D(center, degree, 1.0)
|
||||
self._rotate_img(sample, map_matrix)
|
||||
self._rotate_mask(sample, map_matrix)
|
||||
return sample
|
||||
|
||||
|
||||
@PROCESS.register_module
|
||||
class RandomBlur(object):
|
||||
def __init__(self, applied, cfg=None):
|
||||
self.applied = applied
|
||||
|
||||
def __call__(self, img_group):
|
||||
assert (len(self.applied) == len(img_group))
|
||||
v = random.random()
|
||||
if v < 0.5:
|
||||
out_images = []
|
||||
for img, a in zip(img_group, self.applied):
|
||||
if a:
|
||||
img = cv2.GaussianBlur(img, (5, 5),
|
||||
random.uniform(1e-6, 0.6))
|
||||
out_images.append(img)
|
||||
if len(img.shape) > len(out_images[-1].shape):
|
||||
out_images[-1] = out_images[-1][
|
||||
..., np.newaxis] # single channel image
|
||||
return out_images
|
||||
else:
|
||||
return img_group
|
||||
|
||||
|
||||
@PROCESS.register_module
|
||||
class RandomHorizontalFlip(object):
|
||||
"""Randomly horizontally flips the given numpy Image with a probability of 0.5
|
||||
"""
|
||||
def __init__(self, cfg=None):
|
||||
pass
|
||||
|
||||
def __call__(self, sample):
|
||||
v = random.random()
|
||||
if v < 0.5:
|
||||
sample['img'] = np.fliplr(sample['img'])
|
||||
if 'mask' in sample: sample['mask'] = np.fliplr(sample['mask'])
|
||||
return sample
|
||||
|
||||
|
||||
@PROCESS.register_module
|
||||
class Normalize(object):
|
||||
def __init__(self, img_norm, cfg=None):
|
||||
self.mean = np.array(img_norm['mean'], dtype=np.float32)
|
||||
self.std = np.array(img_norm['std'], dtype=np.float32)
|
||||
|
||||
def __call__(self, sample):
|
||||
m = self.mean
|
||||
s = self.std
|
||||
img = sample['img']
|
||||
if len(m) == 1:
|
||||
img = img - np.array(m) # single channel image
|
||||
img = img / np.array(s)
|
||||
else:
|
||||
img = img - np.array(m)[np.newaxis, np.newaxis, ...]
|
||||
img = img / np.array(s)[np.newaxis, np.newaxis, ...]
|
||||
sample['img'] = img
|
||||
|
||||
return sample
|
||||
|
||||
|
||||
def CLRTransforms(img_h, img_w):
|
||||
return [
|
||||
dict(name='Resize',
|
||||
parameters=dict(size=dict(height=img_h, width=img_w)),
|
||||
p=1.0),
|
||||
dict(name='HorizontalFlip', parameters=dict(p=1.0), p=0.5),
|
||||
dict(name='Affine',
|
||||
parameters=dict(translate_percent=dict(x=(-0.1, 0.1),
|
||||
y=(-0.1, 0.1)),
|
||||
rotate=(-10, 10),
|
||||
scale=(0.8, 1.2)),
|
||||
p=0.7),
|
||||
dict(name='Resize',
|
||||
parameters=dict(size=dict(height=img_h, width=img_w)),
|
||||
p=1.0),
|
||||
]
|
||||
@@ -0,0 +1,54 @@
|
||||
from clrnet.utils import Registry, build_from_cfg
|
||||
|
||||
import torch
|
||||
from functools import partial
|
||||
import numpy as np
|
||||
import random
|
||||
from mmcv.parallel import collate
|
||||
|
||||
DATASETS = Registry('datasets')
|
||||
PROCESS = Registry('process')
|
||||
|
||||
|
||||
def build(cfg, registry, default_args=None):
|
||||
if isinstance(cfg, list):
|
||||
modules = [
|
||||
build_from_cfg(cfg_, registry, default_args) for cfg_ in cfg
|
||||
]
|
||||
return nn.Sequential(*modules)
|
||||
else:
|
||||
return build_from_cfg(cfg, registry, default_args)
|
||||
|
||||
|
||||
def build_dataset(split_cfg, cfg):
|
||||
return build(split_cfg, DATASETS, default_args=dict(cfg=cfg))
|
||||
|
||||
|
||||
def worker_init_fn(worker_id, seed):
|
||||
worker_seed = worker_id + seed
|
||||
np.random.seed(worker_seed)
|
||||
random.seed(worker_seed)
|
||||
|
||||
|
||||
def build_dataloader(split_cfg, cfg, is_train=True):
|
||||
if is_train:
|
||||
shuffle = True
|
||||
else:
|
||||
shuffle = False
|
||||
|
||||
dataset = build_dataset(split_cfg, cfg)
|
||||
|
||||
init_fn = partial(worker_init_fn, seed=cfg.seed)
|
||||
|
||||
samples_per_gpu = cfg.batch_size // cfg.gpus
|
||||
data_loader = torch.utils.data.DataLoader(
|
||||
dataset,
|
||||
batch_size=cfg.batch_size,
|
||||
shuffle=shuffle,
|
||||
num_workers=cfg.workers,
|
||||
pin_memory=False,
|
||||
drop_last=False,
|
||||
collate_fn=partial(collate, samples_per_gpu=samples_per_gpu),
|
||||
worker_init_fn=init_fn)
|
||||
|
||||
return data_loader
|
||||
@@ -0,0 +1,100 @@
|
||||
import os.path as osp
|
||||
import numpy as np
|
||||
import cv2
|
||||
import os
|
||||
import json
|
||||
import torchvision
|
||||
from .base_dataset import BaseDataset
|
||||
from clrnet.utils.tusimple_metric import LaneEval
|
||||
from .registry import DATASETS
|
||||
import logging
|
||||
import random
|
||||
|
||||
SPLIT_FILES = {
|
||||
'trainval':
|
||||
['label_data_0313.json', 'label_data_0601.json', 'label_data_0531.json'],
|
||||
'train': ['label_data_0313.json', 'label_data_0601.json'],
|
||||
'val': ['label_data_0531.json'],
|
||||
'test': ['test_label.json'],
|
||||
}
|
||||
|
||||
|
||||
@DATASETS.register_module
|
||||
class TuSimple(BaseDataset):
|
||||
def __init__(self, data_root, split, processes=None, cfg=None):
|
||||
super().__init__(data_root, split, processes, cfg)
|
||||
self.anno_files = SPLIT_FILES[split]
|
||||
self.load_annotations()
|
||||
self.h_samples = list(range(160, 720, 10))
|
||||
|
||||
def load_annotations(self):
|
||||
self.logger.info('Loading TuSimple annotations...')
|
||||
self.data_infos = []
|
||||
max_lanes = 0
|
||||
for anno_file in self.anno_files:
|
||||
anno_file = osp.join(self.data_root, anno_file)
|
||||
with open(anno_file, 'r') as anno_obj:
|
||||
lines = anno_obj.readlines()
|
||||
for line in lines:
|
||||
data = json.loads(line)
|
||||
y_samples = data['h_samples']
|
||||
gt_lanes = data['lanes']
|
||||
mask_path = data['raw_file'].replace('clips',
|
||||
'seg_label')[:-3] + 'png'
|
||||
lanes = [[(x, y) for (x, y) in zip(lane, y_samples) if x >= 0]
|
||||
for lane in gt_lanes]
|
||||
lanes = [lane for lane in lanes if len(lane) > 0]
|
||||
max_lanes = max(max_lanes, len(lanes))
|
||||
self.data_infos.append({
|
||||
'img_path':
|
||||
osp.join(self.data_root, data['raw_file']),
|
||||
'img_name':
|
||||
data['raw_file'],
|
||||
'mask_path':
|
||||
osp.join(self.data_root, mask_path),
|
||||
'lanes':
|
||||
lanes,
|
||||
})
|
||||
|
||||
if self.training:
|
||||
random.shuffle(self.data_infos)
|
||||
self.max_lanes = max_lanes
|
||||
|
||||
def pred2lanes(self, pred):
|
||||
ys = np.array(self.h_samples) / self.cfg.ori_img_h
|
||||
lanes = []
|
||||
for lane in pred:
|
||||
xs = lane(ys)
|
||||
invalid_mask = xs < 0
|
||||
lane = (xs * self.cfg.ori_img_w).astype(int)
|
||||
lane[invalid_mask] = -2
|
||||
lanes.append(lane.tolist())
|
||||
|
||||
return lanes
|
||||
|
||||
def pred2tusimpleformat(self, idx, pred, runtime):
|
||||
runtime *= 1000. # s to ms
|
||||
img_name = self.data_infos[idx]['img_name']
|
||||
lanes = self.pred2lanes(pred)
|
||||
output = {'raw_file': img_name, 'lanes': lanes, 'run_time': runtime}
|
||||
return json.dumps(output)
|
||||
|
||||
def save_tusimple_predictions(self, predictions, filename, runtimes=None):
|
||||
if runtimes is None:
|
||||
runtimes = np.ones(len(predictions)) * 1.e-3
|
||||
lines = []
|
||||
for idx, (prediction, runtime) in enumerate(zip(predictions,
|
||||
runtimes)):
|
||||
line = self.pred2tusimpleformat(idx, prediction, runtime)
|
||||
lines.append(line)
|
||||
with open(filename, 'w') as output_file:
|
||||
output_file.write('\n'.join(lines))
|
||||
|
||||
def evaluate(self, predictions, output_basedir, runtimes=None):
|
||||
pred_filename = os.path.join(output_basedir,
|
||||
'tusimple_predictions.json')
|
||||
self.save_tusimple_predictions(predictions, pred_filename, runtimes)
|
||||
result, acc = LaneEval.bench_one_submit(pred_filename,
|
||||
self.cfg.test_json_file)
|
||||
self.logger.info(result)
|
||||
return acc
|
||||
Reference in New Issue
Block a user