Instinct323 / wiou

Wise-IoU: Bounding Box Regression Loss with Dynamic Focusing Mechanism
https://arxiv.org/abs/2301.10051
67 stars 4 forks source link

Project Branches

v1

v2 (master)

Example

After initializing your model, initialize the IouLoss and assign it to the model. The instance attribute iou_mean of IouLoss will be output to the state_dict for saving and loading during the training process.

# After initializing your model, integrate the IoU loss module into the model.
m = YourDetectionModel()
m.iouloss = IouLoss(ltype='WIoU', monotonous=False)
# Ensure that `IouLoss` is assigned to the model before executing `load_state_dict`.
m.load_state_dict(torch.load('last.pt'))

Modify the bounding box regression loss in the loss function. Note that the confidence of the bounding box uses $IoU$ instead of $1-\mathcal{L}_{WIoU}$.

m.iouloss.train()
iloss, liou = m.iouloss(xywh2xyxy(pred), xywh2xyxy(gt), ret_iou=True)
# Accumulate the regression loss.
lbox += iloss.mean()
# Adjust the objectness term to use IoU as the bounding box confidence.
tobj = 1 - liou.detach()

Citation

@article{tong2023wise,
  title={Wise-IoU: Bounding Box Regression Loss with Dynamic Focusing Mechanism},
  author={Tong, Zanjia and Chen, Yuhang and Xu, Zewei and Yu, Rong},
  journal={arXiv preprint arXiv:2301.10051},
  year={2023}
}