[loss.py] there is a bug

[qq852518421]

https://github.com/klintan/pytorch-lanenet/blob/43f1bf3c8ef6163357993d2f8869cd2be74741a3/lanenet/model/loss.py#L61-L67

[qq852518421]

we can change **embedding_i = embedding_b[seg_mask_i] ** to embedding_i = embedding_b * seg_mask_i.

but, the code mean_i = torch.mean(embedding_i, dim=0) means ?

we should change

seg_mask_i = (seg_gt_b == lane_idx)   # seg_mask_i is  [1, H, W]   but  embedding_b is  [dim, H, W]
 if not seg_mask_i.any(): 
     continue 
 embedding_i = embedding_b[:, seg_mask_i[0]]    # seg_mask_i[0] is [H, W] .  to index according the [H, W]
  
 mean_i = torch.mean(embedding_i, dim=1)  # embedding_i is [dim , h*w]. we shoud mean to [dim, 1] -> [dim]
 centroid_mean.append(mean_i)

maybe this is ok!

[zhangzhongyuan1996]

I tried your code ,but it is not work.Why?