Skip to content

GitLab

Commit 41b18fd8 authored by zhe chen's avatar zhe chen
Browse files

Use pre-commit to reformat code 


Use pre-commit to reformat code
parent ff20ea39
Hide whitespace changes
Inline Side-by-side
Showing with 400 additions and 420 deletions
autonomous_driving/Online-HD-Map-Construction/src/datasets/pipelines/loading.py
View file @ 41b18fd8
......@@ -2,6 +2,7 @@ import mmcv
import numpy as np
from mmdet.datasets.builder import PIPELINES
@PIPELINES.register_module(force=True)
class LoadMultiViewImagesFromFiles(object):
"""Load multi channel images from a list of separate channel files.
......@@ -56,5 +57,5 @@ class LoadMultiViewImagesFromFiles(object):
def __repr__(self):
"""str: Return a string that describes the module."""
return f'{self.__class__.__name__} (to_float32={self.to_float32}, '\
f"color_type='{self.color_type}')"
return f'{self.__class__.__name__} (to_float32={self.to_float32}, ' \
f"color_type='{self.color_type}')"
autonomous_driving/Online-HD-Map-Construction/src/datasets/pipelines/poly_bbox.py
View file @ 41b18fd8
import numpy as np
from mmdet.datasets.builder import PIPELINES
from shapely.geometry import LineString
@PIPELINES.register_module(force=True)
class PolygonizeLocalMapBbox(object):
"""Pre-Processing used by vectormapnet model.
......@@ -18,7 +18,7 @@ class PolygonizeLocalMapBbox(object):
canvas_size=(200, 100),
coord_dim=2,
num_class=3,
threshold=6/200,
threshold=6 / 200,
):
self.canvas_size = np.array(canvas_size)
......@@ -47,7 +47,7 @@ class PolygonizeLocalMapBbox(object):
polyline_weight = np.ones_like(polyline).reshape(-1)
polyline_weight = np.pad(
polyline_weight, ((0, 1),), constant_values=1.)
polyline_weight = polyline_weight/polyline_weight.sum()
polyline_weight = polyline_weight / polyline_weight.sum()
# flatten and quantilized
fpolyline = quantize_verts(
......@@ -58,7 +58,7 @@ class PolygonizeLocalMapBbox(object):
# reindex starting from 1, and add a zero stopping token(EOS),
fpolyline = \
np.pad(fpolyline + self.coord_dim_start_idx, ((0, 1),),
constant_values=0)
constant_values=0)
fpolyline_msk = np.ones(fpolyline.shape, dtype=np.bool)
polyline_masks.append(fpolyline_msk)
......@@ -98,11 +98,11 @@ class PolygonizeLocalMapBbox(object):
qkp_msks = np.stack(qkp_masks)
# format det
kps = np.stack(kps, axis=0).astype(np.float32)*self.canvas_size
kps = np.stack(kps, axis=0).astype(np.float32) * self.canvas_size
kp_labels = np.array(kp_labels)
# restrict the boundary
kps[..., 0] = np.clip(kps[..., 0], 0.1, self.canvas_size[0]-0.1)
kps[..., 1] = np.clip(kps[..., 1], 0.1, self.canvas_size[1]-0.1)
kps[..., 0] = np.clip(kps[..., 0], 0.1, self.canvas_size[0] - 0.1)
kps[..., 1] = np.clip(kps[..., 1], 0.1, self.canvas_size[1] - 0.1)
# nbox, boxsize(4)*coord_dim(2)
kps = kps.reshape(kps.shape[0], -1)
......@@ -114,7 +114,7 @@ class PolygonizeLocalMapBbox(object):
'''
Process vertices.
'''
vectors = input_dict['vectors']
n_lines = 0
......@@ -157,10 +157,9 @@ class PolygonizeLocalMapBbox(object):
def evaluate_line(polyline):
edge = np.linalg.norm(polyline[1:] - polyline[:-1], axis=-1)
start_end_weight = edge[(0, -1), ].copy()
start_end_weight = edge[(0, -1),].copy()
mid_weight = (edge[:-1] + edge[1:]) * .5
pts_weight = np.concatenate(
......@@ -172,16 +171,16 @@ def evaluate_line(polyline):
pts_weight /= denominator
# add weights for stop index
pts_weight = np.repeat(pts_weight, 2)/2
pts_weight = np.repeat(pts_weight, 2) / 2
pts_weight = np.pad(pts_weight, ((0, 1)),
constant_values=1/(len(polyline)*2))
constant_values=1 / (len(polyline) * 2))
return pts_weight
def quantize_verts(verts, canvas_size, coord_dim):
"""Convert vertices from its original range ([-1,1]) to discrete values in [0, n_bits**2 - 1].
Args:
verts (array): vertices coordinates, shape (seqlen, coords_dim)
canvas_size (tuple): bev feature size
......@@ -196,7 +195,7 @@ def quantize_verts(verts, canvas_size, coord_dim):
range_quantize = np.array(canvas_size) - 1 # (0-199) = 200
verts_ratio = (verts[:, :coord_dim] - min_range) / (
max_range - min_range)
max_range - min_range)
verts_quantize = verts_ratio * range_quantize[:coord_dim]
return verts_quantize.astype('int32')
......@@ -204,11 +203,11 @@ def quantize_verts(verts, canvas_size, coord_dim):
def get_bbox(polyline, threshold):
"""Convert vertices from its original range ([-1,1]) to discrete values in [0, n_bits**2 - 1].
Args:
polyline (array): point coordinates, shape (seqlen, 2)
threshold (float): threshold for minimum bbox size
Returns:
bbox (array): bounding box in xyxy format, shape (2, 2)
"""
......@@ -216,14 +215,14 @@ def get_bbox(polyline, threshold):
polyline = LineString(polyline)
bbox = polyline.bounds
minx, miny, maxx, maxy = bbox
W, H = maxx-minx, maxy-miny
W, H = maxx - minx, maxy - miny
if W < threshold or H < threshold:
remain = max((threshold - min(W, H))/2, eps)
remain = max((threshold - min(W, H)) / 2, eps)
bbox = polyline.buffer(remain).envelope.bounds
minx, miny, maxx, maxy = bbox
bbox_np = np.array([[minx, miny], [maxx, maxy]])
bbox_np = np.clip(bbox_np, 0., 1.)
return bbox_np
\ No newline at end of file
return bbox_np
autonomous_driving/Online-HD-Map-Construction/src/datasets/pipelines/transform.py
View file @ 41b18fd8
import numpy as np
import mmcv
import numpy as np
from mmdet.datasets.builder import PIPELINES
......@@ -82,26 +81,26 @@ class PadMultiViewImages(object):
if self.change_intrinsics:
post_intrinsics, post_ego2imgs = [], []
for img, oshape, cam_intrinsic, ego2img in zip(results['img'], \
original_shape, results['cam_intrinsics'], results['ego2img']):
original_shape, results['cam_intrinsics'],
results['ego2img']):
scaleW = img.shape[1] / oshape[1]
scaleH = img.shape[0] / oshape[0]
rot_resize_matrix = np.array([
[scaleW, 0, 0, 0],
[0, scaleH, 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 1]])
rot_resize_matrix = np.array([
[scaleW, 0, 0, 0],
[0, scaleH, 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 1]])
post_intrinsic = rot_resize_matrix[:3, :3] @ cam_intrinsic
post_ego2img = rot_resize_matrix @ ego2img
post_intrinsics.append(post_intrinsic)
post_ego2imgs.append(post_ego2img)
results.update({
'cam_intrinsics': post_intrinsics,
'ego2img': post_ego2imgs,
})
results['img_shape'] = [img.shape for img in padded_img]
results['img_fixed_size'] = self.size
results['img_size_divisor'] = self.size_divisor
......@@ -135,16 +134,17 @@ class ResizeMultiViewImages(object):
size (tuple, optional): resize target size, (h, w).
change_intrinsics (bool): whether to update intrinsics.
"""
def __init__(self, size, change_intrinsics=True):
self.size = size
self.change_intrinsics = change_intrinsics
def __call__(self, results:dict):
def __call__(self, results: dict):
new_imgs, post_intrinsics, post_ego2imgs = [], [], []
for img, cam_intrinsic, ego2img in zip(results['img'], \
results['cam_intrinsics'], results['ego2img']):
for img, cam_intrinsic, ego2img in zip(results['img'], \
results['cam_intrinsics'], results['ego2img']):
tmp, scaleW, scaleH = mmcv.imresize(img,
# NOTE: mmcv.imresize expect (w, h) shape
(self.size[1], self.size[0]),
......@@ -152,10 +152,10 @@ class ResizeMultiViewImages(object):
new_imgs.append(tmp)
rot_resize_matrix = np.array([
[scaleW, 0, 0, 0],
[0, scaleH, 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 1]])
[scaleW, 0, 0, 0],
[0, scaleH, 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 1]])
post_intrinsic = rot_resize_matrix[:3, :3] @ cam_intrinsic
post_ego2img = rot_resize_matrix @ ego2img
post_intrinsics.append(post_intrinsic)
......@@ -170,10 +170,10 @@ class ResizeMultiViewImages(object):
})
return results
def __repr__(self):
repr_str = self.__class__.__name__
repr_str += f'(size={self.size}, '
repr_str += f'change_intrinsics={self.change_intrinsics})'
return repr_str
\ No newline at end of file
return repr_str
autonomous_driving/Online-HD-Map-Construction/src/datasets/pipelines/vectorize.py
View file @ 41b18fd8
from typing import Dict, List, Tuple, Union
import numpy as np
from mmdet.datasets.builder import PIPELINES
from shapely.geometry import LineString
from numpy.typing import NDArray
from typing import List, Tuple, Union, Dict
from shapely.geometry import LineString
@PIPELINES.register_module(force=True)
class VectorizeMap(object):
......@@ -20,14 +22,14 @@ class VectorizeMap(object):
sample_dist (float): interpolate distance. Set to -1 to ignore.
"""
def __init__(self,
roi_size: Union[Tuple, List],
def __init__(self,
roi_size: Union[Tuple, List],
normalize: bool,
coords_dim: int,
simplify: bool=False,
sample_num: int=-1,
sample_dist: float=-1,
):
simplify: bool = False,
sample_num: int = -1,
sample_dist: float = -1,
):
self.coords_dim = coords_dim
self.sample_num = sample_num
self.sample_dist = sample_dist
......@@ -45,46 +47,46 @@ class VectorizeMap(object):
def interp_fixed_num(self, line: LineString) -> NDArray:
''' Interpolate a line to fixed number of points.
Args:
line (LineString): line
Returns:
points (array): interpolated points, shape (N, 2)
'''
distances = np.linspace(0, line.length, self.sample_num)
sampled_points = np.array([list(line.interpolate(distance).coords)
for distance in distances]).squeeze()
sampled_points = np.array([list(line.interpolate(distance).coords)
for distance in distances]).squeeze()
return sampled_points
def interp_fixed_dist(self, line: LineString) -> NDArray:
''' Interpolate a line at fixed interval.
Args:
line (LineString): line
Returns:
points (array): interpolated points, shape (N, 2)
'''
distances = list(np.arange(self.sample_dist, line.length, self.sample_dist))
# make sure to sample at least two points when sample_dist > line.length
distances = [0,] + distances + [line.length,]
distances = [0, ] + distances + [line.length, ]
sampled_points = np.array([list(line.interpolate(distance).coords)
for distance in distances]).squeeze()
for distance in distances]).squeeze()
return sampled_points
def get_vectorized_lines(self, map_geoms: Dict) -> Dict:
''' Vectorize map elements. Iterate over the input dict and apply the
''' Vectorize map elements. Iterate over the input dict and apply the
specified sample funcion.
Args:
line (LineString): line
Returns:
vectors (array): dict of vectorized map elements.
'''
......@@ -110,22 +112,22 @@ class VectorizeMap(object):
elif geom.geom_type == 'Polygon':
# polygon objects will not be vectorized
continue
else:
raise ValueError('map geoms must be either LineString or Polygon!')
return vectors
def normalize_line(self, line: NDArray) -> NDArray:
''' Convert points to range (0, 1).
Args:
line (LineString): line
Returns:
normalized (array): normalized points.
'''
origin = -np.array([self.roi_size[0]/2, self.roi_size[1]/2])
origin = -np.array([self.roi_size[0] / 2, self.roi_size[1] / 2])
line[:, :2] = line[:, :2] - origin
......@@ -134,7 +136,7 @@ class VectorizeMap(object):
line[:, :2] = line[:, :2] / (self.roi_size + eps)
return line
def __call__(self, input_dict):
map_geoms = input_dict['map_geoms']
......@@ -145,9 +147,9 @@ class VectorizeMap(object):
repr_str = self.__class__.__name__
repr_str += f'(simplify={self.simplify}, '
repr_str += f'sample_num={self.sample_num}), '
repr_str += f'sample_dist={self.sample_dist}), '
repr_str += f'sample_dist={self.sample_dist}), '
repr_str += f'roi_size={self.roi_size})'
repr_str += f'normalize={self.normalize})'
repr_str += f'coords_dim={self.coords_dim})'
return repr_str
\ No newline at end of file
return repr_str
autonomous_driving/Online-HD-Map-Construction/src/models/__init__.py
View file @ 41b18fd8
from .backbones import *
from .heads import *
from .losses import *
from .mapers import *
from .transformer_utils import *
from .assigner import *
autonomous_driving/Online-HD-Map-Construction/src/models/assigner/__init__.py
View file @ 41b18fd8
from .assigner import HungarianLinesAssigner
from .match_cost import MapQueriesCost, BBoxLogitsCost, DynamicLinesCost, IoUCostC, BBoxCostC, LinesCost, LinesFixNumChamferCost, ClsSigmoidCost
autonomous_driving/Online-HD-Map-Construction/src/models/assigner/assigner.py
View file @ 41b18fd8
import torch
from mmdet.core.bbox.assigners import AssignResult, BaseAssigner
from mmdet.core.bbox.builder import BBOX_ASSIGNERS
from mmdet.core.bbox.assigners import AssignResult
from mmdet.core.bbox.assigners import BaseAssigner
from mmdet.core.bbox.match_costs import build_match_cost
try:
......@@ -36,8 +34,8 @@ class HungarianLinesAssigner(BaseAssigner):
type='MapQueriesCost',
cls_cost=dict(type='ClassificationCost', weight=1.),
reg_cost=dict(type='LinesCost', weight=1.0),
),
pc_range=None,
),
pc_range=None,
**kwargs):
self.pc_range = pc_range
......@@ -110,7 +108,8 @@ class HungarianLinesAssigner(BaseAssigner):
matched_row_inds, matched_col_inds = linear_sum_assignment(cost)
except:
print('cost max{}, min{}'.format(cost.max(), cost.min()))
import ipdb; ipdb.set_trace()
import ipdb
ipdb.set_trace()
matched_row_inds = torch.from_numpy(matched_row_inds).to(
preds['lines'].device)
matched_col_inds = torch.from_numpy(matched_col_inds).to(
......@@ -123,4 +122,4 @@ class HungarianLinesAssigner(BaseAssigner):
assigned_gt_inds[matched_row_inds] = matched_col_inds + 1
assigned_labels[matched_row_inds] = gts['labels'][matched_col_inds]
return AssignResult(
num_gts, assigned_gt_inds, None, labels=assigned_labels)
\ No newline at end of file
num_gts, assigned_gt_inds, None, labels=assigned_labels)
autonomous_driving/Online-HD-Map-Construction/src/models/assigner/match_cost.py
View file @ 41b18fd8
import torch
from mmdet.core.bbox.match_costs.builder import MATCH_COST
from mmdet.core.bbox.match_costs import build_match_cost
from mmdet.core.bbox.iou_calculators import bbox_overlaps
from mmdet.core.bbox.match_costs import build_match_cost
from mmdet.core.bbox.match_costs.builder import MATCH_COST
from mmdet.core.bbox.transforms import bbox_cxcywh_to_xyxy
......@@ -83,7 +82,7 @@ class LinesFixNumChamferCost(object):
num_gts, num_bboxes = gt_lines.size(0), lines_pred.size(0)
dist_mat = lines_pred.new_full((num_bboxes, num_gts),
1.0,)
1.0, )
for i in range(num_bboxes):
for j in range(num_gts):
......@@ -212,6 +211,7 @@ class IoUCostC:
iou_cost = -overlaps
return iou_cost * self.weight
@MATCH_COST.register_module()
class DynamicLinesCost(object):
"""LinesL1Cost.
......@@ -273,7 +273,7 @@ class DynamicLinesCost(object):
m1 = m1.unsqueeze(1).sigmoid() > 0.5
m2 = m2.unsqueeze(0)
valid_points_mask = (m1 + m2)/2.
valid_points_mask = (m1 + m2) / 2.
average_factor_mask = valid_points_mask.sum(-1) > 0
average_factor = average_factor_mask.masked_fill(
......@@ -360,8 +360,7 @@ class MapQueriesCost(object):
# Iou
if self.iou_cost is not None:
iou_cost = self.iou_cost(preds['lines'],gts['lines'])
iou_cost = self.iou_cost(preds['lines'], gts['lines'])
cost += iou_cost
return cost
autonomous_driving/Online-HD-Map-Construction/src/models/augmentation/sythesis_det.py
View file @ 41b18fd8
......@@ -5,13 +5,13 @@ import torch.nn.functional as F
class NoiseSythesis(nn.Module):
def __init__(self,
p, scale=0.01, shift_scale=(8,5),
scaling_size=(0.1,0.1), canvas_size=(200, 100),
bbox_type='sce',
poly_coord_dim=2,
bbox_coord_dim=2,
quantify=True):
def __init__(self,
p, scale=0.01, shift_scale=(8, 5),
scaling_size=(0.1, 0.1), canvas_size=(200, 100),
bbox_type='sce',
poly_coord_dim=2,
bbox_coord_dim=2,
quantify=True):
super(NoiseSythesis, self).__init__()
self.p = p
......@@ -37,7 +37,7 @@ class NoiseSythesis(nn.Module):
dtype = bbox.dtype
B = bbox.shape[0]
noise = (torch.rand(B, device=device)*2-1)[:,None,None] # [-1,1]
noise = (torch.rand(B, device=device) * 2 - 1)[:, None, None] # [-1,1]
scale = self.scaling_size.to(device)
scale = (noise * scale) + 1
......@@ -45,7 +45,7 @@ class NoiseSythesis(nn.Module):
# recenterization
coffset = scaled_bbox.mean(-2) - bbox.float().mean(-2)
scaled_bbox = scaled_bbox - coffset[:,None]
scaled_bbox = scaled_bbox - coffset[:, None]
return scaled_bbox.round().type(dtype)
......@@ -60,13 +60,13 @@ class NoiseSythesis(nn.Module):
scale = (bbox.max(1)[0] - bbox.min(1)[0]) * 0.1
scale = torch.where(scale < shift_scale, scale, shift_scale)
noise = (torch.rand(batch_size, 2, device=device)*2-1) # [-1,1]
noise = (torch.rand(batch_size, 2, device=device) * 2 - 1) # [-1,1]
offset = (noise * scale).round().type(bbox.dtype)
shifted_bbox = bbox + offset[:, None]
return shifted_bbox
def gaussian_noise_bbox(self, bbox):
dtype = bbox.dtype
......@@ -80,23 +80,23 @@ class NoiseSythesis(nn.Module):
noisy_bbox = noisy_bbox.round().type(dtype)
# prevent out of bound case
for i in range(self.bbox_coord_dim):
noisy_bbox[...,i] =\
torch.clamp(noisy_bbox[...,0],1,self.canvas_size[i])
noisy_bbox[..., i] = \
torch.clamp(noisy_bbox[..., 0], 1, self.canvas_size[i])
else:
noisy_bbox = noisy_bbox.type(torch.float)
return noisy_bbox
def gaussian_noise_poly(self, polyline, polyline_mask):
device = polyline.device
batchsize = polyline.shape[0]
scale = self.canvas_size * self.scale
polyline = F.pad(polyline,(0,self.poly_coord_dim-1))
polyline = polyline.view(batchsize,-1, self.poly_coord_dim)
mask = F.pad(polyline_mask[:,1:],(0,self.poly_coord_dim))
polyline = F.pad(polyline, (0, self.poly_coord_dim - 1))
polyline = polyline.view(batchsize, -1, self.poly_coord_dim)
mask = F.pad(polyline_mask[:, 1:], (0, self.poly_coord_dim))
noisy_polyline = torch.normal(polyline.type(torch.float), scale)
if self.quantify:
......@@ -104,14 +104,14 @@ class NoiseSythesis(nn.Module):
# prevent out of bound case
for i in range(self.poly_coord_dim):
noisy_polyline[...,i] =\
torch.clamp(noisy_polyline[...,i],0,self.canvas_size[i])
noisy_polyline[..., i] = \
torch.clamp(noisy_polyline[..., i], 0, self.canvas_size[i])
else:
noisy_polyline = noisy_polyline.type(torch.float)
noisy_polyline = noisy_polyline.view(batchsize,-1) * mask
noisy_polyline = noisy_polyline[:,:-(self.poly_coord_dim-1)]
noisy_polyline = noisy_polyline.view(batchsize, -1) * mask
noisy_polyline = noisy_polyline[:, :-(self.poly_coord_dim - 1)]
return noisy_polyline
......@@ -125,11 +125,11 @@ class NoiseSythesis(nn.Module):
bbox = t(bbox)
# prevent out of bound case
bbox[...,0] =\
torch.clamp(bbox[...,0],0,self.canvas_size[0])
bbox[...,1] =\
torch.clamp(bbox[...,1],0,self.canvas_size[1])
bbox[..., 0] = \
torch.clamp(bbox[..., 0], 0, self.canvas_size[0])
bbox[..., 1] = \
torch.clamp(bbox[..., 1], 0, self.canvas_size[1])
return bbox
......@@ -143,8 +143,8 @@ class NoiseSythesis(nn.Module):
bbox = self.gaussian_noise_bbox(bbox)
fbbox_aug = bbox.view(seq_len, -1)
aug_mask = torch.rand(fbbox.shape,device=fbbox.device)
fbbox = torch.where(aug_mask<self.p, fbbox_aug, fbbox)
aug_mask = torch.rand(fbbox.shape, device=fbbox.device)
fbbox = torch.where(aug_mask < self.p, fbbox_aug, fbbox)
elif self.bbox_type == 'rxyxy':
fbbox = self.rbbox_aug(batch)
elif self.bbox_type == 'convex_hull':
......@@ -154,18 +154,18 @@ class NoiseSythesis(nn.Module):
polyline = batch['polylines']
polyline_mask = batch['polyline_masks']
polyline_aug = self.gaussian_noise_poly(polyline, polyline_mask)
aug_mask = torch.rand(polyline.shape,device=polyline.device)
polyline = torch.where(aug_mask<self.p, polyline_aug, polyline)
aug_mask = torch.rand(polyline.shape, device=polyline.device)
polyline = torch.where(aug_mask < self.p, polyline_aug, polyline)
return polyline, fbbox
def rbbox_aug(self, batch):
return None
def convex_hull_aug(self,batch):
def convex_hull_aug(self, batch):
return None
def __call__(self, batch, simple_aug=False):
......@@ -183,5 +183,4 @@ class NoiseSythesis(nn.Module):
aug_bbox_flat = aug_bbox.view(seq_len, -1)
return aug_bbox_flat
autonomous_driving/Online-HD-Map-Construction/src/models/backbones/__init__.py
View file @ 41b18fd8
from .ipm_backbone import IPMEncoder
__all__ = [
'IPMEncoder'
'IPMEncoder'
]
autonomous_driving/Online-HD-Map-Construction/src/models/backbones/internimage.py
View file @ 41b18fd8
......@@ -4,17 +4,19 @@
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
from collections import OrderedDict
import torch
import torch.nn as nn
from collections import OrderedDict
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint
from timm.models.layers import trunc_normal_, DropPath
from mmcv.runner import _load_checkpoint
from mmcv.cnn import constant_init, trunc_normal_init
from mmcv.runner import _load_checkpoint
from mmdet.models.builder import BACKBONES
from mmseg.utils import get_root_logger
from ops_dcnv3 import modules as opsm
import torch.nn.functional as F
from mmdet.models.builder import BACKBONES
from timm.models.layers import DropPath, trunc_normal_
class to_channels_first(nn.Module):
......@@ -84,7 +86,7 @@ class CrossAttention(nn.Module):
attn_head_dim (int, optional): Dimension of attention head.
out_dim (int, optional): Dimension of output.
"""
def __init__(self,
dim,
num_heads=8,
......@@ -176,7 +178,7 @@ class AttentiveBlock(nn.Module):
attn_head_dim (int, optional): Dimension of attention head. Default: None.
out_dim (int, optional): Dimension of output. Default: None.
"""
def __init__(self,
dim,
num_heads,
......@@ -185,7 +187,7 @@ class AttentiveBlock(nn.Module):
drop=0.,
attn_drop=0.,
drop_path=0.,
norm_layer="LN",
norm_layer='LN',
attn_head_dim=None,
out_dim=None):
super().__init__()
......@@ -361,9 +363,9 @@ class InternImageLayer(nn.Module):
layer_scale=None,
offset_scale=1.0,
with_cp=False,
dw_kernel_size=None, # for InternImage-H/G
res_post_norm=False, # for InternImage-H/G
center_feature_scale=False): # for InternImage-H/G
dw_kernel_size=None, # for InternImage-H/G
res_post_norm=False, # for InternImage-H/G
center_feature_scale=False): # for InternImage-H/G
super().__init__()
self.channels = channels
self.groups = groups
......@@ -382,8 +384,8 @@ class InternImageLayer(nn.Module):
offset_scale=offset_scale,
act_layer=act_layer,
norm_layer=norm_layer,
dw_kernel_size=dw_kernel_size, # for InternImage-H/G
center_feature_scale=center_feature_scale) # for InternImage-H/G
dw_kernel_size=dw_kernel_size, # for InternImage-H/G
center_feature_scale=center_feature_scale) # for InternImage-H/G
self.drop_path = DropPath(drop_path) if drop_path > 0. \
else nn.Identity()
self.norm2 = build_norm_layer(channels, 'LN')
......@@ -409,7 +411,7 @@ class InternImageLayer(nn.Module):
if self.post_norm:
x = x + self.drop_path(self.norm1(self.dcn(x)))
x = x + self.drop_path(self.norm2(self.mlp(x)))
elif self.res_post_norm: # for InternImage-H/G
elif self.res_post_norm: # for InternImage-H/G
x = x + self.drop_path(self.res_post_norm1(self.dcn(self.norm1(x))))
x = x + self.drop_path(self.res_post_norm2(self.mlp(self.norm2(x))))
else:
......@@ -464,10 +466,10 @@ class InternImageBlock(nn.Module):
offset_scale=1.0,
layer_scale=None,
with_cp=False,
dw_kernel_size=None, # for InternImage-H/G
post_norm_block_ids=None, # for InternImage-H/G
res_post_norm=False, # for InternImage-H/G
center_feature_scale=False): # for InternImage-H/G
dw_kernel_size=None, # for InternImage-H/G
post_norm_block_ids=None, # for InternImage-H/G
res_post_norm=False, # for InternImage-H/G
center_feature_scale=False): # for InternImage-H/G
super().__init__()
self.channels = channels
self.depth = depth
......@@ -489,15 +491,15 @@ class InternImageBlock(nn.Module):
layer_scale=layer_scale,
offset_scale=offset_scale,
with_cp=with_cp,
dw_kernel_size=dw_kernel_size, # for InternImage-H/G
res_post_norm=res_post_norm, # for InternImage-H/G
center_feature_scale=center_feature_scale # for InternImage-H/G
dw_kernel_size=dw_kernel_size, # for InternImage-H/G
res_post_norm=res_post_norm, # for InternImage-H/G
center_feature_scale=center_feature_scale # for InternImage-H/G
) for i in range(depth)
])
if not self.post_norm or center_feature_scale:
self.norm = build_norm_layer(channels, 'LN')
self.post_norm_block_ids = post_norm_block_ids
if post_norm_block_ids is not None: # for InternImage-H/G
if post_norm_block_ids is not None: # for InternImage-H/G
self.post_norms = nn.ModuleList(
[build_norm_layer(channels, 'LN', eps=1e-6) for _ in post_norm_block_ids]
)
......@@ -509,7 +511,7 @@ class InternImageBlock(nn.Module):
x = blk(x)
if (self.post_norm_block_ids is not None) and (i in self.post_norm_block_ids):
index = self.post_norm_block_ids.index(i)
x = self.post_norms[index](x) # for InternImage-H/G
x = self.post_norms[index](x) # for InternImage-H/G
if not self.post_norm or self.center_feature_scale:
x = self.norm(x)
if return_wo_downsample:
......@@ -575,7 +577,7 @@ class InternImage(nn.Module):
self.num_levels = len(depths)
self.depths = depths
self.channels = channels
self.num_features = int(channels * 2**(self.num_levels - 1))
self.num_features = int(channels * 2 ** (self.num_levels - 1))
self.post_norm = post_norm
self.mlp_ratio = mlp_ratio
self.init_cfg = init_cfg
......@@ -607,10 +609,10 @@ class InternImage(nn.Module):
self.levels = nn.ModuleList()
for i in range(self.num_levels):
post_norm_block_ids = level2_post_norm_block_ids if level2_post_norm and (
i == 2) else None # for InternImage-H/G
i == 2) else None # for InternImage-H/G
level = InternImageBlock(
core_op=getattr(opsm, core_op),
channels=int(channels * 2**i),
channels=int(channels * 2 ** i),
depth=depths[i],
groups=groups[i],
mlp_ratio=self.mlp_ratio,
......@@ -624,9 +626,9 @@ class InternImage(nn.Module):
offset_scale=offset_scale,
with_cp=with_cp,
dw_kernel_size=dw_kernel_size, # for InternImage-H/G
post_norm_block_ids=post_norm_block_ids, # for InternImage-H/G
res_post_norm=res_post_norm, # for InternImage-H/G
center_feature_scale=center_feature_scale # for InternImage-H/G
post_norm_block_ids=post_norm_block_ids, # for InternImage-H/G
res_post_norm=res_post_norm, # for InternImage-H/G
center_feature_scale=center_feature_scale # for InternImage-H/G
)
self.levels.append(level)
......@@ -697,4 +699,4 @@ class InternImage(nn.Module):
x, x_ = level(x, return_wo_downsample=True)
if level_idx in self.out_indices:
seq_out.append(x_.permute(0, 3, 1, 2).contiguous())
return seq_out
\ No newline at end of file
return seq_out
autonomous_driving/Online-HD-Map-Construction/src/models/backbones/ipm_backbone.py
View file @ 41b18fd8
import copy
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from mmdet3d.models.builder import BACKBONES
from mmdet.models import build_backbone, build_neck
class UpsampleBlock(nn.Module):
def __init__(self, ins, outs):
super(UpsampleBlock, self).__init__()
......@@ -17,7 +18,6 @@ class UpsampleBlock(nn.Module):
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
x = self.conv(x)
x = self.relu(self.gn(x))
x = self.upsample2x(x)
......@@ -26,7 +26,7 @@ class UpsampleBlock(nn.Module):
def upsample2x(self, x):
_, _, h, w = x.shape
x = F.interpolate(x, size=(h*2, w*2),
x = F.interpolate(x, size=(h * 2, w * 2),
mode='bilinear', align_corners=True)
return x
......@@ -54,7 +54,7 @@ class Upsample(nn.Module):
continue
tmp = [copy.deepcopy(input_conv), ]
tmp += [copy.deepcopy(inter_conv) for i in range(layer_num-1)]
tmp += [copy.deepcopy(inter_conv) for i in range(layer_num - 1)]
fscale.append(nn.Sequential(*tmp))
self.fscale = nn.ModuleList(fscale)
......@@ -117,21 +117,21 @@ class IPMEncoder(nn.Module):
if self.use_lidar:
self.pp = PointPillarEncoder(lidar_dim, xbound, ybound, zbound)
self.outconvs =\
nn.Conv2d((self.upsample.out_channels+3)*len(heights), out_channels//2,
kernel_size=3, stride=1, padding=1) # same
self.outconvs = \
nn.Conv2d((self.upsample.out_channels + 3) * len(heights), out_channels // 2,
kernel_size=3, stride=1, padding=1) # same
if self.use_image:
_out_channels = out_channels//2
_out_channels = out_channels // 2
else:
_out_channels = out_channels
self.outconvs_lidar =\
nn.Conv2d(lidar_dim, _out_channels,
kernel_size=3, stride=1, padding=1) # same
self.outconvs_lidar = \
nn.Conv2d(lidar_dim, _out_channels,
kernel_size=3, stride=1, padding=1) # same
else:
self.outconvs =\
nn.Conv2d((self.upsample.out_channels+3)*len(heights), out_channels,
kernel_size=3, stride=1, padding=1) # same
self.outconvs = \
nn.Conv2d((self.upsample.out_channels + 3) * len(heights), out_channels,
kernel_size=3, stride=1, padding=1) # same
self.init_weights(pretrained=pretrained)
......@@ -139,11 +139,10 @@ class IPMEncoder(nn.Module):
bev_planes = [construct_plane_grid(
xbound, ybound, h) for h in self.heights]
self.register_buffer('bev_planes', torch.stack(
bev_planes),) # nlvl,bH,bW,2
bev_planes), ) # nlvl,bH,bW,2
self.masked_embeds = nn.Embedding(len(heights), out_channels)
def init_weights(self, pretrained=None):
"""Initialize model weights."""
......@@ -154,12 +153,12 @@ class IPMEncoder(nn.Module):
for p in self.outconvs.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
if self.use_lidar:
for p in self.outconvs_lidar.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
for p in self.pp.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
......@@ -169,7 +168,7 @@ class IPMEncoder(nn.Module):
Extract image feaftures and sum up into one pic
Args:
imgs: B, n_cam, C, iH, iW
Returns:
Returns:
img_feat: B * n_cam, C, H, W
'''
......@@ -188,12 +187,12 @@ class IPMEncoder(nn.Module):
def forward(self, imgs, img_metas, *args, points=None, **kwargs):
'''
Args:
Args:
imgs: torch.Tensor of shape [B, N, 3, H, W]
N: number of cams
img_metas:
img_metas:
# N=6, ['CAM_FRONT', 'CAM_FRONT_RIGHT', 'CAM_FRONT_LEFT', 'CAM_BACK', 'CAM_BACK_LEFT', 'CAM_BACK_RIGHT']
ego2cam: [B, N, 4, 4]
ego2cam: [B, N, 4, 4]
cam_intrinsics: [B, N, 3, 3]
cam2ego_rotations: [B, N, 3, 3]
cam2ego_translations: [B, N, 3]
......@@ -225,7 +224,7 @@ class IPMEncoder(nn.Module):
if self.use_lidar:
lidar_feat = self.get_lidar_feature(points)
if self.use_image:
bev_feat = torch.cat([bev_feat,lidar_feat],dim=1)
bev_feat = torch.cat([bev_feat, lidar_feat], dim=1)
else:
bev_feat = lidar_feat
......@@ -233,7 +232,7 @@ class IPMEncoder(nn.Module):
def ipm(self, cam_feat, ego2cam, img_shape):
'''
inverse project
inverse project
Args:
cam_feat: B*ncam, C, cH, cW
img_shape: tuple(H, W)
......@@ -250,7 +249,7 @@ class IPMEncoder(nn.Module):
# bev_grid_pos: B*ncam, nlvl*bH*bW, 2
bev_grid_pos, bev_cam_mask = get_campos(bev_grid, ego2cam, img_shape)
# B*cam, nlvl*bH, bW, 2
bev_grid_pos = bev_grid_pos.unflatten(-2, (nlvl*bH, bW))
bev_grid_pos = bev_grid_pos.unflatten(-2, (nlvl * bH, bW))
# project feat from 2D to bev plane
projected_feature = F.grid_sample(
......@@ -262,11 +261,11 @@ class IPMEncoder(nn.Module):
# eliminate the ncam
# The bev feature is the sum of the 6 cameras
bev_feat_mask = bev_feat_mask.unsqueeze(2)
projected_feature = (projected_feature*bev_feat_mask).sum(1)
projected_feature = (projected_feature * bev_feat_mask).sum(1)
num_feat = bev_feat_mask.sum(1)
projected_feature = projected_feature / \
num_feat.masked_fill(num_feat == 0, 1)
num_feat.masked_fill(num_feat == 0, 1)
# concatenate a position information
# projected_feature: B, bH, bW, nlvl, C+3
......@@ -287,7 +286,7 @@ class IPMEncoder(nn.Module):
# bev_grid = bev_grid.permute(0, 3, 1, 2)
# lidar_feature = torch.cat(
# (lidar_feature, bev_grid), dim=1)
lidar_feature = self.outconvs_lidar(lidar_feature)
return lidar_feature
......@@ -321,7 +320,7 @@ def construct_plane_grid(xbound, ybound, height: float, dtype=torch.float32):
def get_campos(reference_points, ego2cam, img_shape):
'''
Find the each refence point's corresponding pixel in each camera
Args:
Args:
reference_points: [B, num_query, 3]
ego2cam: (B, num_cam, 4, 4)
Outs:
......@@ -351,7 +350,7 @@ def get_campos(reference_points, ego2cam, img_shape):
eps = 1e-9
mask = (reference_points_cam[..., 2:3] > eps)
reference_points_cam =\
reference_points_cam = \
reference_points_cam[..., 0:2] / \
reference_points_cam[..., 2:3] + eps
......@@ -362,13 +361,13 @@ def get_campos(reference_points, ego2cam, img_shape):
reference_points_cam = (reference_points_cam - 0.5) * 2
mask = (mask & (reference_points_cam[..., 0:1] > -1.0)
& (reference_points_cam[..., 0:1] < 1.0)
& (reference_points_cam[..., 1:2] > -1.0)
& (reference_points_cam[..., 1:2] < 1.0))
& (reference_points_cam[..., 0:1] < 1.0)
& (reference_points_cam[..., 1:2] > -1.0)
& (reference_points_cam[..., 1:2] < 1.0))
# (B, num_cam, num_query)
mask = mask.view(B, num_cam, num_query)
reference_points_cam = reference_points_cam.view(B*num_cam, num_query, 2)
reference_points_cam = reference_points_cam.view(B * num_cam, num_query, 2)
return reference_points_cam, mask
......
autonomous_driving/Online-HD-Map-Construction/src/models/heads/__init__.py
View file @ 41b18fd8
from .base_map_head import BaseMapHead
from .dg_head import DGHead
from .map_element_detector import MapElementDetector
from .polyline_generator import PolylineGenerator
\ No newline at end of file
autonomous_driving/Online-HD-Map-Construction/src/models/heads/base_map_head.py
View file @ 41b18fd8
......@@ -3,7 +3,6 @@ from abc import ABCMeta, abstractmethod
import torch.nn as nn
from mmcv.runner import auto_fp16
from mmcv.utils import print_log
from mmdet.utils import get_root_logger
......@@ -24,10 +23,10 @@ class BaseMapHead(nn.Module, metaclass=ABCMeta):
logger = get_root_logger()
print_log(f'load model from: {pretrained}', logger=logger)
@auto_fp16(apply_to=('img', ))
@auto_fp16(apply_to=('img',))
def forward(self, *args, **kwargs):
pass
@abstractmethod
def loss(self, pred, gt):
'''
......@@ -42,7 +41,7 @@ class BaseMapHead(nn.Module, metaclass=ABCMeta):
)
'''
return
@abstractmethod
def post_process(self, pred):
'''
......
autonomous_driving/Online-HD-Map-Construction/src/models/heads/detgen_utils/causal_trans.py
View file @ 41b18fd8
# the causal layer is credited by the https://github.com/alexmt-scale/causal-transformer-decoder
# we made some change to stick with the polygen.
import torch
import torch.nn as nn
from typing import Optional
from torch import Tensor
import torch
import torch.nn as nn
from mmcv.cnn.bricks.registry import ATTENTION
from mmcv.utils import build_from_cfg
from torch import Tensor
def build_attention(cfg, default_args=None):
......@@ -29,14 +29,14 @@ class CausalTransformerDecoder(nn.TransformerDecoder):
"""
def forward(
self,
tgt: Tensor,
memory: Optional[Tensor] = None,
cache: Optional[Tensor] = None,
memory_mask: Optional[Tensor] = None,
tgt_key_padding_mask: Optional[Tensor] = None,
memory_key_padding_mask: Optional[Tensor] = None,
causal_mask: Optional[Tensor] = None,
self,
tgt: Tensor,
memory: Optional[Tensor] = None,
cache: Optional[Tensor] = None,
memory_mask: Optional[Tensor] = None,
tgt_key_padding_mask: Optional[Tensor] = None,
memory_key_padding_mask: Optional[Tensor] = None,
causal_mask: Optional[Tensor] = None,
) -> Tensor:
"""
Args:
......@@ -58,7 +58,7 @@ class CausalTransformerDecoder(nn.TransformerDecoder):
if self.training:
if cache is not None:
raise ValueError(
"cache parameter should be None in training mode")
'cache parameter should be None in training mode')
for mod in self.layers:
output = mod(
output,
......@@ -132,7 +132,7 @@ class CausalTransformerDecoderLayer(nn.TransformerDecoderLayer):
"""
Args:
see CausalTransformerDecoder
query is not None model will perform query stream
query is not None model will perform query stream
Returns:
Tensor:
If training: embedding of the whole layer: seq_len x bsz x hidden_dim
......@@ -140,23 +140,23 @@ class CausalTransformerDecoderLayer(nn.TransformerDecoderLayer):
"""
if not self.norm_first:
raise ValueError(
"norm_first parameter should be True!")
'norm_first parameter should be True!')
if self.training:
# the official Pytorch implementation
x = tgt
if query is not None:
x = query
x = x + self.res_weight1 * \
self._sa_block(self.norm1(x), self.norm1(tgt), causal_mask,
tgt_key_padding_mask)
tgt_key_padding_mask)
if memory is not None:
x = x + self.res_weight2 * \
self._mha_block(self.norm2(x), memory,
memory_mask, memory_key_padding_mask)
x = x + self.res_weight3*self._ff_block(self.norm3(x))
x = x + self.res_weight3 * self._ff_block(self.norm3(x))
return x
# This part is adapted from the official Pytorch implementation
......@@ -169,14 +169,14 @@ class CausalTransformerDecoderLayer(nn.TransformerDecoderLayer):
if only_last:
x = x[-1:]
if causal_mask is not None:
attn_mask = causal_mask
attn_mask = causal_mask
if only_last:
attn_mask = attn_mask[-1:] # XXX
attn_mask = attn_mask[-1:] # XXX
else:
attn_mask = None
# efficient self attention
x = x + self.res_weight1 * \
self._sa_block(self.norm1(x), self.norm1(tgt), attn_mask,
......@@ -189,7 +189,7 @@ class CausalTransformerDecoderLayer(nn.TransformerDecoderLayer):
memory_mask, memory_key_padding_mask)
# final feed-forward network
x = x + self.res_weight3*self._ff_block(self.norm3(x))
x = x + self.res_weight3 * self._ff_block(self.norm3(x))
return x
......@@ -235,7 +235,8 @@ class PolygenTransformerEncoderLayer(nn.TransformerEncoderLayer):
self.norm_first = norm_first
def forward(self, src: Tensor, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
def forward(self, src: Tensor, src_mask: Optional[Tensor] = None,
src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
r"""Pass the input through the encoder layer.
Args:
src: the sequence to the encoder layer (required).
......@@ -249,13 +250,13 @@ class PolygenTransformerEncoderLayer(nn.TransformerEncoderLayer):
x = src
if self.norm_first:
x = x + self.res_weight1*self._sa_block(self.norm1(x), src_mask,
src_key_padding_mask)
x = x + self.res_weight2*self._ff_block(self.norm2(x))
x = x + self.res_weight1 * self._sa_block(self.norm1(x), src_mask,
src_key_padding_mask)
x = x + self.res_weight2 * self._ff_block(self.norm2(x))
else:
x = self.norm1(
x + self.res_weight1*self._sa_block(x, src_mask, src_key_padding_mask))
x = self.norm2(x + self.res_weight2*self._ff_block(x))
x + self.res_weight1 * self._sa_block(x, src_mask, src_key_padding_mask))
x = self.norm2(x + self.res_weight2 * self._ff_block(x))
return x
......@@ -274,12 +275,12 @@ class PolygenTransformerEncoderLayer(nn.TransformerEncoderLayer):
return self.dropout2(x)
def generate_square_subsequent_mask(sz: int, device: str = "cpu") -> torch.Tensor:
def generate_square_subsequent_mask(sz: int, device: str = 'cpu') -> torch.Tensor:
""" Generate the attention mask for causal decoding """
mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
mask = (
mask.float()
.masked_fill(mask == 0, float("-inf"))
.masked_fill(mask == 1, float(0.0))
.masked_fill(mask == 0, float('-inf'))
.masked_fill(mask == 1, float(0.0))
).to(device=device)
return mask
\ No newline at end of file
return mask
autonomous_driving/Online-HD-Map-Construction/src/models/heads/detgen_utils/utils.py
View file @ 41b18fd8
......@@ -2,18 +2,20 @@ import torch
import torch.nn.functional as F
from torch import Tensor
def generate_square_subsequent_mask(sz: int, condition_len: int = 1, bool_out=False, device: str = "cpu") -> torch.Tensor:
def generate_square_subsequent_mask(sz: int, condition_len: int = 1, bool_out=False,
device: str = 'cpu') -> torch.Tensor:
""" Generate the attention mask for causal decoding """
mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
if condition_len > 1:
mask[:condition_len,:condition_len] = 1
mask[:condition_len, :condition_len] = 1
if not bool_out:
mask = (
mask.float()
.masked_fill(mask == 0, float("-inf"))
.masked_fill(mask == 1, float(0.0)))
.masked_fill(mask == 0, float('-inf'))
.masked_fill(mask == 1, float(0.0)))
return mask.to(device=device)
......@@ -39,10 +41,10 @@ def quantize_verts(
"""
min_range = -1
max_range = 1
range_quantize = canvas_size-1
range_quantize = canvas_size - 1
verts_ratio = (verts - min_range) / (
max_range - min_range)
max_range - min_range)
verts_quantize = verts_ratio * range_quantize
return verts_quantize.type(torch.int32)
......@@ -56,7 +58,7 @@ def top_k_logits(logits, k):
values, _ = torch.topk(logits, k=k)
k_largest = torch.min(values)
logits = torch.where(logits < k_largest,
torch.ones_like(logits)*-1e9, logits)
torch.ones_like(logits) * -1e9, logits)
return logits
......
autonomous_driving/Online-HD-Map-Construction/src/models/heads/detr_bbox.py
View file @ 41b18fd8
import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
from mmcv.cnn import Conv2d, Linear
from mmcv.cnn import Linear
from mmcv.runner import force_fp32
from torch.distributions.categorical import Categorical
from mmdet.core import multi_apply, reduce_mean
from mmdet.models import HEADS
from torch.distributions.categorical import Categorical
from .detr_head import DETRMapFixedNumHead
@HEADS.register_module(force=True)
class DETRBboxHead(DETRMapFixedNumHead):
def __init__(self, *args, canvas_size=(400, 200), discrete_output=True, separate_detect=True,
mode='xyxy', bbox_size=None, coord_dim=2, kp_coord_dim=2,
**kwargs):
def __init__(self, *args, canvas_size=(400, 200), discrete_output=True, separate_detect=True,
mode='xyxy', bbox_size=None, coord_dim=2, kp_coord_dim=2,
**kwargs):
self.canvas_size = canvas_size # hard code
self.separate_detect = separate_detect
self.discrete_output = discrete_output
self.bbox_size = 3 if mode=='sce' else 2
self.bbox_size = 3 if mode == 'sce' else 2
if bbox_size is not None:
self.bbox_size = bbox_size
self.coord_dim = coord_dim # for xyz
......@@ -31,7 +32,7 @@ class DETRBboxHead(DETRMapFixedNumHead):
del self.canvas_size
self.register_buffer('canvas_size', torch.tensor(canvas_size))
self._init_embedding()
def _init_embedding(self):
# for bbox parameter xstart, ystart, xend, yend
......@@ -42,12 +43,12 @@ class DETRBboxHead(DETRMapFixedNumHead):
self.img_coord_embed = nn.Linear(2, self.embed_dims)
def _init_branch(self,):
def _init_branch(self, ):
"""Initialize classification branch and regression branch of head."""
# add sigmoid or not
if self.separate_detect:
if self.cls_out_channels == self.num_classes+1:
if self.cls_out_channels == self.num_classes + 1:
self.cls_out_channels = 2
else:
self.cls_out_channels = 1
......@@ -62,10 +63,10 @@ class DETRBboxHead(DETRMapFixedNumHead):
if self.discrete_output:
reg_branch.append(nn.Linear(
self.embed_dims, max(self.canvas_size), bias=True,))
self.embed_dims, max(self.canvas_size), bias=True, ))
else:
reg_branch.append(nn.Linear(
self.embed_dims, self.bbox_size*self.coord_dim, bias=True,))
self.embed_dims, self.bbox_size * self.coord_dim, bias=True, ))
reg_branch = nn.Sequential(*reg_branch)
......@@ -133,12 +134,12 @@ class DETRBboxHead(DETRMapFixedNumHead):
[nb_dec, bs, num_query, num_points, 2].
'''
(global_context_embedding, sequential_context_embeddings) =\
(global_context_embedding, sequential_context_embeddings) = \
self._prepare_context(batch, context)
if self.separate_detect:
query_embedding = self.query_embedding.weight[None] + \
global_context_embedding[:, None]
global_context_embedding[:, None]
else:
B = sequential_context_embeddings.shape[0]
query_embedding = self.query_embedding.weight[None].repeat(B, 1, 1)
......@@ -166,18 +167,18 @@ class DETRBboxHead(DETRMapFixedNumHead):
pos = []
for i in range(4):
pos_embeds = self.bbox_embedding.weight[i]
_pos = self.pre_branches['reg'](query_feat+pos_embeds)
_pos = self.pre_branches['reg'](query_feat + pos_embeds)
pos.append(_pos)
# # y mask
# _vert_mask = torch.arange(logits.shape[-1], device=logits.device)
# vertices_mask_y = (_vert_mask < self.canvas_size[1]+1)
# logits[:,1::2] = logits[:,1::2]*vertices_mask_y - ~vertices_mask_y*1e9
logits = torch.stack(pos, dim=-2)/1.
logits = torch.stack(pos, dim=-2) / 1.
lines = Categorical(logits=logits)
else:
lines = self.pre_branches['reg'](query_feat).sigmoid()
lines = lines.unflatten(-1, (self.bbox_size, self.coord_dim))*self.canvas_size
lines = lines.unflatten(-1, (self.bbox_size, self.coord_dim)) * self.canvas_size
lines = lines.flatten(-2)
return dict(
......@@ -220,7 +221,7 @@ class DETRBboxHead(DETRMapFixedNumHead):
num_pred_lines = len(lines_pred)
# assigner and sampler
assign_result = self.assigner.assign(preds=dict(lines=lines_pred, scores=score_pred,),
assign_result = self.assigner.assign(preds=dict(lines=lines_pred, scores=score_pred, ),
gts=dict(lines=gt_lines,
labels=gt_labels, ),
gt_bboxes_ignore=gt_bboxes_ignore)
......@@ -232,10 +233,10 @@ class DETRBboxHead(DETRMapFixedNumHead):
# label targets 0: foreground, 1: background
if self.separate_detect:
labels = gt_lines.new_full((num_pred_lines, ), 1, dtype=torch.long)
labels = gt_lines.new_full((num_pred_lines,), 1, dtype=torch.long)
else:
labels = gt_lines.new_full(
(num_pred_lines, ), self.num_classes, dtype=torch.long)
(num_pred_lines,), self.num_classes, dtype=torch.long)
labels[pos_inds] = gt_labels[sampling_result.pos_assigned_gt_inds]
label_weights = gt_lines.new_ones(num_pred_lines)
......@@ -308,11 +309,11 @@ class DETRBboxHead(DETRMapFixedNumHead):
(labels_list, label_weights_list,
lines_targets_list, lines_weights_list,
pos_inds_list, neg_inds_list,pos_gt_inds_list) = multi_apply(
self._get_target_single,
preds['scores'], lines_pred,
class_label, bbox,
gt_bboxes_ignore=gt_bboxes_ignore_list)
pos_inds_list, neg_inds_list, pos_gt_inds_list) = multi_apply(
self._get_target_single,
preds['scores'], lines_pred,
class_label, bbox,
gt_bboxes_ignore=gt_bboxes_ignore_list)
num_total_pos = sum((inds.numel() for inds in pos_inds_list))
num_total_neg = sum((inds.numel() for inds in neg_inds_list))
......@@ -351,7 +352,7 @@ class DETRBboxHead(DETRMapFixedNumHead):
"""
# Get target for each sample
new_gts, num_total_pos, num_total_neg, pos_inds_list, pos_gt_inds_list =\
new_gts, num_total_pos, num_total_neg, pos_inds_list, pos_gt_inds_list = \
self.get_targets(preds, gts, gt_bboxes_ignore_list)
# Batched all data
......@@ -360,7 +361,7 @@ class DETRBboxHead(DETRMapFixedNumHead):
# construct weighted avg_factor to match with the official DETR repo
cls_avg_factor = num_total_pos * 1.0 + \
num_total_neg * self.bg_cls_weight
num_total_neg * self.bg_cls_weight
if self.sync_cls_avg_factor:
cls_avg_factor = reduce_mean(
preds['scores'].new_tensor([cls_avg_factor]))
......@@ -386,7 +387,7 @@ class DETRBboxHead(DETRMapFixedNumHead):
# position NLL loss
if self.discrete_output:
loss_reg = -(preds['lines'].log_prob(new_gts['bboxs']) *
new_gts['bboxs_weights']).sum()/(num_total_pos)
new_gts['bboxs_weights']).sum() / (num_total_pos)
else:
loss_reg = self.reg_loss(
preds['lines'], new_gts['bboxs'], new_gts['bboxs_weights'], avg_factor=num_total_pos)
......@@ -408,9 +409,9 @@ class DETRBboxHead(DETRMapFixedNumHead):
pos_msk = label == 0
neg_msk = ~pos_msk
loss_cls = -(p.log()*pos_msk + (1-p).log()*neg_msk)
loss_cls = -(p.log() * pos_msk + (1 - p).log() * neg_msk)
loss_cls = (loss_cls * weights).sum()/cls_avg_factor
loss_cls = (loss_cls * weights).sum() / cls_avg_factor
return loss_cls
......@@ -465,7 +466,7 @@ class DETRBboxHead(DETRMapFixedNumHead):
result_dict['bbox'].append(det_preds)
result_dict['scores'].append(scores)
result_dict['labels'].append(det_labels)
result_dict['lines_bs_idx'].extend([i]*nline)
result_dict['lines_bs_idx'].extend([i] * nline)
# for down stream polyline
_bboxs = torch.cat(result_dict['bbox'], dim=0)
......@@ -481,4 +482,4 @@ class DETRBboxHead(DETRMapFixedNumHead):
def assign_bev(feat, idx):
return feat[idx]
\ No newline at end of file
return feat[idx]
autonomous_driving/Online-HD-Map-Construction/src/models/heads/detr_head.py
View file @ 41b18fd8
import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
import copy
from mmdet.models import HEADS
from mmcv.cnn import Conv2d
from mmcv.cnn import Linear, build_activation_layer, bias_init_with_prob
from mmcv.cnn import (Conv2d, Linear, bias_init_with_prob,
build_activation_layer)
from mmcv.cnn.bricks.transformer import build_positional_encoding
from mmdet.models.utils import build_transformer
from mmcv.runner import force_fp32
from mmdet.core import (multi_apply, build_assigner, build_sampler,
reduce_mean)
from mmdet.models.utils.transformer import inverse_sigmoid
from mmdet.models import build_loss
from mmdet.core import build_assigner, build_sampler, multi_apply, reduce_mean
from mmdet.models import HEADS, build_loss
from mmdet.models.utils import build_transformer
from .base_map_head import BaseMapHead
......@@ -60,14 +57,14 @@ class DETRMapFixedNumHead(BaseMapHead):
if loss_cls['use_sigmoid']:
self.cls_out_channels = num_classes
else:
self.cls_out_channels = num_classes+1
self.cls_out_channels = num_classes + 1
self.iterative = iterative
self.num_reg_fcs = num_reg_fcs
if patch_size is not None:
self.register_buffer('patch_size', torch.tensor(
(patch_size[1], patch_size[0])),)
(patch_size[1], patch_size[0])), )
self._build_transformer(transformer, positional_encoding)
......@@ -104,7 +101,7 @@ class DETRMapFixedNumHead(BaseMapHead):
self.transformer = build_transformer(transformer)
self.embed_dims = self.transformer.embed_dims
def _init_branch(self,):
def _init_branch(self, ):
"""Initialize classification branch and regression branch of head."""
fc_cls = Linear(self.embed_dims, self.cls_out_channels)
......@@ -114,8 +111,9 @@ class DETRMapFixedNumHead(BaseMapHead):
reg_branch.append(Linear(self.embed_dims, self.embed_dims))
reg_branch.append(nn.LayerNorm(self.embed_dims))
reg_branch.append(nn.ReLU())
reg_branch.append(Linear(self.embed_dims, self.num_points*2))
reg_branch.append(Linear(self.embed_dims, self.num_points * 2))
reg_branch = nn.Sequential(*reg_branch)
# add sigmoid or not
def _get_clones(module, N):
......@@ -185,7 +183,6 @@ class DETRMapFixedNumHead(BaseMapHead):
outputs = []
for i, query_feat in enumerate(outs_dec):
ocls = self.pre_branches['cls'](query_feat)
oreg = self.pre_branches['reg'](query_feat)
oreg = oreg.unflatten(dim=2, sizes=(self.num_points, 2))
......@@ -235,7 +232,7 @@ class DETRMapFixedNumHead(BaseMapHead):
num_pred_lines = lines_pred.size(0)
# assigner and sampler
assign_result = self.assigner.assign(preds=dict(lines=lines_pred, scores=score_pred,),
assign_result = self.assigner.assign(preds=dict(lines=lines_pred, scores=score_pred, ),
gts=dict(lines=gt_lines,
labels=gt_labels, ),
gt_bboxes_ignore=gt_bboxes_ignore)
......@@ -245,7 +242,7 @@ class DETRMapFixedNumHead(BaseMapHead):
neg_inds = sampling_result.neg_inds
# label targets
labels = gt_lines.new_full((num_pred_lines, ),
labels = gt_lines.new_full((num_pred_lines,),
self.num_classes,
dtype=torch.long)
labels[pos_inds] = gt_labels[sampling_result.pos_assigned_gt_inds]
......@@ -297,10 +294,10 @@ class DETRMapFixedNumHead(BaseMapHead):
(labels_list, label_weights_list,
lines_targets_list, lines_weights_list,
pos_inds_list, neg_inds_list) = multi_apply(
self._get_target_single,
preds['scores'], preds['lines'],
gts['lines'], gts['labels'],
gt_bboxes_ignore=gt_bboxes_ignore_list)
self._get_target_single,
preds['scores'], preds['lines'],
gts['lines'], gts['labels'],
gt_bboxes_ignore=gt_bboxes_ignore_list)
num_total_pos = sum((inds.numel() for inds in pos_inds_list))
num_total_neg = sum((inds.numel() for inds in neg_inds_list))
......@@ -319,7 +316,7 @@ class DETRMapFixedNumHead(BaseMapHead):
gts: dict,
gt_bboxes_ignore_list=None,
reduction='none'):
"""
"""
Loss function for outputs from a single decoder layer of a single
feature level.
Args:
......@@ -327,7 +324,7 @@ class DETRMapFixedNumHead(BaseMapHead):
for all images. Shape [bs, num_query, cls_out_channels].
lines_preds (Tensor):
shape [bs, num_query, num_points, 2].
gt_lines_list (list[Tensor]):
gt_lines_list (list[Tensor]):
with shape (num_gts, num_points, 2)
gt_labels_list (list[Tensor]): Ground truth class indices for each
image with shape (num_gts, ).
......@@ -339,7 +336,7 @@ class DETRMapFixedNumHead(BaseMapHead):
"""
# get target for each sample
new_gts, num_total_pos, num_total_neg, pos_inds_list =\
new_gts, num_total_pos, num_total_neg, pos_inds_list = \
self.get_targets(preds, gts, gt_bboxes_ignore_list)
# batched all data
......@@ -348,7 +345,7 @@ class DETRMapFixedNumHead(BaseMapHead):
# construct weighted avg_factor to match with the official DETR repo
cls_avg_factor = num_total_pos * 1.0 + \
num_total_neg * self.bg_cls_weight
num_total_neg * self.bg_cls_weight
if self.sync_cls_avg_factor:
cls_avg_factor = reduce_mean(
preds['scores'].new_tensor([cls_avg_factor]))
......@@ -368,7 +365,8 @@ class DETRMapFixedNumHead(BaseMapHead):
lines_preds = preds['lines'].reshape(-1, self.num_points, 2)
if reduction == 'none': # For performance analysis
loss_reg = self.reg_loss(
lines_preds, new_gts['lines_targets'], new_gts['lines_weights'], reduction_override=reduction, avg_factor=num_total_pos)
lines_preds, new_gts['lines_targets'], new_gts['lines_weights'], reduction_override=reduction,
avg_factor=num_total_pos)
else:
loss_reg = self.reg_loss(
lines_preds, new_gts['lines_targets'], new_gts['lines_weights'], avg_factor=num_total_pos)
......
autonomous_driving/Online-HD-Map-Construction/src/models/heads/dg_head.py
View file @ 41b18fd8
import copy
import numpy as np
import torch
import torch.nn as nn
from mmcv.cnn import Linear, bias_init_with_prob, build_activation_layer
from mmcv.cnn.bricks.transformer import build_positional_encoding
from mmcv.runner import force_fp32
from mmdet.models import HEADS, build_head, build_loss
from mmdet.models import HEADS, build_head
from mmdet.models.utils import build_transformer
from mmdet.models.utils.transformer import inverse_sigmoid
from .base_map_head import BaseMapHead
import numpy as np
from ..augmentation.sythesis_det import NoiseSythesis
from .base_map_head import BaseMapHead
@HEADS.register_module(force=True)
class DGHead(BaseMapHead):
......@@ -46,16 +41,16 @@ class DGHead(BaseMapHead):
self.augmentation = None
if augmentation:
augmentation_kwargs.update({'canvas_size':gen_net_cfg.canvas_size})
augmentation_kwargs.update({'canvas_size': gen_net_cfg.canvas_size})
self.augmentation = NoiseSythesis(**augmentation_kwargs)
self.joint_training = joint_training
def forward(self, batch, img_metas=None, **kwargs):
'''
Args:
Returns:
outs (Dict):
outs (Dict):
'''
if self.training:
......@@ -68,8 +63,8 @@ class DGHead(BaseMapHead):
bbox_dict = self.det_net(context=context)
outs = dict(
bbox=bbox_dict,
)
bbox=bbox_dict,
)
losses_dict, det_match_idxs, det_match_gt_idxs = \
self.loss_det(batch, outs)
......@@ -77,12 +72,12 @@ class DGHead(BaseMapHead):
if only_det: return outs, losses_dict
if self.augmentation is not None:
polylines, bbox_flat =\
self.augmentation(batch['gen'],simple_aug=True)
polylines, bbox_flat = \
self.augmentation(batch['gen'], simple_aug=True)
if bbox_flat is None:
bbox_flat = batch['gen']['bbox_flat']
gen_input = dict(
lines_bs_idx=batch['gen']['lines_bs_idx'],
lines_cls=batch['gen']['lines_cls'],
......@@ -104,32 +99,32 @@ class DGHead(BaseMapHead):
pred_bbox = bbox_dict[-1]['bboxs'].detach()
else:
raise NotImplementedError
# changed to original gt order.
# changed to original gt order.
det_match_idx = det_match_idxs[-1]
det_match_gt_idx = det_match_gt_idxs[-1]
_bboxs = []
for i, (match_idx, bbox) in enumerate(zip(det_match_idx,pred_bbox)):
_bboxs.append(bbox[match_idx])
_bboxs[-1] = _bboxs[-1][torch.argsort(det_match_gt_idx[i])]
for i, (match_idx, bbox) in enumerate(zip(det_match_idx, pred_bbox)):
_bboxs.append(bbox[match_idx])
_bboxs[-1] = _bboxs[-1][torch.argsort(det_match_gt_idx[i])]
_bboxs = torch.cat(_bboxs, dim=0)
# quantize the data
_bboxs = \
torch.round(_bboxs).type(torch.int32)
# gen_input['bbox_flat'] = _bboxs
remain_idx = torch.randperm(_bboxs.shape[0])[:int(_bboxs.shape[0]*0.2)]
remain_idx = torch.randperm(_bboxs.shape[0])[:int(_bboxs.shape[0] * 0.2)]
# for data efficient
for k in gen_input.keys():
if k == 'bbox_flat':
gen_input[k] = torch.cat((_bboxs,gen_input[k][remain_idx]),dim=0)
gen_input[k] = torch.cat((_bboxs, gen_input[k][remain_idx]), dim=0)
else:
gen_input[k] = torch.cat((gen_input[k],gen_input[k][remain_idx]),dim=0)
if isinstance(context['bev_embeddings'],tuple):
gen_input[k] = torch.cat((gen_input[k], gen_input[k][remain_idx]), dim=0)
if isinstance(context['bev_embeddings'], tuple):
context['bev_embeddings'] = context['bev_embeddings'][0]
poly_dict = self.gen_net(gen_input, context=context)
......@@ -141,17 +136,17 @@ class DGHead(BaseMapHead):
if self.joint_training:
for k in batch['gen'].keys():
batch['gen'][k] = \
torch.cat((batch['gen'][k],batch['gen'][k][remain_idx]),dim=0)
torch.cat((batch['gen'][k], batch['gen'][k][remain_idx]), dim=0)
gen_losses_dict = \
self.loss_gen(batch, outs)
losses_dict.update(gen_losses_dict)
losses_dict.update(gen_losses_dict)
return outs, losses_dict
def loss_det(self, gt: dict, pred: dict):
loss_dict = {}
# det
......@@ -159,8 +154,8 @@ class DGHead(BaseMapHead):
self.det_net.loss(gt['det'], pred['bbox'])
for k, v in det_loss_dict.items():
loss_dict['det_'+k] = v
loss_dict['det_' + k] = v
return loss_dict, det_match_idx, det_match_gt_idx
def loss_gen(self, gt: dict, pred: dict):
......@@ -171,34 +166,34 @@ class DGHead(BaseMapHead):
gen_loss_dict = self.gen_net.loss(gt['gen'], pred['polylines'])
for k, v in gen_loss_dict.items():
loss_dict['gen_'+k] = v
loss_dict['gen_' + k] = v
return loss_dict
def loss(self, gt: dict, pred: dict):
pass
@torch.no_grad()
def inference(self, batch: dict={}, context: dict={}, gt_condition=False, **kwargs):
def inference(self, batch: dict = {}, context: dict = {}, gt_condition=False, **kwargs):
'''
num_samples_batch: number of sample per batch (batch size)
'''
outs = {}
bbox_dict = self.det_net(context=context)
bbox_dict = self.det_net.post_process(bbox_dict)
outs.update(bbox_dict)
if len(outs['lines_bs_idx']) == 0:
return None
if isinstance(context['bev_embeddings'],tuple):
if isinstance(context['bev_embeddings'], tuple):
context['bev_embeddings'] = context['bev_embeddings'][0]
poly_dict = self.gen_net(outs,
context=context,
# max_sample_length=self.max_num_vertices,
# max_sample_length=self.max_num_vertices,
max_sample_length=64,
top_p=self.top_p_gen_model,
gt_condition=gt_condition)
......@@ -206,7 +201,7 @@ class DGHead(BaseMapHead):
return outs
def post_process(self, preds: dict, tokens, gts:dict=None, **kwargs):
def post_process(self, preds: dict, tokens, gts: dict = None, **kwargs):
'''
Args:
XXX
......@@ -215,8 +210,8 @@ class DGHead(BaseMapHead):
'''
range_size = self.gen_net.canvas_size.cpu().numpy()
coord_dim = self.gen_net.coord_dim
gen_net_name = self.gen_net.name if hasattr(self.gen_net,'name') else 'gen'
gen_net_name = self.gen_net.name if hasattr(self.gen_net, 'name') else 'gen'
ret_list = []
for batch_idx in range(len(tokens)):
......@@ -227,8 +222,8 @@ class DGHead(BaseMapHead):
det_gt = None
if gts is not None:
det_gt, rec_groundtruth = pack_groundtruth(
batch_idx,gts,tokens,range_size,gen_net_name,coord_dim=coord_dim)
batch_idx, gts, tokens, range_size, gen_net_name, coord_dim=coord_dim)
bbox_res = {
# 'bboxes': preds['bbox'][batch_idx].detach().cpu().numpy(),
# 'det_gt': det_gt,
......@@ -238,7 +233,6 @@ class DGHead(BaseMapHead):
}
ret_dict_single.update(bbox_res)
# for gen results.
batch2seq = np.nonzero(
preds['lines_bs_idx'].cpu().numpy() == batch_idx)[0]
......@@ -249,16 +243,15 @@ class DGHead(BaseMapHead):
})
for i in batch2seq:
pre = preds['polylines'][i].detach().cpu().numpy()
pre_msk = preds['polyline_masks'][i].detach().cpu().numpy()
valid_idx = np.nonzero(pre_msk)[0][:-1]
# From [200,1] to [199,0] to (1,0)
line = (pre[valid_idx].reshape(-1, coord_dim) - 1) / (range_size-1)
line = (pre[valid_idx].reshape(-1, coord_dim) - 1) / (range_size - 1)
ret_dict_single['vectors'].append(line)
# if gts is not None:
# ret_dict_single['groundTruth'] = rec_groundtruth
......@@ -266,8 +259,8 @@ class DGHead(BaseMapHead):
return ret_list
def pack_groundtruth(batch_idx,gts,tokens,range_size,gen_net_name='gen',coord_dim=2):
def pack_groundtruth(batch_idx, gts, tokens, range_size, gen_net_name='gen', coord_dim=2):
if 'keypoints' in gts['det']:
gt_bbox = \
gts['det']['keypoints'][batch_idx].detach().cpu().numpy()
......@@ -281,7 +274,7 @@ def pack_groundtruth(batch_idx,gts,tokens,range_size,gen_net_name='gen',coord_di
batch2seq = np.nonzero(
gts['gen']['lines_bs_idx'].cpu().numpy() == batch_idx)[0]
ret_groundtruth = {
'token': tokens[batch_idx],
'nline': len(batch2seq),
......@@ -290,16 +283,16 @@ def pack_groundtruth(batch_idx,gts,tokens,range_size,gen_net_name='gen',coord_di
}
for i in batch2seq:
gt_line =\
gt_line = \
gts['gen']['polylines'].detach().cpu().numpy()[i]
gt_msk = gts['gen']['polyline_masks'].detach().cpu().numpy()[i]
if gen_net_name == 'gen_gmm':
valid_idx = np.nonzero(gt_msk)[0]
else:
valid_idx = np.nonzero(gt_msk)[0][:-1]
# From [200,1] to [199,0] to (1,0)
line = (gt_line[valid_idx].reshape(-1, coord_dim) - 1) / (range_size-1)
line = (gt_line[valid_idx].reshape(-1, coord_dim) - 1) / (range_size - 1)
ret_groundtruth['lines'].append(line)
return det_gt, ret_groundtruth
autonomous_driving/Online-HD-Map-Construction/src/models/heads/map_element_detector.py
View file @ 41b18fd8
import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
from mmcv.cnn import Conv2d, Linear
from mmcv.runner import force_fp32
from torch.distributions.categorical import Categorical
from mmdet.core import (multi_apply, build_assigner, build_sampler,
reduce_mean)
from mmdet.models import HEADS
from .detr_bbox import DETRBboxHead
from mmdet.models.utils.transformer import inverse_sigmoid
from mmdet.models import build_loss
from mmcv.cnn import Linear, build_activation_layer, bias_init_with_prob
from mmcv.cnn import (Conv2d, Linear, bias_init_with_prob,
build_activation_layer)
from mmcv.cnn.bricks.transformer import build_positional_encoding
from mmcv.runner import force_fp32
from mmdet.core import build_assigner, build_sampler, multi_apply, reduce_mean
from mmdet.models import HEADS, build_loss
from mmdet.models.utils import build_transformer
from mmdet.models.utils.transformer import inverse_sigmoid
@HEADS.register_module(force=True)
class MapElementDetector(nn.Module):
def __init__(self,
canvas_size=(400, 200),
discrete_output=False,
separate_detect=False,
mode='xyxy',
bbox_size=None,
coord_dim=2,
def __init__(self,
canvas_size=(400, 200),
discrete_output=False,
separate_detect=False,
mode='xyxy',
bbox_size=None,
coord_dim=2,
kp_coord_dim=2,
num_classes=3,
in_channels=128,
......@@ -41,8 +38,8 @@ class MapElementDetector(nn.Module):
positional_encoding: dict = None,
loss_cls: dict = None,
loss_reg: dict = None,
train_cfg: dict = None,):
train_cfg: dict = None, ):
super().__init__()
assigner = train_cfg['assigner']
......@@ -65,7 +62,7 @@ class MapElementDetector(nn.Module):
if loss_cls['use_sigmoid']:
self.cls_out_channels = num_classes
else:
self.cls_out_channels = num_classes+1
self.cls_out_channels = num_classes + 1
self.iterative = iterative
self.num_reg_fcs = num_reg_fcs
......@@ -82,7 +79,7 @@ class MapElementDetector(nn.Module):
self.separate_detect = separate_detect
self.discrete_output = discrete_output
self.bbox_size = 3 if mode=='sce' else 2
self.bbox_size = 3 if mode == 'sce' else 2
if bbox_size is not None:
self.bbox_size = bbox_size
self.coord_dim = coord_dim # for xyz
......@@ -115,16 +112,16 @@ class MapElementDetector(nn.Module):
# query_pos_embed & query_embed
self.query_embedding = nn.Embedding(self.num_query,
self.embed_dims*2)
self.embed_dims * 2)
# for bbox parameter xstart, ystart, xend, yend
self.bbox_embedding = nn.Embedding( self.bbox_size,
self.embed_dims*2)
self.bbox_embedding = nn.Embedding(self.bbox_size,
self.embed_dims * 2)
def _init_branch(self,):
def _init_branch(self, ):
"""Initialize classification branch and regression branch of head."""
fc_cls = Linear(self.embed_dims*self.bbox_size, self.cls_out_channels)
fc_cls = Linear(self.embed_dims * self.bbox_size, self.cls_out_channels)
# fc_cls = Linear(self.embed_dims, self.cls_out_channels)
reg_branch = []
......@@ -135,12 +132,13 @@ class MapElementDetector(nn.Module):
if self.discrete_output:
reg_branch.append(nn.Linear(
self.embed_dims, max(self.canvas_size), bias=True,))
self.embed_dims, max(self.canvas_size), bias=True, ))
else:
reg_branch.append(nn.Linear(
self.embed_dims, self.coord_dim, bias=True,))
self.embed_dims, self.coord_dim, bias=True, ))
reg_branch = nn.Sequential(*reg_branch)
# add sigmoid or not
def _get_clones(module, N):
......@@ -240,29 +238,29 @@ class MapElementDetector(nn.Module):
[nb_dec, bs, num_query, num_points, 2].
'''
(global_context_embedding, sequential_context_embeddings) =\
(global_context_embedding, sequential_context_embeddings) = \
self._prepare_context(context)
x = sequential_context_embeddings
B, C, H, W = x.shape
query_embedding = self.query_embedding.weight[None,:,None].repeat(B, 1, self.bbox_size, 1)
query_embedding = self.query_embedding.weight[None, :, None].repeat(B, 1, self.bbox_size, 1)
bbox_embed = self.bbox_embedding.weight
query_embedding = query_embedding + bbox_embed[None,None]
query_embedding = query_embedding.view(B, -1, C*2)
query_embedding = query_embedding + bbox_embed[None, None]
query_embedding = query_embedding.view(B, -1, C * 2)
img_masks = x.new_zeros((B, H, W))
pos_embed = self.positional_encoding(img_masks)
# outs_dec: [nb_dec, bs, num_query, embed_dim]
hs, init_reference, inter_references = self.transformer(
[x,],
[img_masks.type(torch.bool)],
query_embedding,
[pos_embed],
reg_branches= self.reg_branches if self.iterative else None, # noqa:E501
cls_branches= None, # noqa:E501
)
[x, ],
[img_masks.type(torch.bool)],
query_embedding,
[pos_embed],
reg_branches=self.reg_branches if self.iterative else None, # noqa:E501
cls_branches=None, # noqa:E501
)
outs_dec = hs.permute(0, 2, 1, 3)
outputs = []
......@@ -271,23 +269,23 @@ class MapElementDetector(nn.Module):
reference = init_reference
else:
reference = inter_references[i - 1]
outputs.append(self.get_prediction(i,query_feat,reference))
outputs.append(self.get_prediction(i, query_feat, reference))
return outputs
def get_prediction(self, level, query_feat, reference):
bs, num_query, h = query_feat.shape
query_feat = query_feat.view(bs, -1, self.bbox_size,h)
query_feat = query_feat.view(bs, -1, self.bbox_size, h)
ocls = self.pre_branches['cls'][level](query_feat.flatten(-2))
# ocls = ocls.mean(-2)
reference = inverse_sigmoid(reference)
reference = reference.view(bs, -1, self.bbox_size,self.coord_dim)
reference = reference.view(bs, -1, self.bbox_size, self.coord_dim)
tmp = self.pre_branches['reg'][level](query_feat)
tmp[...,:self.kp_coord_dim] = tmp[...,:self.kp_coord_dim] + reference[...,:self.kp_coord_dim]
lines = tmp.sigmoid() # bs, num_query, self.bbox_size,2
tmp[..., :self.kp_coord_dim] = tmp[..., :self.kp_coord_dim] + reference[..., :self.kp_coord_dim]
lines = tmp.sigmoid() # bs, num_query, self.bbox_size,2
lines = lines * self.canvas_size[:self.coord_dim]
lines = lines.flatten(-2)
......@@ -295,7 +293,7 @@ class MapElementDetector(nn.Module):
return dict(
lines=lines, # [bs, num_query, bboxsize*2]
scores=ocls, # [bs, num_query, num_class]
embeddings= query_feat, # [bs, num_query, bbox_size, h]
embeddings=query_feat, # [bs, num_query, bbox_size, h]
)
@force_fp32(apply_to=('score_pred', 'lines_pred', 'gt_lines'))
......@@ -333,7 +331,7 @@ class MapElementDetector(nn.Module):
num_pred_lines = len(lines_pred)
# assigner and sampler
assign_result = self.assigner.assign(preds=dict(lines=lines_pred, scores=score_pred,),
assign_result = self.assigner.assign(preds=dict(lines=lines_pred, scores=score_pred, ),
gts=dict(lines=gt_lines,
labels=gt_labels, ),
gt_bboxes_ignore=gt_bboxes_ignore)
......@@ -345,10 +343,10 @@ class MapElementDetector(nn.Module):
# label targets 0: foreground, 1: background
if self.separate_detect:
labels = gt_lines.new_full((num_pred_lines, ), 1, dtype=torch.long)
labels = gt_lines.new_full((num_pred_lines,), 1, dtype=torch.long)
else:
labels = gt_lines.new_full(
(num_pred_lines, ), self.num_classes, dtype=torch.long)
(num_pred_lines,), self.num_classes, dtype=torch.long)
labels[pos_inds] = gt_labels[sampling_result.pos_assigned_gt_inds]
label_weights = gt_lines.new_ones(num_pred_lines)
......@@ -421,11 +419,11 @@ class MapElementDetector(nn.Module):
(labels_list, label_weights_list,
lines_targets_list, lines_weights_list,
pos_inds_list, neg_inds_list,pos_gt_inds_list) = multi_apply(
self._get_target_single,
preds['scores'], lines_pred,
class_label, bbox,
gt_bboxes_ignore=gt_bboxes_ignore_list)
pos_inds_list, neg_inds_list, pos_gt_inds_list) = multi_apply(
self._get_target_single,
preds['scores'], lines_pred,
class_label, bbox,
gt_bboxes_ignore=gt_bboxes_ignore_list)
num_total_pos = sum((inds.numel() for inds in pos_inds_list))
num_total_neg = sum((inds.numel() for inds in neg_inds_list))
......@@ -464,7 +462,7 @@ class MapElementDetector(nn.Module):
"""
# Get target for each sample
new_gts, num_total_pos, num_total_neg, pos_inds_list, pos_gt_inds_list =\
new_gts, num_total_pos, num_total_neg, pos_inds_list, pos_gt_inds_list = \
self.get_targets(preds, gts, gt_bboxes_ignore_list)
# Batched all data
......@@ -473,7 +471,7 @@ class MapElementDetector(nn.Module):
# construct weighted avg_factor to match with the official DETR repo
cls_avg_factor = num_total_pos * 1.0 + \
num_total_neg * self.bg_cls_weight
num_total_neg * self.bg_cls_weight
if self.sync_cls_avg_factor:
cls_avg_factor = reduce_mean(
preds['scores'].new_tensor([cls_avg_factor]))
......@@ -499,7 +497,7 @@ class MapElementDetector(nn.Module):
# position NLL loss
if self.discrete_output:
loss_reg = -(preds['lines'].log_prob(new_gts['bboxs']) *
new_gts['bboxs_weights']).sum()/(num_total_pos)
new_gts['bboxs_weights']).sum() / (num_total_pos)
else:
loss_reg = self.reg_loss(
preds['lines'], new_gts['bboxs'], new_gts['bboxs_weights'], avg_factor=num_total_pos)
......@@ -613,7 +611,7 @@ class MapElementDetector(nn.Module):
result_dict['bbox'].append(det_preds)
result_dict['scores'].append(scores)
result_dict['labels'].append(det_labels)
result_dict['lines_bs_idx'].extend([i]*nline)
result_dict['lines_bs_idx'].extend([i] * nline)
# for down stream polyline
_bboxs = torch.cat(result_dict['bbox'], dim=0)
......@@ -625,4 +623,4 @@ class MapElementDetector(nn.Module):
result_dict['lines_bs_idx'] = torch.tensor(
result_dict['lines_bs_idx'], device=device).long()
return result_dict
\ No newline at end of file
return result_dict
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment