YOLOv5-6D Pose: Advancing 6-DoF Instrument Pose Estimation in Variable X-Ray Imaging Geometries

6-DoF Pose estimation based on the YOLO framework and tailored to our application of instrument pose estimation in X-Ray images.

Project Page | Paper | arXiv |

Results

Currently the fastest single-shot object 6-DoF pose estimation method at this level of accuracy.

Quick Start

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.

Download the weights and data

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 

Detailed Setup

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.

TODO List

• 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

Citation

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}}

License

All assets and code are under the AGPL-3.0 License - in line with YOLOv5 - unless specified otherwise.