How rafat2205/surgikneevision-3d-reconstruction-and-registration-of-knee-for-surgical-navigation works
1. What this repo is
SurgiKneeVision is a research-oriented Python project for knee surface capture and 3D registration during total knee arthroplasty. The live pipeline uses Microsoft Azure Kinect for RGB/depth acquisition, YOLO for object detection, SAM for segmentation, and then evaluates the segmented knee depth/point cloud against registration methods.
2. Top-level architecture split
The repository is split into two clear halves:
- Runtime inference scripts that run on camera input and produce detections, masks, depth/3D overlays, and registration results.
- Offline experimentation/training artifacts in notebooks and training folders for model development, evaluation, visualization, and dataset inspection.
3. Primary live pipeline
The main data flow is:
Azure Kinect frame stream → YOLO knee detection → SAM mask refinement → segmented depth extraction → 3D visualization / registration.
The scripts expose toggles for showing the YOLO box, SAM mask, and 3D bounding-box overlays, which suggests the system is designed as an interactive live demo rather than a batch-only pipeline.
4. Main runtime entry points
There are several focused CLI scripts, each wrapping a slightly different end-to-end flow:
- surgikneevision.py: the full integrated application.
- yolo_inference_az.py: YOLO-only inference on Azure Kinect input.
- o3d_sam_inference_az.py: segmentation plus 3D/point-cloud-focused output.
- yolosam_inference_az.py: combined YOLO + SAM inference with visualization options.
These form the operational boundary of the repo and are the scripts most likely to appear in a deployment or demo diagram.
5. Detection and segmentation subsystem
YOLO and SAM are treated as separate model stages:
- YOLO proposes the knee region via bounding boxes.
- SAM uses those boxes as prompts to generate a more precise mask.
This separation is important architecturally because detection narrows the search space and SAM turns a coarse localization into a pixel-level region that can be projected into depth/3D.
6. 3D registration subsystem
The repo contains multiple registration approaches, indicating a benchmarking-style architecture:
- ICP_reg.py
- RANSAC_reg.py
- global_reg.py
- Robust_ICP.py
- DCP_model.py
- DCP_reg.ipynb, DCP_ICP.ipynb
This shows a comparison between classical point-cloud registration and deep learning-based registration. The segmented knee point cloud is the common input boundary for these algorithms.
7. Offline research and analysis layer
Many notebooks are used for exploratory work rather than production runtime:
- dataset_inspection.ipynb
- plytoxyz.ipynb
- CT_to_pc.ipynb
- SICP_vis.ipynb
- SAM_*training*.ipynb, SAM_*testing*.ipynb
These notebooks support dataset QA, format conversion, registration experiments, and model training/validation. They are a separate analysis layer, not part of the real-time inference path.
8. Model training and artifact management
YOLO training is isolated under yolo_train/, including:
- train_yolo.py
- runs/detect/train/...
- saved weights such as best.pt and last.pt
This indicates the repo keeps both training code and experiment outputs. The README also says SAM weights are loaded from an external Google Drive folder, so the runtime depends on local model artifacts copied into a models/ directory.
9. Shared utilities
utils.py likely contains helper functions shared across inference and registration scripts: image/point-cloud conversion, drawing, coordinate handling, or model-loading helpers. In architecture terms, it acts as the glue layer between camera input, model inference, and 3D processing.
10. Data and external dependencies
The project is centered on proprietary research data and includes a sample point cloud file (scan_knee.ply) plus model outputs. It also depends on external model files that are not committed in the repo. The presence of separate requirements_pip.txt and requirements_conda.txt suggests environment setup is handled explicitly, which matters because camera SDKs, YOLO/SAM, and 3D processing stacks often require careful version alignment.
11. How to draw the architecture
A good diagram would show:
- Input layer: Azure Kinect RGB + depth stream
- Perception layer: YOLO detector → SAM segmenter
- Geometry layer: masked depth → point cloud / 3D overlay
- Registration layer: ICP / RANSAC / global / DCP variants
- Offline layer: training notebooks, evaluation notebooks, dataset conversion, and saved weights
- Artifacts: trained YOLO weights, SAM weights, and research point clouds
The key boundary is between live inference and offline experimentation/training.
Main entry points: utils.py Utils · surgikneevision.py Main app · yolosam_inference_az.py YOLO + SAM · o3d_sam_inference_az.py 3D SAM
Components 21
Live inference
Interactive Azure Kinect pipelines for detection, segmentation, depth extraction, and visualization.
- Azure Kinect · RGB and depth input source for live capture.
- Main app · Integrated end-to-end knee capture, inference, and 3D registration demo.
surgikneevision.py - YOLO infer · YOLO-only Azure Kinect inference entry point.
yolo_inference_az.py - YOLO + SAM · Combined detection and segmentation pipeline with visualization toggles.
yolosam_inference_az.py - 3D SAM · Segmentation-focused live pipeline emphasizing Open3D and 3D output.
o3d_sam_inference_az.py - YOLO · Finds the knee region with bounding boxes.
- SAM · Refines the YOLO box into a pixel-accurate knee mask.
- Masked depth · Projects the segmentation onto depth data to isolate knee surface points.
- 3D view · Overlay and point-cloud visualization for live inspection.
Registration methods
Classical and deep registration variants used to align segmented knee point clouds.
- ICP · Iterative closest point registration baseline.
ICP_reg.py - RANSAC · Robust feature-based registration approach.
RANSAC_reg.py - Global reg · Global alignment method for point clouds.
global_reg.py - Robust ICP · ICP variant designed for noisy or partial scans.
Robust_ICP.py - DCP · Deep Closest Point model for learned point-cloud registration.
DCP_model.py
Offline research
Notebooks for dataset inspection, conversion, training, testing, and visualization.
- Reg notebooks · Notebook-based registration experiments and comparisons.
- Data notebooks · Dataset inspection and format conversion notebooks.
- SAM notebooks · Notebook workflows for SAM training and prompt-based evaluation.
Model artifacts
YOLO training code and saved experiment outputs; external SAM weights are referenced at runtime.
- Train YOLO · YOLO training entry point for the custom knee dataset.
yolo_train/train_yolo.py - YOLO weights · Saved detector checkpoints from training runs.
yolo_train/runs/detect/train/weights/best.pt
Connections 31
- Azure Kinect → Main app: frames · Live RGB/depth stream enters the integrated demo.
- Azure Kinect → YOLO infer: frames · Camera input feeds the detector-only script.
- Azure Kinect → YOLO + SAM: frames · Camera input feeds the combined perception pipeline.
- Azure Kinect → 3D SAM: frames · Camera input feeds the 3D-focused segmentation pipeline.
- Main app → YOLO: detect · Main app invokes YOLO for localization.
- YOLO infer → YOLO: detect · Detector-only runtime stage.
- YOLO + SAM → YOLO: detect · YOLO proposes the knee bounding box.
- YOLO → SAM: prompt · Detected box is used to prompt SAM.
- YOLO + SAM → SAM: segment · Combined pipeline runs mask refinement.
- 3D SAM → SAM: segment · 3D pipeline uses SAM for knee isolation.
- SAM → Masked depth: mask depth · Mask is projected into depth space to extract knee surface points.
- Masked depth → 3D view: render · Segmented geometry is visualized as 3D output.
- Main app → 3D view: overlay · Integrated demo shows live overlays and 3D bounding boxes.
- Masked depth → ICP: register · Masked point cloud is aligned with classical registration baselines.
- Masked depth → RANSAC: register · Masked point cloud is aligned with RANSAC-based registration.
- Masked depth → Global reg: register · Masked point cloud is aligned with global registration.
- Masked depth → Robust ICP: register · Masked point cloud is aligned with a robust ICP variant.
- Masked depth → DCP: register · Masked point cloud is also evaluated with learned registration.
- Reg notebooks → ICP: study · Notebooks compare registration approaches.
- Reg notebooks → RANSAC: study · Notebooks compare registration approaches.
- Reg notebooks → Global reg: study · Notebooks compare registration approaches.
- Reg notebooks → Robust ICP: study · Notebooks compare registration approaches.
- Reg notebooks → DCP: study · Notebooks compare registration approaches.
- Data notebooks → Scan knee: inspect · Dataset and point-cloud notebooks use sample scans and conversions.
- SAM notebooks → SAM: train/test · SAM notebook workflows develop and validate segmentation behavior.
- Train YOLO → YOLO weights: produce · Training script outputs detector checkpoints.
- YOLO weights → YOLO: load · Live inference consumes trained YOLO weights.
- Utils → Main app: shared · Common helper functions support the integrated app.
- Utils → YOLO + SAM: shared · Common helper functions support the combined pipeline.
- Utils → 3D SAM: shared · Common helper functions support the 3D pipeline.
- Utils → ICP: shared · Common helper functions support registration scripts.