An AI-powered medical imaging pipeline for automated Cardio-Thoracic Ratio (CTR) analysis from fetal ultrasound video streams.
This repository implements a two-stage deep learning pipeline for automated cardiac measurement from fetal ultrasound:
- Semantic Segmentation — Identifies cardiac and thorax regions using Attention U-Net or HRNet architectures
- Landmark Regression — Detects anatomical keypoints for precise CTR measurement
- Quality Assessment — Automatically selects high-quality frames for reliable analysis
- 🎥 Video Processing Pipeline — Processes ultrasound video streams frame-by-frame
- 🧠 Multi-Model Support — Supports Attention U-Net, HRNet, and TransUNet architectures
- 📊 Quality-Aware Selection — Automatically identifies optimal frames based on brightness, contrast, and anatomical visibility
- ⚡ GPU Acceleration — CUDA-enabled inference for real-time processing
- 🔧 Modular Design — Easily swap models and configure thresholds
┌─────────────────────────────────────────────────────────────────────────┐
│ Video Input │
└────────────────────────────────┬────────────────────────────────────────┘
│
▼
┌──────────────────────────────────┐
│ Frame Preprocessing │
│ (Resize → Normalize → Tensor) │
└──────────────────────────────────┘
│
▼
┌──────────────────────────────────┐
│ Segmentation Model │
│ (Attention U-Net / HRNet) │
│ │
│ Classes: Background, Cardiac, │
│ Thorax │
└──────────────────────────────────┘
│
▼
┌──────────────────────────────────┐
│ Quality Assessment │
│ • Brightness/Contrast Check │
│ • Chamber Detection │
│ • Edge Energy Analysis │
└──────────────────────────────────┘
│
(If quality threshold met)
│
▼
┌──────────────────────────────────┐
│ Landmark Detection │
│ (HRNet-based Regression) │
│ │
│ 8 Keypoints for CTR Calc │
└──────────────────────────────────┘
│
▼
┌──────────────────────────────────┐
│ CTR Measurement │
└──────────────────────────────────┘
Origin_Medical_CV/
├── main.py # Main entry point - video processing pipeline
├── config.py # Central configuration (thresholds, paths, device)
├── inference_engine.py # Unified inference for segmentation & landmarks
├── model_loader.py # Dynamic model loading with weight management
├── quality_analyzer.py # Image quality assessment heuristics
├── requirements.txt # Python dependencies
│
├── models/ # Model architectures
│ ├── attention_unet.py # Attention U-Net for segmentation
│ ├── hrnet_models.py # HRNet for segmentation & landmark detection
│ └── hrnet_inference.py # HRNet inference utilities
│
├── Segmentation/ # Segmentation-specific training & inference
│ ├── train/ # Training scripts
│ ├── inference/ # Inference utilities
│ └── weights/ # Trained model weights
│
├── landmark_regression/ # Landmark detection module
│ ├── train/ # Training scripts
│ ├── inference/ # Inference utilities
│ └── weights/ # Trained model weights
│
├── AI_Pipeline/ # End-to-end pipeline integration
│ └── Results/ # Pipeline outputs
│
├── utils/ # Utility scripts
│ ├── cleaning_data.py # Dataset cleaning utilities
│ ├── generate_masks.py # Mask generation from annotations
│ ├── evaluate.py # Model evaluation scripts
│ └── train_with_augmentation.py # Augmented training utilities
│
└── testing/ # Test images and validation data
# Clone the repository
git clone https://github.com/Kirthik1824/Origin_Medical_CV.git
cd Origin_Medical_CV
# Create virtual environment (recommended)
python -m venv venv
source venv/bin/activate # Linux/Mac
# or: venv\Scripts\activate # Windows
# Install dependencies
pip install -r requirements.txtpython main.py \
--seg-model AttentionUNet \
--seg-weights path/to/segmentation_weights.pth \
--landmark-model HRNet \
--landmark-weights path/to/landmark_weights.pth \
--input-video path/to/ultrasound_video.mp4 \
--quality-threshold 0.5 \
--frame-skip 1| Argument | Description | Default |
|---|---|---|
--seg-model |
Segmentation model (AttentionUNet, HRNet) |
Required |
--seg-weights |
Path to segmentation model weights | Required |
--landmark-model |
Landmark model (HRNet) |
Required |
--landmark-weights |
Path to landmark model weights | Required |
--input-video |
Path to input ultrasound video | Required |
--quality-threshold |
Minimum quality score (0-1) | 0.5 |
--frame-skip |
Process every nth frame | 1 |
python quality_analyzer.py path/to/image_folder/| Model | Description | Classes |
|---|---|---|
| Attention U-Net | U-Net with attention gates for improved feature selection | 3 (Background, Cardiac, Thorax) |
| HRNet | High-Resolution Network maintaining spatial precision | 3 |
| TransUNet | Transformer-based U-Net (optional) | 3 |
| Model | Description | Keypoints |
|---|---|---|
| HRNet | Regression-based landmark detection | 8 (4 points × 2 coords) |
Key parameters in config.py:
# Image dimensions
IMG_HEIGHT = 512
IMG_WIDTH = 512
# Model configuration
NUM_CLASSES_SEG = 3 # Background, Cardiac, Thorax
NUM_LANDMARKS = 8 # 4 keypoints (x, y each)
# Quality thresholds
QUALITY_THRESHOLD = 0.5
MIN_CHAMBER_AREA = 100 # pixels
# Normalization (ImageNet)
MEAN = [0.485, 0.456, 0.406]
STD = [0.229, 0.224, 0.225]The pipeline automatically filters frames based on:
- Brightness — Mean pixel intensity within acceptable range (20-230)
- Contrast — Standard deviation above minimum threshold (25)
- Chamber Detection — Presence of at least 2 chamber-like regions
- Edge Energy — Sufficient vertical and horizontal edge structure
Results are saved to the outputs/ directory:
- Annotated frames with segmentation masks
- Landmark coordinates
- Quality metrics
- Python 3.8+
- PyTorch 1.9+
- OpenCV
- Albumentations
- NumPy
- tqdm
This project is part of Origin Medical Research Laboratory's computer vision initiatives.
If you use this work in your research, please cite:
@software{origin_medical_cv,
title = {Origin Medical CV: Automated Cardio-Thoracic Ratio Analysis},
author = {OMRL Team},
year = {2024},
url = {https://github.com/Kirthik1824/Origin_Medical_CV}
}