1.在ultralytics/data/base.py 中
298行def get_image_and_label 函数中,模仿增加meta项
# label['img'], label['ori_shape'], label['resized_shape']= self.load_image(index)
label['img'], label['ori_shape'], label['resized_shape'],label['meta'] = self.load_image(index)#sxx
与此同时,相应地,修改load_image
2.在ultralytics/data/augment.py文件中 找到 _cat_labels 函数,并修改
final_labels = {
'im_file': mosaic_labels[0]['im_file'],
'ori_shape': mosaic_labels[0]['ori_shape'],
'resized_shape': (imgsz, imgsz),
'cls': np.concatenate(cls, 0),
'instances': Instances.concatenate(instances, axis=0),
'mosaic_border': self.border,
'meta':mosaic_labels[0]['meta']
} # final_labels
找到Format类的__call__l函数,因为我的数据和img比较相似,所以模仿img处理meta,如果在运行程序时下面4修改的地方没报错就不用管
meta = labels.pop('meta') labels['meta'] = self._format_img(meta)#把ndaray改成tensor
在ultralytics/data/dataset.py文件中
collate_fn函数中,要模仿img 将batch个tensor的list转成tensor
if k == 'meta':#sxx value = torch.stack(value, 0)
3.最终,在ultralytics/engine/trainer.py 的 _do_train()函数中进行debug,
在该语句之后for i, batch in pbar:
batch 中就有meta信息了
4.ultralytics/models/yolo/detect/train.py文件的
preprocess_batch函数中 模仿对img的操作将meta到GPU
batch['meta'] = torch.Tensor(batch['meta'])#sxx 转成tensor batch['meta'] = batch['meta'].to(self.device, non_blocking=True)#sxx
5.模型解析过程中修改
ultralytics/nn/tasks.py文件中BaseModal类的
forward,predict,predict_once函数都添加一个参数x0=None
在DetectionModal类中覆盖父类的
def loss(self, batch, preds=None): """ Compute loss Args: batch (dict): Batch to compute loss on preds (torch.Tensor | List[torch.Tensor]): Predictions. """ if not hasattr(self, 'criterion'): self.criterion = self.init_criterion() preds = self.forward(batch['img'],batch['meta']) if preds is None else preds return self.criterion(preds, batch)
def _predict_once(self, x, x0=None,profile=False, visualize=False):也要覆盖