6-DoF Pose estimation based on the YOLO framework and tailored to our application of instrument pose estimation in X-Ray images.
Currently the fastest single-shot object 6-DoF pose estimation method at this level of accuracy.
Setup
Clone repo and install requirements.txt in a Python>=3.9.0 environment, including PyTorch==1.9.
git clone https://github.com/cviviers/YOLOv5-6D-Pose.git # clone
cd YOLOv5-6D-Pose
pip install -r requirements.txt # install
cd utils
python3 ./setup.py build_ext --inplace # To compute ADD-S metric
Alternatively, build the dockerfile included or pull our image sudochris/yolo6dpose:v2.
Inference
python detect.py --weights path/to/weights/linemod/cat/best.pt --img 640 --conf 0.25 --source ../data/LINEMOD/cat/JPEGImages --static-camera configs/linemod/linemod_camera.json --mesh-data path/to/data/LINEMOD/cat/cat.ply
Training
python train.py --batch 32 --epochs 5000 --cfg yolov5xv6_pose_bifpn.yaml --hyp configs/hyp.single.yaml --weights yolov5x.pt --data configs/linemod/benchvise.yaml --rect --cache --optimizer Adam
For a detailed description on how to set up YOLOv5-6D, follow our tutorial tutorial. We have also created an easy guide on how to create a single object dataset for using with YOLOv5-6D.
- Release code: code will be released upon paper acceptance.
- Release weights: will be released upon paper acceptance.
- Linemod occlusion results
- Multi-object Pose
- Release code & results for YOLOv8-6D pose
Please consider citing our paper if the project helps your research with the following BibTex:
@ARTICLE{10478293,
author={Viviers, Christiaan G.A. and Filatova, Lena and Termeer, Maurice and De With, Peter H.N. and Sommen, Fons van der},
journal={IEEE Transactions on Image Processing},
title={Advancing 6-DoF Instrument Pose Estimation in Variable X-Ray Imaging Geometries},
year={2024},
keywords={X-ray instrument detection;6-DoF pose estimation;surgical vision;imaging geometry;deep learning},
doi={10.1109/TIP.2024.3378469}}
All assets and code are under the AGPL-3.0 License - in line with YOLOv5 - unless specified otherwise.