feat: HSAP platform v2 — modular navigation, quality review, audit log, world model simulation
Major changes: - New frontend (platform/web/): Vite + React 18 + TypeScript + Tailwind - 4-module navigation: 数据送标 / 模型管理 / 车队管理 / 系统管理 - Data catalog with charts (DMS/ADAS/Lane 3-tab view) - Quality review workflow (标注质检): Good/Fine/Bad scoring with auto-advance - Audit enhancements: batch operations, rejection categories, Feishu notifications - Operation audit log (操作日志) - World model simulation studio (仿真工坊) - Dataset version management with snapshots and diff - ADAS 7-class dataset integration (138K images organized + compressed) - User management with Feishu integration and pagination - CRUD/search/filter on all pages, card layout redesign - PIL-optimized image overlay rendering - Auto-snapshot on build, in_review workflow stage - Removed embedded algorithm code (now in workspace)
This commit is contained in:
@@ -0,0 +1,13 @@
|
||||
from .segmentation import PASCAL_VOC_Segmentation, CityscapesSegmentation, SYNTHIA_Segmentation, GTAV_Segmentation
|
||||
from .lane_as_segmentation import TuSimpleAsSegmentation, CULaneAsSegmentation, LLAMAS_AsSegmentation
|
||||
from .lane_as_bezier import TuSimpleAsBezier, CULaneAsBezier, LLAMAS_AsBezier, Curvelanes_AsBezier
|
||||
from .tusimple import TuSimple
|
||||
from .tusimple_vis import TuSimpleVis
|
||||
from .culane import CULane
|
||||
from .culane_vis import CULaneVis
|
||||
from .llamas import LLAMAS
|
||||
from .llamas_vis import LLAMAS_Vis
|
||||
from .image_folder import ImageFolderDataset
|
||||
from .video import VideoLoader
|
||||
from .utils import dict_collate_fn
|
||||
from .builder import DATASETS
|
||||
@@ -0,0 +1,22 @@
|
||||
import torchvision
|
||||
import os
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
# BDD100K direct loading (work with the segmentation style lists)
|
||||
class CULane(torchvision.datasets.VisionDataset):
|
||||
def __init__(self, root, image_set, transforms=None, transform=None, target_transform=None,
|
||||
ppl=0, gap=0, start=0):
|
||||
super().__init__(root, transforms, transform, target_transform)
|
||||
pass
|
||||
|
||||
def __getitem__(self, index):
|
||||
# Return x (input image) & y (L lane with N coordinates (x, y) as np.array (L x N x 2))
|
||||
# Empty coordinates are marked by (-2, -2)
|
||||
# If just testing,
|
||||
# y is the filename to store prediction
|
||||
pass
|
||||
|
||||
def __len__(self):
|
||||
pass
|
||||
@@ -0,0 +1,3 @@
|
||||
from ..registry import SimpleRegistry
|
||||
|
||||
DATASETS = SimpleRegistry()
|
||||
@@ -0,0 +1,74 @@
|
||||
import os
|
||||
import pickle
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
|
||||
from .utils import LaneKeypointDataset
|
||||
from .builder import DATASETS
|
||||
|
||||
|
||||
# CULane direct loading (work with the segmentation style lists)
|
||||
@DATASETS.register()
|
||||
class CULane(LaneKeypointDataset):
|
||||
colors = [
|
||||
[0, 0, 0], # background
|
||||
[0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 0],
|
||||
[0, 0, 0] # ignore
|
||||
]
|
||||
|
||||
def __init__(self, root, image_set, transforms=None, transform=None, target_transform=None,
|
||||
ppl=31, gap=10, start=290, padding_mask=False, is_process=True):
|
||||
super().__init__(root, transforms, transform, target_transform, ppl, gap, start, padding_mask, image_set,
|
||||
is_process)
|
||||
|
||||
self._check()
|
||||
|
||||
# Data list
|
||||
with open(os.path.join(root, 'lists', image_set + '.txt'), "r") as f:
|
||||
contents = [x.strip() for x in f.readlines()]
|
||||
|
||||
# Load filenames
|
||||
if image_set == 'test' or image_set == 'val': # Test
|
||||
self.images = [os.path.join(root, x + '.jpg') for x in contents]
|
||||
self.targets = [os.path.join('./output', x + '.lines.txt') for x in contents]
|
||||
else: # Train
|
||||
self.images = [os.path.join(root, x[:x.find(' ')] + '.jpg') for x in contents]
|
||||
self.targets = []
|
||||
print('Loading targets into memory...')
|
||||
processed_file = os.path.join(root, 'train_processed_targets')
|
||||
if os.path.exists(processed_file):
|
||||
with open(processed_file, 'rb') as f:
|
||||
self.targets = pickle.load(f)
|
||||
else:
|
||||
print('Pre-processing will only be performed for 1 time, please wait ~10 minutes.')
|
||||
for x in tqdm(contents):
|
||||
with open(os.path.join(root, x[:x.find(' ')] + '.lines.txt'), 'r') as f:
|
||||
self.targets.append(self._load_target(f.readlines()))
|
||||
with open(processed_file, 'wb') as f:
|
||||
pickle.dump(self.targets, f)
|
||||
print('Loading complete.')
|
||||
|
||||
assert len(self.targets) == len(self.images)
|
||||
|
||||
def _load_target(self, lines):
|
||||
# Read file content to lists (file content could be empty or variable number of lanes)
|
||||
target = np.array([[[-2.0, self.start + i * self.gap] for i in range(self.ppl)]
|
||||
for _ in range(len(lines))], dtype=np.float32)
|
||||
for i in range(len(lines)): # lines=[] will end this immediately
|
||||
temp = [float(k) for k in lines[i].strip().split(' ')]
|
||||
for j in range(int(len(temp) / 2)):
|
||||
x = temp[2 * j]
|
||||
y = temp[2 * j + 1]
|
||||
target[i][target[i][:, 1] == y] = [x, y]
|
||||
|
||||
return target
|
||||
|
||||
@staticmethod
|
||||
def load_target_xy(lines):
|
||||
# A direct loading of JSON file to a list of N x 2 numpy arrays
|
||||
target = []
|
||||
for line in lines:
|
||||
temp = [float(x) for x in line.strip().split(' ')]
|
||||
target.append(np.array(temp).reshape(-1, 2))
|
||||
|
||||
return target
|
||||
@@ -0,0 +1,47 @@
|
||||
import os
|
||||
|
||||
from .image_folder_lane_base import ImageFolderLaneBase
|
||||
from .builder import DATASETS
|
||||
|
||||
|
||||
# Visualization version of CULane
|
||||
@DATASETS.register()
|
||||
class CULaneVis(ImageFolderLaneBase):
|
||||
def __init__(self, root_dataset, root_output, root_keypoint, image_set, transforms=None,
|
||||
keypoint_process_fn=None, use_gt=True):
|
||||
super().__init__(root_dataset, root_output, transforms, keypoint_process_fn)
|
||||
self.image_set = image_set
|
||||
|
||||
self._check()
|
||||
|
||||
# Data list
|
||||
with open(os.path.join(root_dataset, 'lists', image_set + '.txt'), "r") as f:
|
||||
contents = [x.strip() for x in f.readlines()]
|
||||
|
||||
# Load filenames
|
||||
if image_set == 'test' or image_set == 'val': # Test
|
||||
self.images = [os.path.join(root_dataset, x + '.jpg') for x in contents]
|
||||
if use_gt:
|
||||
self.gt_keypoints = [os.path.join(root_dataset, x + '.lines.txt') for x in contents]
|
||||
self.filenames = [x + '.jpg' for x in contents]
|
||||
if root_keypoint is not None:
|
||||
self.keypoints = [os.path.join(root_keypoint, x + '.lines.txt') for x in contents]
|
||||
else: # Train
|
||||
self.images = [os.path.join(root_dataset, x[:x.find(' ')] + '.jpg') for x in contents]
|
||||
if use_gt:
|
||||
self.gt_keypoints = [os.path.join(root_dataset, x[:x.find(' ')] + '.lines.txt') for x in contents]
|
||||
self.filenames = [x[:x.find(' ')] + '.jpg' for x in contents]
|
||||
if root_keypoint is not None:
|
||||
self.keypoints = [os.path.join(root_keypoint, x + '.lines.txt') for x in contents]
|
||||
|
||||
self.make_sub_dirs()
|
||||
|
||||
assert len(self.images) == len(self.gt_keypoints)
|
||||
if self.keypoints is not None:
|
||||
assert len(self.images) == len(self.keypoints)
|
||||
|
||||
def _check(self):
|
||||
# Checks
|
||||
if self.image_set not in ['train', 'val', 'test']:
|
||||
raise ValueError
|
||||
assert self.output_dir != self.root, 'Avoid overwriting your dataset!'
|
||||
@@ -0,0 +1,75 @@
|
||||
import torchvision
|
||||
import os
|
||||
from PIL import Image
|
||||
|
||||
from ..transforms import functional as F, ToTensor
|
||||
from .builder import DATASETS
|
||||
from .image_folder_lane_base import ImageFolderLaneBase
|
||||
|
||||
|
||||
# Load a directory of images for inference
|
||||
@DATASETS.register()
|
||||
class ImageFolderDataset(torchvision.datasets.VisionDataset):
|
||||
def __init__(self, root_image, root_output, root_target=None, transforms=None,
|
||||
target_process_fn=None, image_suffix='', target_suffix=''):
|
||||
super().__init__(root_image, transforms, None, None)
|
||||
self.output_dir = root_output
|
||||
self.filenames = []
|
||||
self.images = []
|
||||
self.targets = None if root_target is None else []
|
||||
self.target_process_fn = target_process_fn
|
||||
for filename in sorted(os.listdir(root_image)):
|
||||
suffix_pos = filename.rfind(image_suffix)
|
||||
if suffix_pos != -1:
|
||||
middle_name = filename[:suffix_pos]
|
||||
self.filenames.append(filename)
|
||||
self.images.append(os.path.join(root_image, filename))
|
||||
if self.targets is not None:
|
||||
self.targets.append(os.path.join(root_target, middle_name + target_suffix))
|
||||
|
||||
def __getitem__(self, index):
|
||||
# Return transformed image / original image / save filename / label (if exist)
|
||||
img = Image.open(self.images[index]).convert('RGB')
|
||||
filename = os.path.join(self.output_dir, self.filenames[index])
|
||||
original_img = F.to_tensor(img).clone()
|
||||
|
||||
# Transforms
|
||||
if self.transforms is not None:
|
||||
img = self.transforms(img)
|
||||
|
||||
# Process potential target
|
||||
target = None
|
||||
if self.targets is not None:
|
||||
target = self.target_process_fn(self.targets[index])
|
||||
|
||||
return img, original_img, {
|
||||
'filename': filename,
|
||||
'target': target
|
||||
}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.images)
|
||||
|
||||
|
||||
# Load a directory of images for lane inference
|
||||
@DATASETS.register()
|
||||
class ImageFolderLaneDataset(ImageFolderLaneBase):
|
||||
def __init__(self, root_image, root_output, root_keypoint=None, root_gt_keypoint=None, root_mask=None,
|
||||
transforms=None, keypoint_process_fn=None,
|
||||
image_suffix='', keypoint_suffix='.txt', gt_keypoint_suffix='.txt', mask_suffix=''):
|
||||
super().__init__(root_image, root_output, transforms, keypoint_process_fn)
|
||||
self.keypoints = None if root_keypoint is None else []
|
||||
self.gt_keypoints = None if root_gt_keypoint is None else []
|
||||
self.masks = None if root_mask is None else []
|
||||
for filename in sorted(os.listdir(root_image)):
|
||||
suffix_pos = filename.rfind(image_suffix)
|
||||
if suffix_pos != -1:
|
||||
middle_name = filename[:suffix_pos]
|
||||
self.filenames.append(filename)
|
||||
self.images.append(os.path.join(root_image, filename))
|
||||
if self.keypoints is not None:
|
||||
self.keypoints.append(os.path.join(root_keypoint, middle_name + keypoint_suffix))
|
||||
if self.gt_keypoints is not None:
|
||||
self.gt_keypoints.append(os.path.join(root_gt_keypoint, middle_name + gt_keypoint_suffix))
|
||||
if self.masks is not None:
|
||||
self.masks.append(os.path.join(root_mask, middle_name + mask_suffix))
|
||||
@@ -0,0 +1,60 @@
|
||||
import torchvision
|
||||
import os
|
||||
from PIL import Image
|
||||
|
||||
from ..transforms import functional as F, ToTensor
|
||||
|
||||
|
||||
# Base class for lane image folder datasets (usually for visualizations)
|
||||
class ImageFolderLaneBase(torchvision.datasets.VisionDataset):
|
||||
def __init__(self, root=None, root_output=None, transforms=None, keypoint_process_fn=None):
|
||||
super().__init__(root, transforms, None, None)
|
||||
self.output_dir = root_output
|
||||
os.makedirs(self.output_dir, exist_ok=True)
|
||||
self.filenames = []
|
||||
self.images = []
|
||||
self.keypoints = None
|
||||
self.gt_keypoints = None
|
||||
self.masks = None
|
||||
self.keypoint_process_fn = keypoint_process_fn
|
||||
|
||||
def __getitem__(self, index):
|
||||
# Return transformed image / original image / save filename / labels (if exist)
|
||||
img = Image.open(self.images[index]).convert('RGB')
|
||||
filename = os.path.join(self.output_dir, self.filenames[index])
|
||||
original_img = F.to_tensor(img).clone()
|
||||
mask = None
|
||||
if self.masks is not None:
|
||||
w, h = F._get_image_size(img)
|
||||
mask = ToTensor.label_to_tensor(
|
||||
F.resize(Image.open(self.masks[index]), size=[h, w], interpolation=Image.NEAREST)
|
||||
)
|
||||
|
||||
# Transforms
|
||||
if self.transforms is not None:
|
||||
img = self.transforms(img)
|
||||
|
||||
# Process potential target
|
||||
keypoint = None
|
||||
gt_keypoint = None
|
||||
if self.keypoints is not None:
|
||||
keypoint = self.keypoint_process_fn(self.keypoints[index])
|
||||
if self.gt_keypoints is not None:
|
||||
gt_keypoint = self.keypoint_process_fn(self.gt_keypoints[index])
|
||||
|
||||
return img, original_img, {
|
||||
'filename': filename,
|
||||
'keypoint': keypoint,
|
||||
'gt_keypoint': gt_keypoint,
|
||||
'mask': mask
|
||||
}
|
||||
|
||||
def make_sub_dirs(self):
|
||||
# Make sub dirs
|
||||
for f in self.filenames:
|
||||
dir_name = os.path.join(self.output_dir, f[:f.rfind('/')])
|
||||
if not os.path.exists(dir_name):
|
||||
os.makedirs(dir_name)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.images)
|
||||
@@ -0,0 +1,206 @@
|
||||
import os
|
||||
import torch
|
||||
import torchvision
|
||||
import json
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from .builder import DATASETS
|
||||
from ..curve_utils import BezierSampler, get_valid_points
|
||||
|
||||
|
||||
class _BezierLaneDataset(torchvision.datasets.VisionDataset):
|
||||
# BezierLaneNet dataset, includes binary seg labels
|
||||
keypoint_color = [0, 0, 0]
|
||||
|
||||
def __init__(self, root, image_set='train', transforms=None, transform=None, target_transform=None,
|
||||
order=3, num_sample_points=100, aux_segmentation=False):
|
||||
super().__init__(root, transforms, transform, target_transform)
|
||||
self.aux_segmentation = aux_segmentation
|
||||
self.bezier_sampler = BezierSampler(order=order, num_sample_points=num_sample_points)
|
||||
if image_set == 'valfast':
|
||||
raise NotImplementedError('valfast Not supported yet!')
|
||||
elif image_set == 'test' or image_set == 'val': # Different format (without lane existence annotations)
|
||||
self.test = 2
|
||||
elif image_set == 'val_train':
|
||||
self.test = 3
|
||||
else:
|
||||
self.test = 0
|
||||
|
||||
self.init_dataset(root)
|
||||
|
||||
if image_set != 'valfast':
|
||||
self.bezier_labels = os.path.join(self.bezier_labels_dir, image_set + '_' + str(order) + '.json')
|
||||
elif image_set == 'valfast':
|
||||
raise ValueError
|
||||
|
||||
self.image_set = image_set
|
||||
self.splits_dir = os.path.join(root, 'lists')
|
||||
self._init_all()
|
||||
|
||||
def init_dataset(self, root):
|
||||
raise NotImplementedError
|
||||
|
||||
def __getitem__(self, index):
|
||||
# Return x (input image) & y (mask image, i.e. pixel-wise supervision) & lane existence (a list),
|
||||
# if not just testing,
|
||||
# else just return input image.
|
||||
img = Image.open(self.images[index]).convert('RGB')
|
||||
if self.test >= 2:
|
||||
target = self.masks[index]
|
||||
else:
|
||||
if self.aux_segmentation:
|
||||
target = {'keypoints': self.beziers[index],
|
||||
'segmentation_mask': Image.open(self.masks[index])}
|
||||
else:
|
||||
target = {'keypoints': self.beziers[index]}
|
||||
|
||||
# Transforms
|
||||
if self.transforms is not None:
|
||||
img, target = self.transforms(img, target)
|
||||
|
||||
if self.test == 0:
|
||||
target = self._post_process(target)
|
||||
|
||||
return img, target
|
||||
|
||||
def __len__(self):
|
||||
return len(self.images)
|
||||
|
||||
def loader_bezier(self):
|
||||
results = []
|
||||
with open(self.bezier_labels, 'r') as f:
|
||||
results += [json.loads(x.strip()) for x in f.readlines()]
|
||||
beziers = []
|
||||
for lanes in results:
|
||||
temp_lane = []
|
||||
for lane in lanes['bezier_control_points']:
|
||||
temp_cps = []
|
||||
for i in range(0, len(lane), 2):
|
||||
temp_cps.append([lane[i], lane[i + 1]])
|
||||
temp_lane.append(temp_cps)
|
||||
beziers.append(np.array(temp_lane, dtype=np.float32))
|
||||
return beziers
|
||||
|
||||
def _init_all(self):
|
||||
# Got the lists from 4 datasets to be in the same format
|
||||
data_list = 'train.txt' if self.image_set == 'val_train' else self.image_set + '.txt'
|
||||
split_f = os.path.join(self.splits_dir, data_list)
|
||||
with open(split_f, "r") as f:
|
||||
contents = [x.strip() for x in f.readlines()]
|
||||
if self.test == 2: # Test
|
||||
self.images = [os.path.join(self.image_dir, x + self.image_suffix) for x in contents]
|
||||
self.masks = [os.path.join(self.output_prefix, x + self.output_suffix) for x in contents]
|
||||
elif self.test == 3: # Test
|
||||
self.images = [os.path.join(self.image_dir, x[:x.find(' ')] + self.image_suffix) for x in contents]
|
||||
self.masks = [os.path.join(self.output_prefix, x[:x.find(' ')] + self.output_suffix) for x in contents]
|
||||
elif self.test == 1: # Val
|
||||
self.images = [os.path.join(self.image_dir, x[:x.find(' ')] + self.image_suffix) for x in contents]
|
||||
self.masks = [os.path.join(self.mask_dir, x[:x.find(' ')] + '.png') for x in contents]
|
||||
else: # Train
|
||||
self.images = [os.path.join(self.image_dir, x[:x.find(' ')] + self.image_suffix) for x in contents]
|
||||
if self.aux_segmentation:
|
||||
self.masks = [os.path.join(self.mask_dir, x[:x.find(' ')] + '.png') for x in contents]
|
||||
self.beziers = self.loader_bezier()
|
||||
|
||||
def _post_process(self, target, ignore_seg_index=255):
|
||||
# Get sample points and delete invalid lines (< 2 points)
|
||||
if target['keypoints'].numel() != 0: # No-lane cases can be handled in loss computation
|
||||
sample_points = self.bezier_sampler.get_sample_points(target['keypoints'])
|
||||
valid_lanes = get_valid_points(sample_points).sum(dim=-1) >= 2
|
||||
target['keypoints'] = target['keypoints'][valid_lanes]
|
||||
target['sample_points'] = sample_points[valid_lanes]
|
||||
else:
|
||||
target['sample_points'] = torch.tensor([], dtype=target['keypoints'].dtype)
|
||||
|
||||
if 'segmentation_mask' in target.keys(): # Map to binary (0 1 255)
|
||||
positive_mask = (target['segmentation_mask'] > 0) * (target['segmentation_mask'] != ignore_seg_index)
|
||||
target['segmentation_mask'][positive_mask] = 1
|
||||
|
||||
return target
|
||||
|
||||
|
||||
# TuSimple
|
||||
@DATASETS.register()
|
||||
class TuSimpleAsBezier(_BezierLaneDataset):
|
||||
colors = [
|
||||
[0, 0, 0], # background
|
||||
[255, 0, 255], [0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 0], [0, 255, 255],
|
||||
[0, 0, 0] # ignore
|
||||
]
|
||||
|
||||
def init_dataset(self, root):
|
||||
self.image_dir = os.path.join(root, 'clips')
|
||||
self.bezier_labels_dir = os.path.join(root, 'bezier_labels')
|
||||
self.mask_dir = os.path.join(root, 'segGT6')
|
||||
self.output_prefix = 'clips'
|
||||
self.output_suffix = '.jpg'
|
||||
self.image_suffix = '.jpg'
|
||||
|
||||
|
||||
# CULane
|
||||
@DATASETS.register()
|
||||
class CULaneAsBezier(_BezierLaneDataset):
|
||||
colors = [
|
||||
[0, 0, 0], # background
|
||||
[0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 0],
|
||||
[0, 0, 0] # ignore
|
||||
]
|
||||
|
||||
def init_dataset(self, root):
|
||||
self.image_dir = root
|
||||
self.bezier_labels_dir = os.path.join(root, 'bezier_labels')
|
||||
self.mask_dir = os.path.join(root, 'laneseg_label_w16')
|
||||
self.output_prefix = './output'
|
||||
self.output_suffix = '.lines.txt'
|
||||
self.image_suffix = '.jpg'
|
||||
if not os.path.exists(self.output_prefix):
|
||||
os.makedirs(self.output_prefix)
|
||||
|
||||
|
||||
# LLAMAS
|
||||
@DATASETS.register()
|
||||
class LLAMAS_AsBezier(_BezierLaneDataset):
|
||||
colors = [
|
||||
[0, 0, 0], # background
|
||||
[0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 0],
|
||||
[0, 0, 0] # ignore
|
||||
]
|
||||
|
||||
def init_dataset(self, root):
|
||||
self.image_dir = os.path.join(root, 'color_images')
|
||||
self.bezier_labels_dir = os.path.join(root, 'bezier_labels')
|
||||
self.mask_dir = os.path.join(root, 'laneseg_labels')
|
||||
self.output_prefix = './output'
|
||||
self.output_suffix = '.lines.txt'
|
||||
self.image_suffix = '.png'
|
||||
if not os.path.exists(self.output_prefix):
|
||||
os.makedirs(self.output_prefix)
|
||||
|
||||
|
||||
# Curvelanes
|
||||
@DATASETS.register()
|
||||
class Curvelanes_AsBezier(CULaneAsBezier):
|
||||
# TODO: Match formats
|
||||
colors = []
|
||||
|
||||
def _init_all(self):
|
||||
# Got the lists from 4 datasets to be in the same format
|
||||
data_list = 'train.txt' if self.image_set == 'val_train' else self.image_set + '.txt'
|
||||
split_f = os.path.join(self.splits_dir, data_list)
|
||||
with open(split_f, "r") as f:
|
||||
contents = [x.strip() for x in f.readlines()]
|
||||
if self.test == 2: # Test
|
||||
self.images = [os.path.join(self.image_dir, x + self.image_suffix) for x in contents]
|
||||
self.masks = [os.path.join(self.output_prefix, x + self.output_suffix) for x in contents]
|
||||
elif self.test == 3: # Test
|
||||
self.images = [os.path.join(self.image_dir, x[:x.find(' ')] + self.image_suffix) for x in contents]
|
||||
self.masks = [os.path.join(self.output_prefix, x[:x.find(' ')] + self.output_suffix) for x in contents]
|
||||
elif self.test == 1: # Val
|
||||
self.images = [os.path.join(self.image_dir, x[:x.find(' ')] + self.image_suffix) for x in contents]
|
||||
self.masks = [os.path.join(self.mask_dir, x[:x.find(' ')] + '.png') for x in contents]
|
||||
else: # Train
|
||||
self.images = [os.path.join(self.image_dir, x + self.image_suffix) for x in contents]
|
||||
if self.aux_segmentation:
|
||||
self.masks = [os.path.join(self.mask_dir, x[:x.find(' ')] + '.png') for x in contents]
|
||||
self.beziers = self.loader_bezier()
|
||||
@@ -0,0 +1,120 @@
|
||||
import os
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torchvision.datasets import VisionDataset
|
||||
|
||||
from .builder import DATASETS
|
||||
|
||||
|
||||
# Lane detection as segmentation
|
||||
class _StandardLaneDetectionDataset(VisionDataset):
|
||||
keypoint_color = [0, 0, 0]
|
||||
|
||||
def __init__(self, root, image_set, transforms=None):
|
||||
super().__init__(root, transforms, None, None)
|
||||
if image_set == 'valfast':
|
||||
self.test = 1
|
||||
elif image_set == 'test' or image_set == 'val': # Different format (without lane existence annotations)
|
||||
self.test = 2
|
||||
else:
|
||||
self.test = 0
|
||||
|
||||
self.init_dataset(root)
|
||||
self.image_set = image_set
|
||||
self.splits_dir = os.path.join(root, 'lists')
|
||||
|
||||
self._init_all()
|
||||
|
||||
assert (len(self.images) == len(self.masks))
|
||||
|
||||
def init_dataset(self, root):
|
||||
raise NotImplementedError
|
||||
|
||||
def __getitem__(self, index):
|
||||
# Return x (input image) & y (mask image, i.e. pixel-wise supervision) & lane existence (a list),
|
||||
# if not just testing,
|
||||
# else just return input image.
|
||||
img = Image.open(self.images[index]).convert('RGB')
|
||||
if self.test == 2:
|
||||
target = self.masks[index]
|
||||
elif self.test == 1:
|
||||
target = Image.open(self.masks[index])
|
||||
else:
|
||||
target = Image.open(self.masks[index])
|
||||
lane_existence = torch.tensor(self.lane_existences[index]).float()
|
||||
|
||||
# Transforms
|
||||
if self.transforms is not None:
|
||||
img, target = self.transforms(img, target)
|
||||
if self.test > 0:
|
||||
return img, target
|
||||
else:
|
||||
return img, target, lane_existence
|
||||
|
||||
def __len__(self):
|
||||
return len(self.images)
|
||||
|
||||
def _init_all(self):
|
||||
# Got the lists from 2 datasets to be in the same format
|
||||
split_f = os.path.join(self.splits_dir, self.image_set + '.txt')
|
||||
with open(split_f, "r") as f:
|
||||
contents = [x.strip() for x in f.readlines()]
|
||||
if self.test == 2: # Test
|
||||
self.images = [os.path.join(self.image_dir, x + self.image_suffix) for x in contents]
|
||||
self.masks = [os.path.join(self.output_prefix, x + self.output_suffix) for x in contents]
|
||||
elif self.test == 1: # Val
|
||||
self.images = [os.path.join(self.image_dir, x[:x.find(' ')] + self.image_suffix) for x in contents]
|
||||
self.masks = [os.path.join(self.mask_dir, x[:x.find(' ')] + '.png') for x in contents]
|
||||
else: # Train
|
||||
self.images = [os.path.join(self.image_dir, x[:x.find(' ')] + self.image_suffix) for x in contents]
|
||||
self.masks = [os.path.join(self.mask_dir, x[:x.find(' ')] + '.png') for x in contents]
|
||||
self.lane_existences = [list(map(int, x[x.find(' '):].split())) for x in contents]
|
||||
|
||||
|
||||
# TuSimple
|
||||
@DATASETS.register()
|
||||
class TuSimpleAsSegmentation(_StandardLaneDetectionDataset):
|
||||
colors = [
|
||||
[0, 0, 0], # background
|
||||
[255, 0, 255], [0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 0], [0, 255, 255],
|
||||
[0, 0, 0] # ignore
|
||||
]
|
||||
|
||||
def init_dataset(self, root):
|
||||
self.image_dir = os.path.join(root, 'clips')
|
||||
self.mask_dir = os.path.join(root, 'segGT6')
|
||||
self.output_prefix = 'clips'
|
||||
self.output_suffix = '.jpg'
|
||||
self.image_suffix = '.jpg'
|
||||
|
||||
|
||||
# CULane
|
||||
@DATASETS.register()
|
||||
class CULaneAsSegmentation(_StandardLaneDetectionDataset):
|
||||
colors = [
|
||||
[0, 0, 0], # background
|
||||
[0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 0],
|
||||
[0, 0, 0] # ignore
|
||||
]
|
||||
|
||||
def init_dataset(self, root):
|
||||
self.image_dir = root
|
||||
self.mask_dir = os.path.join(root, 'laneseg_label_w16')
|
||||
self.output_prefix = './output'
|
||||
self.output_suffix = '.lines.txt'
|
||||
self.image_suffix = '.jpg'
|
||||
if not os.path.exists(self.output_prefix):
|
||||
os.makedirs(self.output_prefix)
|
||||
|
||||
|
||||
# LLAMAS
|
||||
@DATASETS.register()
|
||||
class LLAMAS_AsSegmentation(CULaneAsSegmentation):
|
||||
def init_dataset(self, root):
|
||||
self.image_dir = os.path.join(root, 'color_images')
|
||||
self.mask_dir = os.path.join(root, 'laneseg_labels')
|
||||
self.output_prefix = './output'
|
||||
self.output_suffix = '.lines.txt'
|
||||
self.image_suffix = '.png'
|
||||
if not os.path.exists(self.output_prefix):
|
||||
os.makedirs(self.output_prefix)
|
||||
@@ -0,0 +1,74 @@
|
||||
import os
|
||||
import pickle
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
|
||||
from .utils import LaneKeypointDataset
|
||||
from .builder import DATASETS
|
||||
|
||||
|
||||
# LLAMAS direct loading (similar with culane)
|
||||
@DATASETS.register()
|
||||
class LLAMAS(LaneKeypointDataset):
|
||||
colors = [
|
||||
[0, 0, 0], # background
|
||||
[0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 0],
|
||||
[0, 0, 0] # ignore
|
||||
]
|
||||
|
||||
def __init__(self, root, image_set, transforms=None, transform=None, target_transform=None,
|
||||
ppl=417, gap=1, start=300, padding_mask=False):
|
||||
super().__init__(root, transforms, transform, target_transform, ppl, gap, start, padding_mask, image_set)
|
||||
|
||||
self._check()
|
||||
|
||||
self.images_path = os.path.join(root, 'color_images')
|
||||
# Data list
|
||||
with open(os.path.join(root, 'lists', image_set + '.txt'), "r") as f:
|
||||
contents = [x.strip() for x in f.readlines()]
|
||||
|
||||
# Load filenames
|
||||
if image_set == 'test' or image_set == 'val': # Test
|
||||
self.images = [os.path.join(self.images_path, x + '.png') for x in contents]
|
||||
self.targets = [os.path.join('./output', x + '.lines.txt') for x in contents]
|
||||
else: # Train
|
||||
self.images = [os.path.join(self.images_path, x[:x.find(' ')] + '.png') for x in contents]
|
||||
self.targets = []
|
||||
print('Loading targets into memory...')
|
||||
processed_file = os.path.join(root, 'train_processed_targets')
|
||||
if os.path.exists(processed_file):
|
||||
with open(processed_file, 'rb') as f:
|
||||
self.targets = pickle.load(f)
|
||||
else:
|
||||
print('Pre-processing will only be performed for 1 time, please wait ~10 minutes.')
|
||||
for x in tqdm(contents):
|
||||
with open(os.path.join(self.images_path, x[:x.find(' ')] + '.lines.txt'), 'r') as f:
|
||||
self.targets.append(self._load_target(f.readlines()))
|
||||
with open(processed_file, 'wb') as f:
|
||||
pickle.dump(self.targets, f)
|
||||
print('Loading complete.')
|
||||
|
||||
assert len(self.targets) == len(self.images)
|
||||
|
||||
def _load_target(self, lines):
|
||||
# Read file content to lists (file content could be empty or variable number of lanes)
|
||||
target = np.array([[[-2.0, self.start + i * self.gap] for i in range(self.ppl)]
|
||||
for _ in range(len(lines))], dtype=np.float32)
|
||||
for i in range(len(lines)): # lines=[] will end this immediately
|
||||
temp = [float(k) for k in lines[i].strip().split(' ')]
|
||||
for j in range(int(len(temp) / 2)):
|
||||
x = temp[2 * j]
|
||||
y = temp[2 * j + 1]
|
||||
target[i][target[i][:, 1] == y] = [x, y]
|
||||
|
||||
return target
|
||||
|
||||
@staticmethod
|
||||
def load_target_xy(lines):
|
||||
# A direct loading of JSON file to a list of N x 2 numpy arrays
|
||||
target = []
|
||||
for line in lines:
|
||||
temp = [float(x) for x in line.strip().split(' ')]
|
||||
target.append(np.array(temp).reshape(-1, 2))
|
||||
|
||||
return target
|
||||
@@ -0,0 +1,50 @@
|
||||
import os
|
||||
|
||||
from .image_folder_lane_base import ImageFolderLaneBase
|
||||
from .builder import DATASETS
|
||||
|
||||
|
||||
# Visualization version of CULane
|
||||
@DATASETS.register()
|
||||
class LLAMAS_Vis(ImageFolderLaneBase):
|
||||
def __init__(self, root_dataset, root_output, root_keypoint, image_set, transforms=None,
|
||||
keypoint_process_fn=None, use_gt=True):
|
||||
super().__init__(root_dataset, root_output, transforms, keypoint_process_fn)
|
||||
self.image_set = image_set
|
||||
self.images_path = os.path.join(root_dataset, 'color_images')
|
||||
|
||||
self._check()
|
||||
|
||||
# Data list
|
||||
with open(os.path.join(root_dataset, 'lists', image_set + '.txt'), "r") as f:
|
||||
contents = [x.strip() for x in f.readlines()]
|
||||
|
||||
# Load filenames
|
||||
if image_set == 'test' or image_set == 'val': # Test
|
||||
self.images = [os.path.join(self.images_path, x + '.png') for x in contents]
|
||||
if use_gt:
|
||||
self.gt_keypoints = [os.path.join(self.images_path, x + '.lines.txt') for x in contents]
|
||||
self.filenames = [x + '.png' for x in contents]
|
||||
if root_keypoint is not None:
|
||||
self.keypoints = [os.path.join(root_keypoint, x.replace("_color_rect", "") + '.lines.txt')
|
||||
for x in contents]
|
||||
else: # Train
|
||||
self.images = [os.path.join(self.images_path, x[:x.find(' ')] + '.png') for x in contents]
|
||||
if use_gt:
|
||||
self.gt_keypoints = [os.path.join(self.images_path, x[:x.find(' ')] + '.lines.txt') for x in contents]
|
||||
self.filenames = [x[:x.find(' ')] + '.png' for x in contents]
|
||||
if root_keypoint is not None:
|
||||
self.keypoints = [os.path.join(root_keypoint, x.replace("_color_rect", "") + '.lines.txt')
|
||||
for x in contents]
|
||||
|
||||
self.make_sub_dirs()
|
||||
|
||||
assert len(self.images) == len(self.gt_keypoints)
|
||||
if self.keypoints is not None:
|
||||
assert len(self.images) == len(self.keypoints)
|
||||
|
||||
def _check(self):
|
||||
# Checks
|
||||
if self.image_set not in ['train', 'val', 'test']:
|
||||
raise ValueError
|
||||
assert self.output_dir != self.root, 'Avoid overwriting your dataset!'
|
||||
@@ -0,0 +1,154 @@
|
||||
import os
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from torchvision.datasets import VisionDataset
|
||||
|
||||
from .builder import DATASETS
|
||||
|
||||
|
||||
# Reimplemented based on torchvision.datasets.VOCSegmentation
|
||||
class _StandardSegmentationDataset(VisionDataset):
|
||||
def __init__(self, root, image_set, transforms=None, mask_type='.png'):
|
||||
super().__init__(root, transforms, None, None)
|
||||
self.mask_type = mask_type
|
||||
self.images = self.masks = []
|
||||
self.init_dataset(root, image_set)
|
||||
assert (len(self.images) == len(self.masks))
|
||||
|
||||
def init_dataset(self, root, image_set):
|
||||
raise NotImplementedError
|
||||
|
||||
def __getitem__(self, index):
|
||||
img = Image.open(self.images[index]).convert('RGB')
|
||||
# Return x(input image) & y(mask images as a list)
|
||||
# Supports .png & .npy
|
||||
target = Image.open(self.masks[index]) if '.png' in self.masks[index] else np.load(self.masks[index])
|
||||
|
||||
# Transforms
|
||||
if self.transforms is not None:
|
||||
img, target = self.transforms(img, target)
|
||||
|
||||
return img, target
|
||||
|
||||
def __len__(self):
|
||||
return len(self.images)
|
||||
|
||||
|
||||
# VOC
|
||||
@DATASETS.register()
|
||||
class PASCAL_VOC_Segmentation(_StandardSegmentationDataset):
|
||||
categories = [
|
||||
'Background',
|
||||
'Aeroplane', 'Bicycle', 'Bird', 'Boat',
|
||||
'Bottle', 'Bus', 'Car', 'Cat',
|
||||
'Chair', 'Cow', 'Diningtable', 'Dog',
|
||||
'Horse', 'Motorbike', 'Person', 'Pottedplant',
|
||||
'Sheep', 'Sofa', 'Train', 'Tvmonitor'
|
||||
]
|
||||
|
||||
colors = [
|
||||
[0, 0, 0],
|
||||
[128, 0, 0], [0, 128, 0], [128, 128, 0], [0, 0, 128],
|
||||
[128, 0, 128], [0, 128, 128], [128, 128, 128], [64, 0, 0],
|
||||
[192, 0, 0], [64, 128, 0], [192, 128, 0], [64, 0, 128],
|
||||
[192, 0, 128], [64, 128, 128], [192, 128, 128], [0, 64, 0],
|
||||
[128, 64, 0], [0, 192, 0], [128, 192, 0], [0, 64, 128],
|
||||
[255, 255, 255] # last color for ignore
|
||||
]
|
||||
|
||||
def init_dataset(self, root, image_set):
|
||||
image_dir = os.path.join(root, 'JPEGImages')
|
||||
mask_dir = os.path.join(root, 'SegmentationClassAug')
|
||||
splits_dir = os.path.join(root, 'ImageSets/Segmentation')
|
||||
split_f = os.path.join(splits_dir, image_set + '.txt')
|
||||
with open(split_f, "r") as f:
|
||||
file_names = [x.strip() for x in f.readlines()]
|
||||
|
||||
self.images = [os.path.join(image_dir, x + ".jpg") for x in file_names]
|
||||
self.masks = [os.path.join(mask_dir, x + self.mask_type) for x in file_names]
|
||||
|
||||
|
||||
# Cityscapes
|
||||
@DATASETS.register()
|
||||
class CityscapesSegmentation(_StandardSegmentationDataset):
|
||||
categories = [
|
||||
'road', 'sidewalk', 'building', 'wall',
|
||||
'fence', 'pole', 'traffic light', 'traffic sign',
|
||||
'vegetation', 'terrain', 'sky', 'person',
|
||||
'rider', 'car', 'truck', 'bus',
|
||||
'train', 'motorcycle', 'bicycle'
|
||||
]
|
||||
|
||||
colors = [
|
||||
[128, 64, 128], [244, 35, 232], [70, 70, 70], [102, 102, 156],
|
||||
[190, 153, 153], [153, 153, 153], [250, 170, 30], [220, 220, 0],
|
||||
[107, 142, 35], [152, 251, 152], [70, 130, 180], [220, 20, 60],
|
||||
[255, 0, 0], [0, 0, 142], [0, 0, 70], [0, 60, 100],
|
||||
[0, 80, 100], [0, 0, 230], [119, 11, 32],
|
||||
[0, 0, 0] # last color for ignore
|
||||
]
|
||||
|
||||
cities = [
|
||||
'aachen', 'bremen', 'darmstadt', 'erfurt', 'hanover',
|
||||
'krefeld', 'strasbourg', 'tubingen', 'weimar', 'bochum',
|
||||
'cologne', 'dusseldorf', 'hamburg', 'jena', 'monchengladbach',
|
||||
'stuttgart', 'ulm', 'zurich'
|
||||
]
|
||||
|
||||
def init_dataset(self, root, image_set):
|
||||
image_dir = os.path.join(root, 'leftImg8bit')
|
||||
mask_dir = os.path.join(root, 'gtFine')
|
||||
|
||||
if image_set == 'val':
|
||||
image_dir = os.path.join(image_dir, image_set)
|
||||
mask_dir = os.path.join(mask_dir, image_set)
|
||||
else:
|
||||
image_dir = os.path.join(image_dir, 'train')
|
||||
mask_dir = os.path.join(mask_dir, 'train')
|
||||
|
||||
# We first generate data lists before all this, so we can do this easier
|
||||
splits_dir = os.path.join(root, 'data_lists')
|
||||
split_f = os.path.join(splits_dir, image_set + '.txt')
|
||||
with open(split_f, "r") as f:
|
||||
file_names = [x.strip() for x in f.readlines()]
|
||||
|
||||
self.images = [os.path.join(image_dir, x + "_leftImg8bit.png") for x in file_names]
|
||||
self.masks = [os.path.join(mask_dir, x + "_gtFine_labelIds" + self.mask_type) for x in file_names]
|
||||
|
||||
|
||||
# GTAV
|
||||
@DATASETS.register()
|
||||
class GTAV_Segmentation(CityscapesSegmentation):
|
||||
cities = None
|
||||
|
||||
def init_dataset(self, root, image_set):
|
||||
image_dir = os.path.join(root, 'images')
|
||||
mask_dir = os.path.join(root, 'labels')
|
||||
|
||||
# We first generate data lists before all this, so we can do this easier
|
||||
splits_dir = os.path.join(root, 'data_lists')
|
||||
split_f = os.path.join(splits_dir, image_set + '.txt')
|
||||
with open(split_f, "r") as f:
|
||||
file_names = [x.strip() for x in f.readlines()]
|
||||
|
||||
self.images = [os.path.join(image_dir, x + ".png") for x in file_names]
|
||||
self.masks = [os.path.join(mask_dir, x + self.mask_type) for x in file_names]
|
||||
|
||||
|
||||
# SYNTHIA
|
||||
@DATASETS.register()
|
||||
class SYNTHIA_Segmentation(CityscapesSegmentation):
|
||||
cities = None
|
||||
|
||||
def init_dataset(self, root, image_set):
|
||||
image_dir = os.path.join(root, 'RGB', image_set)
|
||||
mask_dir = os.path.join(root, 'GT/LABELS_CONVERTED', image_set)
|
||||
|
||||
# We first generate data lists before all this, so we can do this easier
|
||||
splits_dir = os.path.join(root, 'data_lists')
|
||||
split_f = os.path.join(splits_dir, image_set + '.txt')
|
||||
with open(split_f, "r") as f:
|
||||
file_names = [x.strip() for x in f.readlines()]
|
||||
|
||||
self.images = [os.path.join(image_dir, x + ".png") for x in file_names]
|
||||
self.masks = [os.path.join(mask_dir, x + self.mask_type) for x in file_names]
|
||||
@@ -0,0 +1,66 @@
|
||||
import os
|
||||
try:
|
||||
import ujson as json
|
||||
except ImportError:
|
||||
import json
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
|
||||
from .utils import LaneKeypointDataset
|
||||
from .builder import DATASETS
|
||||
|
||||
|
||||
# TuSimple direct loading
|
||||
@DATASETS.register()
|
||||
class TuSimple(LaneKeypointDataset):
|
||||
colors = [
|
||||
[0, 0, 0], # background
|
||||
[0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 0],
|
||||
[0, 0, 0] # ignore
|
||||
]
|
||||
|
||||
def __init__(self, root, image_set, transforms=None, transform=None, target_transform=None,
|
||||
ppl=56, gap=10, start=160, padding_mask=False, is_process=True):
|
||||
super().__init__(root, transforms, transform, target_transform, ppl, gap, start, padding_mask, image_set,
|
||||
is_process)
|
||||
|
||||
self._check()
|
||||
|
||||
# Data list
|
||||
with open(os.path.join(root, 'lists', image_set + '.txt'), "r") as f:
|
||||
contents = [x.strip() for x in f.readlines()]
|
||||
|
||||
# Load image filenames and lanes
|
||||
if image_set == 'test' or image_set == 'val': # Test
|
||||
self.images = [os.path.join(root, 'clips', x + '.jpg') for x in contents]
|
||||
self.targets = [os.path.join('clips', x + '.jpg') for x in contents]
|
||||
else: # Train
|
||||
self.images = [os.path.join(root, 'clips', x[:x.find(' ')] + '.jpg') for x in contents]
|
||||
|
||||
# Load target lanes (small dataset, directly load all of them in the memory)
|
||||
print('Loading targets into memory...')
|
||||
target_files = [os.path.join(root, 'label_data_0313.json'),
|
||||
os.path.join(root, 'label_data_0601.json')]
|
||||
json_contents = self.concat_jsons(target_files)
|
||||
self.targets = []
|
||||
for i in tqdm(range(len(json_contents))):
|
||||
lines = json_contents[i]['lanes']
|
||||
h_samples = json_contents[i]['h_samples']
|
||||
temp = np.array([[[-2.0, self.start + j * self.gap] for j in range(self.ppl)]
|
||||
for _ in range(len(lines))], dtype=np.float32)
|
||||
for j in range(len(h_samples)):
|
||||
for k in range(len(lines)):
|
||||
temp[k][temp[k][:, 1] == h_samples[j]] = [float(lines[k][j]), h_samples[j]]
|
||||
self.targets.append(temp)
|
||||
|
||||
assert len(self.targets) == len(self.images)
|
||||
|
||||
@staticmethod
|
||||
def concat_jsons(filenames):
|
||||
# Concat tusimple lists in jsons (actually only each line is json)
|
||||
results = []
|
||||
for filename in filenames:
|
||||
with open(filename, 'r') as f:
|
||||
results += [json.loads(x.strip()) for x in f.readlines()]
|
||||
|
||||
return results
|
||||
@@ -0,0 +1,80 @@
|
||||
import os
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
|
||||
from .image_folder_lane_base import ImageFolderLaneBase
|
||||
from .tusimple import TuSimple
|
||||
from .builder import DATASETS
|
||||
|
||||
|
||||
def dummy_keypoint_process_fn(label):
|
||||
return label
|
||||
|
||||
|
||||
# Visualization version of CULane
|
||||
@DATASETS.register()
|
||||
class TuSimpleVis(ImageFolderLaneBase):
|
||||
def __init__(self, root_dataset, root_output, keypoint_json, image_set, transforms=None,
|
||||
keypoint_process_fn=None, use_gt=True):
|
||||
super().__init__(root_dataset, root_output, transforms, dummy_keypoint_process_fn)
|
||||
self.image_set = image_set
|
||||
|
||||
self._check()
|
||||
|
||||
# Data list
|
||||
with open(os.path.join(root_dataset, 'lists', image_set + '.txt'), "r") as f:
|
||||
contents = [x.strip() for x in f.readlines()]
|
||||
|
||||
# Load filenames
|
||||
if image_set == 'test' or image_set == 'val': # Test
|
||||
self.images = [os.path.join(root_dataset, 'clips', x + '.jpg') for x in contents]
|
||||
self.filenames = [os.path.join('clips', x + '.jpg') for x in contents]
|
||||
if use_gt:
|
||||
self.gt_keypoints = []
|
||||
target_files = [os.path.join(root_dataset, 'label_data_0531.json')]
|
||||
if image_set == 'test':
|
||||
target_files = [os.path.join(root_dataset, 'test_label.json')]
|
||||
json_contents = TuSimple.concat_jsons(target_files)
|
||||
self.gt_keypoints = self.preload_tusimple_labels(json_contents)
|
||||
if keypoint_json is not None:
|
||||
json_contents = TuSimple.concat_jsons([keypoint_json])
|
||||
self.keypoints = self.preload_tusimple_labels(json_contents)
|
||||
else: # Train
|
||||
self.images = [os.path.join(root_dataset, 'clips', x[:x.find(' ')] + '.jpg') for x in contents]
|
||||
self.filenames = [os.path.join('clips', x[:x.find(' ')] + '.jpg') for x in contents]
|
||||
if use_gt:
|
||||
self.gt_keypoints = []
|
||||
target_files = [os.path.join(root_dataset, 'label_data_0313.json'),
|
||||
os.path.join(root_dataset, 'label_data_0601.json')]
|
||||
json_contents = TuSimple.concat_jsons(target_files)
|
||||
self.gt_keypoints = self.preload_tusimple_labels(json_contents)
|
||||
if keypoint_json is not None:
|
||||
json_contents = TuSimple.concat_jsons([keypoint_json])
|
||||
self.keypoints = self.preload_tusimple_labels(json_contents)
|
||||
|
||||
self.make_sub_dirs()
|
||||
|
||||
assert len(self.images) == len(self.gt_keypoints)
|
||||
if self.keypoints is not None:
|
||||
assert len(self.images) == len(self.keypoints)
|
||||
|
||||
def _check(self):
|
||||
# Checks
|
||||
if self.image_set not in ['train', 'val', 'test']:
|
||||
raise ValueError
|
||||
assert self.output_dir != self.root, 'Avoid overwriting your dataset!'
|
||||
|
||||
@staticmethod
|
||||
def preload_tusimple_labels(json_contents):
|
||||
# Load a TuSimple label json's content
|
||||
print('Loading json annotation/prediction...')
|
||||
targets = []
|
||||
for i in tqdm(range(len(json_contents))):
|
||||
lines = json_contents[i]['lanes']
|
||||
h_samples = json_contents[i]['h_samples']
|
||||
temp = []
|
||||
for j in range(len(lines)):
|
||||
temp.append(np.array([[float(x), float(y)] for x, y in zip(lines[j], h_samples)]))
|
||||
targets.append(temp)
|
||||
|
||||
return targets
|
||||
@@ -0,0 +1,139 @@
|
||||
import os
|
||||
import collections.abc
|
||||
import torch
|
||||
import torchvision
|
||||
from PIL import Image
|
||||
from utils.transforms import functional_pil as f_pil
|
||||
# from torch._six import container_abcs, string_classes, int_classes
|
||||
from torch.utils.data._utils.collate import default_collate_err_msg_format, np_str_obj_array_pattern
|
||||
|
||||
|
||||
string_classes = (str, bytes)
|
||||
int_classes = int
|
||||
container_abcs = collections.abc
|
||||
|
||||
|
||||
def dict_collate_fn(batch):
|
||||
# To keep each image's label as separate dictionaries, default pytorch behaviour will stack each key
|
||||
# Only modified one line of the pytorch 1.6.0 default collate function
|
||||
|
||||
elem = batch[0]
|
||||
elem_type = type(elem)
|
||||
if isinstance(elem, torch.Tensor):
|
||||
out = None
|
||||
if torch.utils.data.get_worker_info() is not None:
|
||||
# If we're in a background process, concatenate directly into a
|
||||
# shared memory tensor to avoid an extra copy
|
||||
numel = sum([x.numel() for x in batch])
|
||||
storage = elem.storage()._new_shared(numel)
|
||||
out = elem.new(storage)
|
||||
return torch.stack(batch, 0, out=out)
|
||||
elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \
|
||||
and elem_type.__name__ != 'string_':
|
||||
elem = batch[0]
|
||||
if elem_type.__name__ == 'ndarray':
|
||||
# array of string classes and object
|
||||
if np_str_obj_array_pattern.search(elem.dtype.str) is not None:
|
||||
raise TypeError(default_collate_err_msg_format.format(elem.dtype))
|
||||
|
||||
return dict_collate_fn([torch.as_tensor(b) for b in batch])
|
||||
elif elem.shape == (): # scalars
|
||||
return torch.as_tensor(batch)
|
||||
elif isinstance(elem, float):
|
||||
return torch.tensor(batch, dtype=torch.float64)
|
||||
elif isinstance(elem, int_classes):
|
||||
return torch.tensor(batch)
|
||||
elif isinstance(elem, string_classes):
|
||||
return batch
|
||||
elif isinstance(elem, container_abcs.Mapping):
|
||||
return batch # !Only modified this line
|
||||
elif isinstance(elem, tuple) and hasattr(elem, '_fields'): # namedtuple
|
||||
return elem_type(*(dict_collate_fn(samples) for samples in zip(*batch)))
|
||||
elif isinstance(elem, container_abcs.Sequence):
|
||||
# check to make sure that the elements in batch have consistent size
|
||||
it = iter(batch)
|
||||
elem_size = len(next(it))
|
||||
if not all(len(elem) == elem_size for elem in it):
|
||||
raise RuntimeError('each element in list of batch should be of equal size')
|
||||
transposed = zip(*batch)
|
||||
return [dict_collate_fn(samples) for samples in transposed]
|
||||
|
||||
raise TypeError(default_collate_err_msg_format.format(elem_type))
|
||||
|
||||
|
||||
def generate_lane_label_dict(target):
|
||||
# target: {'keypoints': Tensor, L x N x 2, ...}
|
||||
# Although non-existent keypoints are marked as (-2, y), it is safer to check with > 0
|
||||
|
||||
# Drop invalid lanes (lanes with less than 2 keypoints are seen as invalid)
|
||||
target['lowers'] = torch.tensor([], dtype=target['keypoints'].dtype)
|
||||
target['uppers'] = torch.tensor([], dtype=target['keypoints'].dtype)
|
||||
target['labels'] = torch.tensor([], dtype=torch.int64)
|
||||
if target['keypoints'].numel() > 0:
|
||||
valid_lanes = (target['keypoints'][:, :, 0] > 0).sum(dim=-1) >= 2
|
||||
target['keypoints'] = target['keypoints'][valid_lanes]
|
||||
if target['keypoints'].numel() > 0: # Still has lanes
|
||||
# Append lowest & highest y coordinates (coordinates start at top-left corner), labels (all 1)
|
||||
# Looks better than giving MIN values
|
||||
target['lowers'] = torch.stack([l[l[:, 0] > 0][:, 1].max() for l in target['keypoints']])
|
||||
target['uppers'] = torch.stack([l[l[:, 0] > 0][:, 1].min() for l in target['keypoints']])
|
||||
target['labels'] = torch.ones(target['keypoints'].shape[0],
|
||||
device=target['keypoints'].device, dtype=torch.int64)
|
||||
|
||||
return target
|
||||
|
||||
|
||||
# Lanes as keypoints
|
||||
class LaneKeypointDataset(torchvision.datasets.VisionDataset):
|
||||
keypoint_color = [0, 0, 0]
|
||||
|
||||
def __init__(self, root, transforms, transform, target_transform,
|
||||
ppl, gap, start, padding_mask, image_set, is_process):
|
||||
super().__init__(root, transforms, transform, target_transform)
|
||||
self.ppl = ppl # Sampled points-per-lane
|
||||
self.gap = gap # y gap between sample points
|
||||
self.start = start # y coordinate to start annotation
|
||||
self.padding_mask = padding_mask # Padding mask for transformer
|
||||
self.process_points = image_set == 'train' # Add lowest & highest y coordinates, lane class labels
|
||||
self.is_process = is_process
|
||||
self.images = [] # placeholder
|
||||
self.targets = [] # placeholder
|
||||
self.image_set = image_set
|
||||
|
||||
def _check(self):
|
||||
# Checks
|
||||
if not os.path.exists('./output'):
|
||||
os.makedirs('./output')
|
||||
if self.image_set not in ['train', 'val', 'test']:
|
||||
raise ValueError
|
||||
|
||||
def __getitem__(self, index):
|
||||
# Return x (input image) & y (L lane with N coordinates (x, y) as np.array (L x N x 2))
|
||||
# Invalid coordinates are marked by (-2, y)
|
||||
# If just testing,
|
||||
# y is the filename to store prediction
|
||||
img = Image.open(self.images[index]).convert('RGB')
|
||||
if type(self.targets[index]) == str: # Load as paths
|
||||
target = self.targets[index]
|
||||
else: # Load as dict
|
||||
target = {'keypoints': self.targets[index]}
|
||||
|
||||
if (self.padding_mask or self.process_points) and type(target) == str:
|
||||
print('Testing does not require target padding_mask or process_point!')
|
||||
raise ValueError
|
||||
|
||||
# Add padding mask
|
||||
if self.padding_mask:
|
||||
target['padding_mask'] = Image.new("L", f_pil._get_image_size(img), 0)
|
||||
|
||||
# Transforms
|
||||
if self.transforms is not None:
|
||||
img, target = self.transforms(img, target)
|
||||
|
||||
if self.process_points and self.is_process:
|
||||
target = generate_lane_label_dict(target)
|
||||
|
||||
return img, target
|
||||
|
||||
def __len__(self):
|
||||
return len(self.images)
|
||||
@@ -0,0 +1,43 @@
|
||||
import torch
|
||||
from mmcv import VideoReader
|
||||
|
||||
from ..transforms import Compose
|
||||
from .builder import DATASETS
|
||||
|
||||
|
||||
# Load a video for inference
|
||||
@DATASETS.register()
|
||||
class VideoLoader(object):
|
||||
def __init__(self, filename, transforms=None, batch_size=1, *args, **kwargs):
|
||||
# Don't need ToTensor here
|
||||
self.transforms = Compose(transforms=[t for t in transforms.transforms if t.__class__.__name__ != 'ToTensor'])
|
||||
self.batch_size = batch_size
|
||||
self.video = VideoReader(filename)
|
||||
self.resolution = self.video.resolution
|
||||
self.fps = self.video.fps
|
||||
self.i = 0
|
||||
|
||||
def __next__(self):
|
||||
# Return transformed images / original images
|
||||
# Numpy can suffer a index OOB
|
||||
if self.i >= len(self):
|
||||
raise StopIteration
|
||||
|
||||
images_numpy = self.video[self.i * self.batch_size: (self.i + 1) * self.batch_size]
|
||||
images = torch.stack([torch.from_numpy(img) for img in images_numpy])
|
||||
images = images[..., [2, 1, 0]].permute(0, 3, 1, 2) / 255.0 # BHWC-rgb uint8 -> BCHW-rgb float
|
||||
original_images = images.clone()
|
||||
|
||||
# Transforms
|
||||
if self.transforms is not None:
|
||||
images = self.transforms(images)
|
||||
|
||||
self.i += 1
|
||||
|
||||
return images, original_images
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __len__(self):
|
||||
return len(self.video) // self.batch_size
|
||||
Reference in New Issue
Block a user