Installation¶
CASTLE requires Python 3.10+ and an NVIDIA GPU with CUDA 12.x for full functionality. CPU-only mode is possible but significantly slower.
Method 1: Local Installation (Recommended)¶
Prerequisites¶
- Python 3.10+
- NVIDIA GPU with CUDA 12.x drivers (recommended)
- ~4 GB disk space for model checkpoints
- Git
Step 1: Clone the Repository¶
Step 2: Create a Virtual Environment¶
Step 3: Install Dependencies¶
Do NOT use pip install castle-ai
The requirements.txt is designed for installation from source. Install directly from the cloned repository — the PyPI package may be outdated.
Key dependencies
The main dependencies include:
- PyTorch with CUDA 12.6 support (
torch,torchvision) - cuml-cu12 — NVIDIA RAPIDS for GPU-accelerated clustering
- Gradio — web UI framework
- xFormers — efficient attention for vision models
- OpenCV, NumPy, Matplotlib, Plotly — data processing and visualization
Step 4: Download Model Checkpoints¶
CASTLE uses several pretrained models. You can download them automatically or manually.
Option A: Automatic Download (Script)¶
Option B: Manual Download¶
Download each file and place it in the ckpt/ directory:
| Checkpoint | Source | Download Link |
|---|---|---|
sam_vit_b_01ec64.pth |
Meta (SAM) | Download |
dinov2_vitb14_reg4_pretrain.pth |
Meta (DINOv2) | Download |
R50_DeAOTL_PRE_YTB_DAV.pth |
Google Drive (DeAOT) | Download |
SwinB_DeAOTL_PRE_YTB_DAV.pth |
Google Drive (DeAOT) | Download |
dinov3_vitb16_pretrain_lvd1689m-73cec8be.pth |
Google Drive (DINOv3) | Download |
dinov3_vitl16_pretrain_lvd1689m-8aa4cbdd.pth |
Google Drive (DINOv3) | Download |
Google Drive download issues
If Google Drive links fail due to quota limits, try:
- Open the link in a browser and download manually
- Wait a few hours and retry (quota resets periodically)
- Use
gdownwith the--fuzzyflag:gdown --fuzzy '<URL>'
Step 5: Verify Installation¶
If everything is set up correctly, a Gradio web UI will launch at http://localhost:7860.
Method 2: Docker¶
Docker provides a self-contained environment with all dependencies pre-configured.
Prerequisites¶
- Docker
- NVIDIA Container Toolkit (for GPU support)
Step 1: Clone and Prepare Checkpoints¶
git clone https://github.com/CASTLE-ai/castle-ai.git
cd castle-ai
pip install gdown
bash download_ckpt.sh
Note
Checkpoints must be in the ckpt/ directory before running the container. They are mounted as a volume, so you only need to download them once.
Step 2: Build the Image¶
Step 3: Run the Container¶
Step 4: Access the UI¶
Open http://localhost:7860 in your browser.
The Docker setup automatically:
- Mounts
ckpt/for model checkpoints (persistent) - Mounts
projects/for your analysis data (persistent) - Mounts
demo/for demo videos - Reserves 1 NVIDIA GPU
Method 3: Google Colab¶
The easiest way to try CASTLE — no local setup required.
Performance
Free-tier Colab GPUs are significantly slower than local setups. For production use, a local installation or Docker is recommended.
A video walkthrough is also available:
Troubleshooting¶
CUDA version mismatch¶
Ensure your NVIDIA driver supports CUDA 12.x. Check with nvidia-smi.
PyTorch not finding GPU¶
import torch
print(torch.cuda.is_available()) # Should print True
print(torch.cuda.get_device_name(0))
If False, reinstall PyTorch with the correct CUDA version from pytorch.org.
gdown failing (quota exceeded)¶
Google Drive has daily download quotas. Options:
- Download manually via browser
- Wait and retry later
- Try:
gdown --fuzzy '<URL>'
cuml installation failures¶
cuml-cu12 requires an NVIDIA GPU and CUDA 12.x. If installation fails:
- Verify CUDA is installed:
nvcc --version - Check CUDA version matches (12.x required)
- On systems without GPU, you may need to skip cuml (some clustering features will be unavailable)
ffmpeg not found¶
Some video processing features require ffmpeg: