add collect_key for condlane training

configs/condlane/resnet101_culane.py

@@ -163,7 +163,7 @@
     ),
     #dict(type='Resize', size=(img_width, img_height)),
     dict(type='Normalize', img_norm=img_norm),
-    dict(type='ToTensor', keys=['img', 'gt_hm']),
+    dict(type='ToTensor', keys=['img', 'gt_hm'], collect_keys=['img_metas']),
 ]
 
 

configs/condlane/resnet50_culane.py

@@ -162,7 +162,7 @@
     ),
     #dict(type='Resize', size=(img_width, img_height)),
     dict(type='Normalize', img_norm=img_norm),
-    dict(type='ToTensor', keys=['img', 'gt_hm']),
+    dict(type='ToTensor', keys=['img', 'gt_hm'], collect_keys=['img_metas']),
 ]
 
 

lanedet/datasets/process/transforms.py

@@ -37,17 +37,22 @@ class ToTensor(object):
 
     Args:
         keys (Sequence[str]): Keys that need to be converted to Tensor.
+        collect_keys (Sequence[str]): Keys that need to keep, but not to Tensor.
     """
 
-    def __init__(self, keys=['img', 'mask'], cfg=None):
+    def __init__(self, keys=['img', 'mask'], collect_keys=[], cfg=None):
         self.keys = keys
+        self.collect_keys = collect_keys
 
     def __call__(self, sample):
         data = {}
         if len(sample['img'].shape) < 3:
             sample['img'] = np.expand_dims(sample['img'], -1)
-        for key in self.keys:
-            data[key] = to_tensor(sample[key])
+        for key in sample.keys():
+            if key in self.keys:
+                data[key] = to_tensor(sample[key])
+            if key in self.collect_keys:
+                data[key] = sample[key]
         data['img'] = data['img'].permute(2, 0, 1)
         return data
 

lanedet/models/heads/condlane.py

@@ -713,6 +713,7 @@ def __init__(self,
 
     def loss(self, output, batch):
         img_metas = batch['img_metas']
+        batch.pop('meta')
         return self.loss_impl(output, img_metas, **batch)