The anchor feature points generated by cubic polynomial fitting are used. In the given polyfit_drawfunction, degreethe parameters represent the degree of the fitting polynomial.
Specifically, when we use np.polyfita function to perform a polynomial fit to data points, we need to specify a degree. This degree determines the complexity of the polynomial. For example:
-
degree = 1: Linear fitting, which is the simplest straight line fitting. The fitted polynomial form is f(y)=ax+b. -
degree = 2: Quadratic polynomial fitting. The fitted polynomial form is f(y)=ax2+bx+c. -
degree = 3: Cubic polynomial fitting. The fitted polynomial form is f(y)=ax3+bx2+cx+d.
...and so on.
The higher the degree, the more complex the polynomial is and can fit the data points more accurately, but it is also easier to overfit (that is, the model is too complex, relies too much on the training data, and has poor adaptability to new data).
import torch, os, cv2
from utils.dist_utils import dist_print
import torch, os
from utils.common import merge_config, get_model
import tqdm
import torchvision.transforms as transforms
from data.dataset import LaneTestDataset
def pred2coords(pred, row_anchor, col_anchor, local_width = 1, original_image_width = 1640, original_image_height = 590):
batch_size, num_grid_row, num_cls_row, num_lane_row = pred['loc_row'].shape
batch_size, num_grid_col, num_cls_col, num_lane_col = pred['loc_col'].shape
max_indices_row = pred['loc_row'].argmax(1).cpu()
# n , num_cls, num_lanes
valid_row = pred['exist_row'].argmax(1).cpu()
# n, num_cls, num_lanes
max_indices_col = pred['loc_col'].argmax(1).cpu()
# n , num_cls, num_lanes
valid_col = pred['exist_col'].argmax(1).cpu()
# n, num_cls, num_lanes
pred['loc_row'] = pred['loc_row'].cpu()
pred['loc_col'] = pred['loc_col'].cpu()
coords = []
row_lane_idx = [1,2]
col_lane_idx = [0,3]
for i in row_lane_idx:
tmp = []
if valid_row[0,:,i].sum() > num_cls_row / 2:
for k in range(valid_row.shape[1]):
if valid_row[0,k,i]:
all_ind = torch.tensor(list(range(max(0,max_indices_row[0,k,i] - local_width), min(num_grid_row-1, max_indices_row[0,k,i] + local_width) + 1)))
out_tmp = (pred['loc_row'][0,all_ind,k,i].softmax(0) * all_ind.float()).sum() + 0.5
out_tmp = out_tmp / (num_grid_row-1) * original_image_width
tmp.append((int(out_tmp), int(row_anchor[k] * original_image_height)))
coords.append(tmp)
for i in col_lane_idx:
tmp = []
if valid_col[0,:,i].sum() > num_cls_col / 4:
for k in range(valid_col.shape[1]):
if valid_col[0,k,i]:
all_ind = torch.tensor(list(range(max(0,max_indices_col[0,k,i] - local_width), min(num_grid_col-1, max_indices_col[0,k,i] + local_width) + 1)))
out_tmp = (pred['loc_col'][0,all_ind,k,i].softmax(0) * all_ind.float()).sum() + 0.5
out_tmp = out_tmp / (num_grid_col-1) * original_image_height
tmp.append((int(col_anchor[k] * original_image_width), int(out_tmp)))
coords.append(tmp)
return coords
def polyfit_draw(img, coords, degree=3, color=(144, 238, 144), thickness=2):
"""
对车道线坐标进行多项式拟合并在图像上绘制曲线。
:param img: 输入图像
:param coords: 车道线坐标列表
:param degree: 拟合的多项式的度数
:param color: 曲线的颜色
:param thickness: 曲线的宽度
:return: 绘制了曲线的图像
"""
if len(coords) == 0:
return img
x = [point[0] for point in coords]
y = [point[1] for point in coords]
# 对点进行多项式拟合
coefficients = np.polyfit(y, x, degree)
poly = np.poly1d(coefficients)
ys = np.linspace(min(y), max(y), 100)
xs = poly(ys)
for i in range(len(ys) - 1):
start_point = (int(xs[i]), int(ys[i]))
end_point = (int(xs[i+1]), int(ys[i+1]))
cv2.line(img, start_point, end_point, color, thickness)
return img
if __name__ == "__main__":
torch.backends.cudnn.benchmark = True
args, cfg = merge_config()
cfg.batch_size = 1
print('setting batch_size to 1 for demo generation')
dist_print('start testing...')
assert cfg.backbone in ['18','34','50','101','152','50next','101next','50wide','101wide']
if cfg.dataset == 'CULane':
cls_num_per_lane = 18
elif cfg.dataset == 'Tusimple':
cls_num_per_lane = 56
else:
raise NotImplementedError
net = get_model(cfg)
state_dict = torch.load(cfg.test_model, map_location='cpu')['model']
compatible_state_dict = {}
for k, v in state_dict.items():
if 'module.' in k:
compatible_state_dict[k[7:]] = v
else:
compatible_state_dict[k] = v
net.load_state_dict(compatible_state_dict, strict=False)
net.eval()
img_transforms = transforms.Compose([
transforms.Resize((int(cfg.train_height / cfg.crop_ratio), cfg.train_width)),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
if cfg.dataset == 'CULane':
splits = ['test0_normal.txt']
datasets = [LaneTestDataset(cfg.data_root,os.path.join(cfg.data_root, 'list/test_split/'+split),img_transform = img_transforms, crop_size = cfg.train_height) for split in splits]
img_w, img_h = 1570, 660
elif cfg.dataset == 'Tusimple':
splits = ['test.txt']
datasets = [LaneTestDataset(cfg.data_root,os.path.join(cfg.data_root, split),img_transform = img_transforms, crop_size = cfg.train_height) for split in splits]
img_w, img_h = 1280, 720
else:
raise NotImplementedError
for split, dataset in zip(splits, datasets):
loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle = False, num_workers=1)
fourcc = cv2.VideoWriter_fourcc(*'MJPG')
print(split[:-3]+'avi')
vout = cv2.VideoWriter('4.'+'avi', fourcc , 30.0, (img_w, img_h))
for i, data in enumerate(tqdm.tqdm(loader)):
imgs, names = data
imgs = imgs.cuda()
with torch.no_grad():
pred = net(imgs)
vis = cv2.imread(os.path.join(cfg.data_root,names[0]))
coords = pred2coords(pred, cfg.row_anchor, cfg.col_anchor, original_image_width = img_w, original_image_height = img_h)
for lane in coords:
# for coord in lane:
# cv2.circle(vis,coord,1,(0,255,0),-1)
# vis = draw_lanes(vis, coords)
# polyfit_draw(vis, lane)
vis = polyfit_draw(vis, lane) # 对每一条车道线都使用polyfit_draw函数
vout.write(vis)
vout.release()
ps：
Before optimization
Optimized
Video memory utilization
