Spaces:
Paused
Paused
Georg commited on
Commit ·
24857f8
1
Parent(s): 9f5db88
initial commit
Browse files- .gitignore +46 -0
- DEPLOYMENT.md +184 -0
- Dockerfile +66 -0
- QUICKSTART.md +270 -0
- README.md +111 -5
- STATUS.md +383 -0
- app.py +569 -0
- client.py +212 -0
- deploy.sh +108 -0
- download_weights.py +73 -0
- estimator.py +413 -0
- requirements.txt +19 -0
- test_local.py +264 -0
.gitignore
ADDED
|
@@ -0,0 +1,46 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Python
|
| 2 |
+
__pycache__/
|
| 3 |
+
*.py[cod]
|
| 4 |
+
*$py.class
|
| 5 |
+
*.so
|
| 6 |
+
.Python
|
| 7 |
+
build/
|
| 8 |
+
develop-eggs/
|
| 9 |
+
dist/
|
| 10 |
+
downloads/
|
| 11 |
+
eggs/
|
| 12 |
+
.eggs/
|
| 13 |
+
lib/
|
| 14 |
+
lib64/
|
| 15 |
+
parts/
|
| 16 |
+
sdist/
|
| 17 |
+
var/
|
| 18 |
+
wheels/
|
| 19 |
+
*.egg-info/
|
| 20 |
+
.installed.cfg
|
| 21 |
+
*.egg
|
| 22 |
+
|
| 23 |
+
# Virtual environments
|
| 24 |
+
venv/
|
| 25 |
+
env/
|
| 26 |
+
ENV/
|
| 27 |
+
|
| 28 |
+
# IDE
|
| 29 |
+
.vscode/
|
| 30 |
+
.idea/
|
| 31 |
+
*.swp
|
| 32 |
+
*.swo
|
| 33 |
+
|
| 34 |
+
# Model weights and data
|
| 35 |
+
weights/
|
| 36 |
+
*.pth
|
| 37 |
+
*.ckpt
|
| 38 |
+
*.safetensors
|
| 39 |
+
|
| 40 |
+
# Gradio cache
|
| 41 |
+
gradio_cached_examples/
|
| 42 |
+
flagged/
|
| 43 |
+
|
| 44 |
+
# Test images
|
| 45 |
+
test_images/
|
| 46 |
+
reference_images/
|
DEPLOYMENT.md
ADDED
|
@@ -0,0 +1,184 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# FoundationPose Hugging Face Space Deployment Guide
|
| 2 |
+
|
| 3 |
+
This directory contains the code for deploying FoundationPose on Hugging Face Spaces with ZeroGPU support.
|
| 4 |
+
|
| 5 |
+
## Current Status
|
| 6 |
+
|
| 7 |
+
- ✅ Gradio app structure created
|
| 8 |
+
- ✅ API endpoints defined (/initialize, /estimate)
|
| 9 |
+
- ✅ ZeroGPU decorators added (@spaces.GPU)
|
| 10 |
+
- ✅ Client library for API calls created
|
| 11 |
+
- ⚠️ FoundationPose model integration incomplete (placeholder code)
|
| 12 |
+
|
| 13 |
+
## Next Steps
|
| 14 |
+
|
| 15 |
+
### 1. Complete FoundationPose Integration
|
| 16 |
+
|
| 17 |
+
The current `app.py` has placeholder code marked with `# TODO` comments. You need to:
|
| 18 |
+
|
| 19 |
+
1. **Install FoundationPose in the Space**:
|
| 20 |
+
- Add FoundationPose installation to requirements.txt or use a custom Dockerfile
|
| 21 |
+
- Download pre-trained weights (need to be included in the Space or downloaded at startup)
|
| 22 |
+
|
| 23 |
+
2. **Implement model initialization** (line ~40 in app.py):
|
| 24 |
+
```python
|
| 25 |
+
# Replace the TODO with actual FoundationPose initialization
|
| 26 |
+
from FoundationPose import FoundationPoseEstimator
|
| 27 |
+
self.model = FoundationPoseEstimator(device=self.device)
|
| 28 |
+
```
|
| 29 |
+
|
| 30 |
+
3. **Implement object registration** (line ~70):
|
| 31 |
+
```python
|
| 32 |
+
# Replace the TODO with actual registration
|
| 33 |
+
self.model.register_object(object_id, reference_images, camera_intrinsics)
|
| 34 |
+
```
|
| 35 |
+
|
| 36 |
+
4. **Implement pose estimation** (line ~120):
|
| 37 |
+
```python
|
| 38 |
+
# Replace the TODO with actual inference
|
| 39 |
+
result = self.model.estimate_pose(object_id, query_image, camera_intrinsics)
|
| 40 |
+
```
|
| 41 |
+
|
| 42 |
+
### 2. Handle Model Weights
|
| 43 |
+
|
| 44 |
+
FoundationPose requires pre-trained weights. Options:
|
| 45 |
+
|
| 46 |
+
**Option A: Git LFS (Recommended)**
|
| 47 |
+
```bash
|
| 48 |
+
cd foundationpose
|
| 49 |
+
git lfs install
|
| 50 |
+
mkdir weights
|
| 51 |
+
# Download weights from FoundationPose repo
|
| 52 |
+
wget https://... -O weights/model.pth
|
| 53 |
+
git lfs track "weights/*.pth"
|
| 54 |
+
git add weights/model.pth .gitattributes
|
| 55 |
+
git commit -m "Add model weights"
|
| 56 |
+
```
|
| 57 |
+
|
| 58 |
+
**Option B: Download at Runtime**
|
| 59 |
+
Add to `app.py`:
|
| 60 |
+
```python
|
| 61 |
+
def download_weights():
|
| 62 |
+
from huggingface_hub import hf_hub_download
|
| 63 |
+
weights_path = hf_hub_download(
|
| 64 |
+
repo_id="NVlabs/FoundationPose",
|
| 65 |
+
filename="model.pth"
|
| 66 |
+
)
|
| 67 |
+
return weights_path
|
| 68 |
+
```
|
| 69 |
+
|
| 70 |
+
### 3. Test Locally
|
| 71 |
+
|
| 72 |
+
Before deploying, test the Space locally:
|
| 73 |
+
|
| 74 |
+
```bash
|
| 75 |
+
cd foundationpose
|
| 76 |
+
pip install -r requirements.txt
|
| 77 |
+
python app.py
|
| 78 |
+
```
|
| 79 |
+
|
| 80 |
+
This will start a local Gradio server at http://localhost:7860
|
| 81 |
+
|
| 82 |
+
### 4. Deploy to Hugging Face
|
| 83 |
+
|
| 84 |
+
```bash
|
| 85 |
+
cd foundationpose
|
| 86 |
+
git add .
|
| 87 |
+
git commit -m "Add FoundationPose inference implementation"
|
| 88 |
+
git push
|
| 89 |
+
```
|
| 90 |
+
|
| 91 |
+
The Space will automatically rebuild and deploy.
|
| 92 |
+
|
| 93 |
+
### 5. Monitor GPU Usage
|
| 94 |
+
|
| 95 |
+
After deployment:
|
| 96 |
+
1. Check the Space logs for GPU allocation messages
|
| 97 |
+
2. Monitor inference times (cold start vs warm)
|
| 98 |
+
3. Adjust `@spaces.GPU(duration=X)` parameters if needed
|
| 99 |
+
|
| 100 |
+
### 6. Integrate with Training Pipeline
|
| 101 |
+
|
| 102 |
+
Once the Space is working, update the training code:
|
| 103 |
+
|
| 104 |
+
**In training/nova_sim_trainer/perception/foundation_pose_wrapper.py**:
|
| 105 |
+
```python
|
| 106 |
+
from foundationpose.client import FoundationPoseClient
|
| 107 |
+
|
| 108 |
+
class FoundationPoseWrapper(PoseEstimator):
|
| 109 |
+
def __init__(self, api_url: str, ...):
|
| 110 |
+
self.client = FoundationPoseClient(api_url)
|
| 111 |
+
# Initialize with reference images
|
| 112 |
+
ref_images = load_reference_images(reference_dir)
|
| 113 |
+
self.client.initialize(object_id, ref_images)
|
| 114 |
+
|
| 115 |
+
def estimate_poses(self, frame, camera_intrinsics, scene_objects):
|
| 116 |
+
poses = self.client.estimate_pose(self.object_id, frame, camera_intrinsics)
|
| 117 |
+
return [DetectedPose(**pose) for pose in poses]
|
| 118 |
+
```
|
| 119 |
+
|
| 120 |
+
**In training/observations.yaml**:
|
| 121 |
+
```yaml
|
| 122 |
+
perception:
|
| 123 |
+
enabled: true
|
| 124 |
+
model: foundation_pose
|
| 125 |
+
api_url: https://gpue-foundationpose.hf.space
|
| 126 |
+
tracked_objects:
|
| 127 |
+
- object_id: target_cube
|
| 128 |
+
reference_images_dir: ./perception/reference/target_cube
|
| 129 |
+
```
|
| 130 |
+
|
| 131 |
+
## Performance Considerations
|
| 132 |
+
|
| 133 |
+
### ZeroGPU Latency
|
| 134 |
+
- **Cold start**: 15-30 seconds (GPU allocation + model loading)
|
| 135 |
+
- **Warm inference**: 0.5-2 seconds per query
|
| 136 |
+
- **GPU duration**: Tune the `duration` parameter in `@spaces.GPU` decorators
|
| 137 |
+
|
| 138 |
+
### Recommended Usage
|
| 139 |
+
- ✅ **Batch processing**: Process multiple frames in one GPU allocation
|
| 140 |
+
- ✅ **Validation**: Check perception quality on recorded episodes
|
| 141 |
+
- ✅ **Demos**: Show 6D pose estimation capabilities
|
| 142 |
+
- ⚠️ **Real-time training**: Too slow for 30 Hz control loop - use dummy estimator instead
|
| 143 |
+
|
| 144 |
+
### Optimization Tips
|
| 145 |
+
1. **Batch multiple queries** to amortize cold start time
|
| 146 |
+
2. **Keep GPU warm** by sending periodic keep-alive requests
|
| 147 |
+
3. **Use lower resolution** if inference is too slow
|
| 148 |
+
4. **Cache results** for static scenes
|
| 149 |
+
|
| 150 |
+
## Troubleshooting
|
| 151 |
+
|
| 152 |
+
### Space won't start
|
| 153 |
+
- Check Space logs for errors
|
| 154 |
+
- Verify all dependencies in requirements.txt
|
| 155 |
+
- Check Python version compatibility (3.12)
|
| 156 |
+
|
| 157 |
+
### GPU timeout
|
| 158 |
+
- Increase `duration` in `@spaces.GPU(duration=X)`
|
| 159 |
+
- Optimize model inference code
|
| 160 |
+
- Reduce image resolution
|
| 161 |
+
|
| 162 |
+
### Out of memory
|
| 163 |
+
- Reduce batch size
|
| 164 |
+
- Use smaller model variant
|
| 165 |
+
- Request more GPU memory in Space settings
|
| 166 |
+
|
| 167 |
+
## Alternative: Docker Deployment
|
| 168 |
+
|
| 169 |
+
If ZeroGPU is too restrictive, consider running locally with Docker:
|
| 170 |
+
|
| 171 |
+
```bash
|
| 172 |
+
cd foundationpose
|
| 173 |
+
docker build -t foundationpose .
|
| 174 |
+
docker run -p 7860:7860 --gpus all foundationpose
|
| 175 |
+
```
|
| 176 |
+
|
| 177 |
+
Then set `api_url: http://localhost:7860` in observations.yaml.
|
| 178 |
+
|
| 179 |
+
## References
|
| 180 |
+
|
| 181 |
+
- [FoundationPose GitHub](https://github.com/NVlabs/FoundationPose)
|
| 182 |
+
- [Hugging Face Spaces](https://huggingface.co/docs/hub/spaces)
|
| 183 |
+
- [ZeroGPU Documentation](https://huggingface.co/docs/hub/spaces-gpus-zerogpu)
|
| 184 |
+
- [Gradio Documentation](https://www.gradio.app/docs)
|
Dockerfile
ADDED
|
@@ -0,0 +1,66 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# FoundationPose Dockerfile for Hugging Face Spaces with ZeroGPU
|
| 2 |
+
FROM nvidia/cuda:12.1.0-devel-ubuntu22.04
|
| 3 |
+
|
| 4 |
+
# Set environment variables
|
| 5 |
+
ENV DEBIAN_FRONTEND=noninteractive
|
| 6 |
+
ENV CUDA_HOME=/usr/local/cuda
|
| 7 |
+
ENV PATH=${CUDA_HOME}/bin:${PATH}
|
| 8 |
+
ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
|
| 9 |
+
|
| 10 |
+
# Install system dependencies
|
| 11 |
+
RUN apt-get update && apt-get install -y \
|
| 12 |
+
git \
|
| 13 |
+
wget \
|
| 14 |
+
python3.9 \
|
| 15 |
+
python3.9-dev \
|
| 16 |
+
python3-pip \
|
| 17 |
+
libgl1-mesa-glx \
|
| 18 |
+
libglib2.0-0 \
|
| 19 |
+
libeigen3-dev \
|
| 20 |
+
&& rm -rf /var/lib/apt/lists/*
|
| 21 |
+
|
| 22 |
+
# Set Python 3.9 as default
|
| 23 |
+
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.9 1
|
| 24 |
+
RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.9 1
|
| 25 |
+
|
| 26 |
+
# Upgrade pip
|
| 27 |
+
RUN python3 -m pip install --upgrade pip setuptools wheel
|
| 28 |
+
|
| 29 |
+
# Set working directory
|
| 30 |
+
WORKDIR /app
|
| 31 |
+
|
| 32 |
+
# Copy requirements first for better caching
|
| 33 |
+
COPY requirements.txt .
|
| 34 |
+
|
| 35 |
+
# Install Python dependencies
|
| 36 |
+
RUN pip install --no-cache-dir -r requirements.txt
|
| 37 |
+
|
| 38 |
+
# Install FoundationPose dependencies
|
| 39 |
+
RUN pip install --no-cache-dir \
|
| 40 |
+
git+https://github.com/NVlabs/nvdiffrast.git
|
| 41 |
+
|
| 42 |
+
# Install Kaolin (if using model-free mode)
|
| 43 |
+
RUN pip install --no-cache-dir kaolin==0.15.0 -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.0.0_cu118.html
|
| 44 |
+
|
| 45 |
+
# Install PyTorch3D
|
| 46 |
+
RUN pip install --no-cache-dir pytorch3d
|
| 47 |
+
|
| 48 |
+
# Clone FoundationPose repository
|
| 49 |
+
RUN git clone https://github.com/NVlabs/FoundationPose.git /app/FoundationPose
|
| 50 |
+
|
| 51 |
+
# Build FoundationPose C++ extensions
|
| 52 |
+
WORKDIR /app/FoundationPose
|
| 53 |
+
RUN bash build_all.sh || echo "Build completed with warnings"
|
| 54 |
+
|
| 55 |
+
# Copy application files
|
| 56 |
+
WORKDIR /app
|
| 57 |
+
COPY . .
|
| 58 |
+
|
| 59 |
+
# Download model weights
|
| 60 |
+
RUN python3 download_weights.py
|
| 61 |
+
|
| 62 |
+
# Expose port for Gradio
|
| 63 |
+
EXPOSE 7860
|
| 64 |
+
|
| 65 |
+
# Run the application
|
| 66 |
+
CMD ["python3", "app.py"]
|
QUICKSTART.md
ADDED
|
@@ -0,0 +1,270 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# FoundationPose Quick Start Guide
|
| 2 |
+
|
| 3 |
+
## Overview
|
| 4 |
+
|
| 5 |
+
This Hugging Face Space provides two modes:
|
| 6 |
+
|
| 7 |
+
1. **Placeholder Mode** (default) - Returns empty results, useful for testing the API without GPU requirements
|
| 8 |
+
2. **Real Mode** - Uses actual FoundationPose model for 6D pose estimation (requires GPU and model weights)
|
| 9 |
+
|
| 10 |
+
## Testing Locally (Placeholder Mode)
|
| 11 |
+
|
| 12 |
+
The easiest way to test the API structure:
|
| 13 |
+
|
| 14 |
+
```bash
|
| 15 |
+
cd foundationpose
|
| 16 |
+
pip install -r requirements.txt
|
| 17 |
+
python app.py
|
| 18 |
+
```
|
| 19 |
+
|
| 20 |
+
Visit http://localhost:7860 to see the UI.
|
| 21 |
+
|
| 22 |
+
## Deploying to Hugging Face Spaces
|
| 23 |
+
|
| 24 |
+
### Option 1: Placeholder Mode (No Setup Required)
|
| 25 |
+
|
| 26 |
+
Just push to your Space:
|
| 27 |
+
|
| 28 |
+
```bash
|
| 29 |
+
cd foundationpose
|
| 30 |
+
git add .
|
| 31 |
+
git commit -m "Deploy FoundationPose Space"
|
| 32 |
+
git push
|
| 33 |
+
```
|
| 34 |
+
|
| 35 |
+
The Space will run in placeholder mode by default. This is useful for:
|
| 36 |
+
- Testing the API structure
|
| 37 |
+
- Developing client integrations
|
| 38 |
+
- Demos without GPU costs
|
| 39 |
+
|
| 40 |
+
### Option 2: Real FoundationPose (Requires Setup)
|
| 41 |
+
|
| 42 |
+
**Step 1: Clone FoundationPose Repository**
|
| 43 |
+
|
| 44 |
+
```bash
|
| 45 |
+
# Inside your local foundationpose directory
|
| 46 |
+
git clone https://github.com/NVlabs/FoundationPose.git
|
| 47 |
+
```
|
| 48 |
+
|
| 49 |
+
**Step 2: Download Model Weights**
|
| 50 |
+
|
| 51 |
+
Download weights from the official Google Drive:
|
| 52 |
+
https://drive.google.com/drive/folders/1GCyGE-LbFGgRC-FuGsF3a1zeBuzsQ1Da
|
| 53 |
+
|
| 54 |
+
Extract to:
|
| 55 |
+
```
|
| 56 |
+
foundationpose/weights/
|
| 57 |
+
├── 2023-10-28-18-33-37/ (refiner weights)
|
| 58 |
+
└── 2024-01-11-20-02-45/ (scorer weights)
|
| 59 |
+
```
|
| 60 |
+
|
| 61 |
+
**Step 3: Add Weights to Git LFS**
|
| 62 |
+
|
| 63 |
+
```bash
|
| 64 |
+
git lfs install
|
| 65 |
+
git lfs track "weights/**/*.pth"
|
| 66 |
+
git lfs track "weights/**/*.ckpt"
|
| 67 |
+
git add .gitattributes
|
| 68 |
+
git add weights/
|
| 69 |
+
git commit -m "Add model weights"
|
| 70 |
+
```
|
| 71 |
+
|
| 72 |
+
**Step 4: Enable Real Mode**
|
| 73 |
+
|
| 74 |
+
Add to your Space settings (or use .env file locally):
|
| 75 |
+
```
|
| 76 |
+
USE_REAL_MODEL=true
|
| 77 |
+
```
|
| 78 |
+
|
| 79 |
+
**Step 5: Push to HF**
|
| 80 |
+
|
| 81 |
+
```bash
|
| 82 |
+
git push
|
| 83 |
+
```
|
| 84 |
+
|
| 85 |
+
## Using the API
|
| 86 |
+
|
| 87 |
+
### Python Client
|
| 88 |
+
|
| 89 |
+
```python
|
| 90 |
+
from foundationpose.client import FoundationPoseClient
|
| 91 |
+
import cv2
|
| 92 |
+
import numpy as np
|
| 93 |
+
|
| 94 |
+
# Initialize client
|
| 95 |
+
client = FoundationPoseClient("https://gpue-foundationpose.hf.space")
|
| 96 |
+
|
| 97 |
+
# Load reference images
|
| 98 |
+
ref_images = []
|
| 99 |
+
for i in range(1, 16):
|
| 100 |
+
img = cv2.imread(f"reference/image_{i:03d}.jpg")
|
| 101 |
+
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
| 102 |
+
ref_images.append(img)
|
| 103 |
+
|
| 104 |
+
# Register object
|
| 105 |
+
client.initialize("target_cube", ref_images)
|
| 106 |
+
|
| 107 |
+
# Estimate pose
|
| 108 |
+
query_img = cv2.imread("query.jpg")
|
| 109 |
+
query_img = cv2.cvtColor(query_img, cv2.COLOR_BGR2RGB)
|
| 110 |
+
|
| 111 |
+
poses = client.estimate_pose("target_cube", query_img)
|
| 112 |
+
print(f"Detected {len(poses)} poses")
|
| 113 |
+
for pose in poses:
|
| 114 |
+
print(f"Position: {pose['position']}")
|
| 115 |
+
print(f"Orientation: {pose['orientation']}")
|
| 116 |
+
print(f"Confidence: {pose['confidence']}")
|
| 117 |
+
```
|
| 118 |
+
|
| 119 |
+
### Direct HTTP API
|
| 120 |
+
|
| 121 |
+
```bash
|
| 122 |
+
# Initialize
|
| 123 |
+
curl -X POST https://gpue-foundationpose.hf.space/api/initialize \
|
| 124 |
+
-H "Content-Type: application/json" \
|
| 125 |
+
-d '{
|
| 126 |
+
"object_id": "target_cube",
|
| 127 |
+
"reference_images_b64": ["'$(base64 -w 0 ref1.jpg)'", "'$(base64 -w 0 ref2.jpg)'"],
|
| 128 |
+
"camera_intrinsics": "{\"fx\": 500, \"fy\": 500, \"cx\": 320, \"cy\": 240}"
|
| 129 |
+
}'
|
| 130 |
+
|
| 131 |
+
# Estimate
|
| 132 |
+
curl -X POST https://gpue-foundationpose.hf.space/api/estimate \
|
| 133 |
+
-H "Content-Type: application/json" \
|
| 134 |
+
-d '{
|
| 135 |
+
"object_id": "target_cube",
|
| 136 |
+
"query_image_b64": "'$(base64 -w 0 query.jpg)'"
|
| 137 |
+
}'
|
| 138 |
+
```
|
| 139 |
+
|
| 140 |
+
## Integration with robot-ml Training
|
| 141 |
+
|
| 142 |
+
Update `/training/nova_sim_trainer/perception/foundation_pose_wrapper.py`:
|
| 143 |
+
|
| 144 |
+
```python
|
| 145 |
+
from foundationpose.client import FoundationPoseClient
|
| 146 |
+
|
| 147 |
+
class FoundationPoseWrapper(PoseEstimator):
|
| 148 |
+
def __init__(self, api_url: str, tracked_objects: List[Dict], **kwargs):
|
| 149 |
+
super().__init__()
|
| 150 |
+
self.client = FoundationPoseClient(api_url)
|
| 151 |
+
|
| 152 |
+
# Initialize each tracked object
|
| 153 |
+
for obj_config in tracked_objects:
|
| 154 |
+
object_id = obj_config["object_id"]
|
| 155 |
+
ref_dir = Path(obj_config["reference_images_dir"])
|
| 156 |
+
|
| 157 |
+
# Load reference images
|
| 158 |
+
ref_images = []
|
| 159 |
+
for img_path in sorted(ref_dir.glob("*.jpg")):
|
| 160 |
+
img = cv2.imread(str(img_path))
|
| 161 |
+
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
| 162 |
+
ref_images.append(img)
|
| 163 |
+
|
| 164 |
+
# Register object
|
| 165 |
+
self.client.initialize(object_id, ref_images)
|
| 166 |
+
logger.info(f"Registered {object_id} with {len(ref_images)} images")
|
| 167 |
+
|
| 168 |
+
def estimate_poses(self, frame, camera_intrinsics, scene_objects):
|
| 169 |
+
# Call API for pose estimation
|
| 170 |
+
poses = self.client.estimate_pose(
|
| 171 |
+
self.tracked_objects[0]["object_id"], # For now, single object
|
| 172 |
+
frame,
|
| 173 |
+
camera_intrinsics
|
| 174 |
+
)
|
| 175 |
+
|
| 176 |
+
# Convert to DetectedPose format
|
| 177 |
+
return [DetectedPose(**pose) for pose in poses]
|
| 178 |
+
```
|
| 179 |
+
|
| 180 |
+
Update `observations.yaml`:
|
| 181 |
+
|
| 182 |
+
```yaml
|
| 183 |
+
perception:
|
| 184 |
+
enabled: true
|
| 185 |
+
model: foundation_pose
|
| 186 |
+
api_url: https://gpue-foundationpose.hf.space
|
| 187 |
+
tracked_objects:
|
| 188 |
+
- object_id: target_cube
|
| 189 |
+
reference_images_dir: ./perception/reference/target_cube
|
| 190 |
+
```
|
| 191 |
+
|
| 192 |
+
## Performance Tips
|
| 193 |
+
|
| 194 |
+
### Cold Start Latency
|
| 195 |
+
- First request takes 15-30s (GPU allocation + model loading)
|
| 196 |
+
- Subsequent requests: 0.5-2s
|
| 197 |
+
|
| 198 |
+
### Keeping GPU Warm
|
| 199 |
+
Send periodic keep-alive requests:
|
| 200 |
+
```python
|
| 201 |
+
import time
|
| 202 |
+
import threading
|
| 203 |
+
|
| 204 |
+
def keep_warm():
|
| 205 |
+
while True:
|
| 206 |
+
try:
|
| 207 |
+
client.estimate_pose("target_cube", dummy_image)
|
| 208 |
+
except:
|
| 209 |
+
pass
|
| 210 |
+
time.sleep(60) # Every minute
|
| 211 |
+
|
| 212 |
+
threading.Thread(target=keep_warm, daemon=True).start()
|
| 213 |
+
```
|
| 214 |
+
|
| 215 |
+
### Batch Processing
|
| 216 |
+
For recorded episodes, process all frames in one session:
|
| 217 |
+
```python
|
| 218 |
+
# Initialize once
|
| 219 |
+
client.initialize("target_cube", ref_images)
|
| 220 |
+
|
| 221 |
+
# Process all frames
|
| 222 |
+
poses_list = []
|
| 223 |
+
for frame in frames:
|
| 224 |
+
poses = client.estimate_pose("target_cube", frame)
|
| 225 |
+
poses_list.append(poses)
|
| 226 |
+
```
|
| 227 |
+
|
| 228 |
+
## Troubleshooting
|
| 229 |
+
|
| 230 |
+
### Space shows "Placeholder mode"
|
| 231 |
+
- Set `USE_REAL_MODEL=true` in Space secrets
|
| 232 |
+
- Verify weights are uploaded correctly
|
| 233 |
+
- Check Space logs for errors
|
| 234 |
+
|
| 235 |
+
### "Model weights not found"
|
| 236 |
+
- Ensure weights are in `weights/` directory
|
| 237 |
+
- Check git-lfs tracked files: `git lfs ls-files`
|
| 238 |
+
- Re-upload if needed
|
| 239 |
+
|
| 240 |
+
### GPU timeout
|
| 241 |
+
- Increase `@spaces.GPU(duration=X)` in app.py
|
| 242 |
+
- Reduce image resolution
|
| 243 |
+
- Process fewer reference images
|
| 244 |
+
|
| 245 |
+
### Out of memory
|
| 246 |
+
- Use lower resolution images
|
| 247 |
+
- Process fewer objects simultaneously
|
| 248 |
+
- Request more GPU resources in Space settings
|
| 249 |
+
|
| 250 |
+
## Cost Optimization
|
| 251 |
+
|
| 252 |
+
ZeroGPU is free but has usage limits:
|
| 253 |
+
|
| 254 |
+
- **Development**: Use placeholder mode
|
| 255 |
+
- **Testing**: Enable real mode for specific tests only
|
| 256 |
+
- **Production**: Consider dedicated GPU deployment (RunPod, Modal, etc.)
|
| 257 |
+
|
| 258 |
+
## Next Steps
|
| 259 |
+
|
| 260 |
+
1. Test locally in placeholder mode
|
| 261 |
+
2. Upload weights for real mode
|
| 262 |
+
3. Integrate with robot-ml training pipeline
|
| 263 |
+
4. Monitor GPU usage and costs
|
| 264 |
+
5. Optimize batch processing for your use case
|
| 265 |
+
|
| 266 |
+
## Support
|
| 267 |
+
|
| 268 |
+
- **Issues**: https://github.com/gpuschel/robot-ml/issues
|
| 269 |
+
- **FoundationPose**: https://github.com/NVlabs/FoundationPose
|
| 270 |
+
- **HF Spaces**: https://huggingface.co/docs/hub/spaces
|
README.md
CHANGED
|
@@ -1,13 +1,119 @@
|
|
| 1 |
---
|
| 2 |
-
title:
|
| 3 |
-
emoji:
|
| 4 |
-
colorFrom:
|
| 5 |
-
colorTo:
|
| 6 |
sdk: gradio
|
| 7 |
sdk_version: 6.4.0
|
| 8 |
python_version: '3.12'
|
| 9 |
app_file: app.py
|
| 10 |
pinned: false
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 11 |
---
|
| 12 |
|
| 13 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
---
|
| 2 |
+
title: FoundationPose Inference
|
| 3 |
+
emoji: 🎯
|
| 4 |
+
colorFrom: blue
|
| 5 |
+
colorTo: purple
|
| 6 |
sdk: gradio
|
| 7 |
sdk_version: 6.4.0
|
| 8 |
python_version: '3.12'
|
| 9 |
app_file: app.py
|
| 10 |
pinned: false
|
| 11 |
+
tags:
|
| 12 |
+
- computer-vision
|
| 13 |
+
- 6D-pose
|
| 14 |
+
- object-detection
|
| 15 |
+
- robotics
|
| 16 |
+
- zero-gpu
|
| 17 |
---
|
| 18 |
|
| 19 |
+
# FoundationPose Inference Server
|
| 20 |
+
|
| 21 |
+
This Hugging Face Space provides 6D object pose estimation using [FoundationPose](https://github.com/NVlabs/FoundationPose) with ZeroGPU support.
|
| 22 |
+
|
| 23 |
+
## Features
|
| 24 |
+
|
| 25 |
+
- **6D Pose Estimation**: Detect object position and orientation in 3D space
|
| 26 |
+
- **Reference-based Tracking**: Register objects using multiple reference images
|
| 27 |
+
- **REST API**: Easy integration with robotics pipelines
|
| 28 |
+
- **ZeroGPU**: On-demand GPU allocation for efficient inference
|
| 29 |
+
|
| 30 |
+
## Usage
|
| 31 |
+
|
| 32 |
+
### Web Interface
|
| 33 |
+
|
| 34 |
+
1. **Initialize Tab**: Upload reference images of your object from different angles (16-20 recommended)
|
| 35 |
+
2. **Estimate Tab**: Upload a query image to detect the object's 6D pose
|
| 36 |
+
|
| 37 |
+
### HTTP API
|
| 38 |
+
|
| 39 |
+
#### Initialize Object
|
| 40 |
+
|
| 41 |
+
```bash
|
| 42 |
+
curl -X POST https://gpue-foundationpose.hf.space/api/initialize \
|
| 43 |
+
-H "Content-Type: application/json" \
|
| 44 |
+
-d '{
|
| 45 |
+
"object_id": "target_cube",
|
| 46 |
+
"reference_images_b64": ["<base64-jpeg>", ...],
|
| 47 |
+
"camera_intrinsics": "{\"fx\": 500, \"fy\": 500, \"cx\": 320, \"cy\": 240}"
|
| 48 |
+
}'
|
| 49 |
+
```
|
| 50 |
+
|
| 51 |
+
#### Estimate Pose
|
| 52 |
+
|
| 53 |
+
```bash
|
| 54 |
+
curl -X POST https://gpue-foundationpose.hf.space/api/estimate \
|
| 55 |
+
-H "Content-Type: application/json" \
|
| 56 |
+
-d '{
|
| 57 |
+
"object_id": "target_cube",
|
| 58 |
+
"query_image_b64": "<base64-jpeg>",
|
| 59 |
+
"camera_intrinsics": "{\"fx\": 500, \"fy\": 500, \"cx\": 320, \"cy\": 240}"
|
| 60 |
+
}'
|
| 61 |
+
```
|
| 62 |
+
|
| 63 |
+
#### Response Format
|
| 64 |
+
|
| 65 |
+
```json
|
| 66 |
+
{
|
| 67 |
+
"success": true,
|
| 68 |
+
"poses": [
|
| 69 |
+
{
|
| 70 |
+
"object_id": "target_cube",
|
| 71 |
+
"position": {"x": 0.5, "y": 0.3, "z": 0.1},
|
| 72 |
+
"orientation": {"w": 1.0, "x": 0.0, "y": 0.0, "z": 0.0},
|
| 73 |
+
"confidence": 0.95,
|
| 74 |
+
"dimensions": [0.1, 0.1, 0.1]
|
| 75 |
+
}
|
| 76 |
+
]
|
| 77 |
+
}
|
| 78 |
+
```
|
| 79 |
+
|
| 80 |
+
## Integration with robot-ml
|
| 81 |
+
|
| 82 |
+
This Space is designed to work with the [robot-ml](https://github.com/gpuschel/robot-ml) training pipeline:
|
| 83 |
+
|
| 84 |
+
1. Capture reference images: `make capture-reference`
|
| 85 |
+
2. Configure perception in `observations.yaml`:
|
| 86 |
+
```yaml
|
| 87 |
+
perception:
|
| 88 |
+
enabled: true
|
| 89 |
+
model: foundation_pose
|
| 90 |
+
api_url: https://gpue-foundationpose.hf.space
|
| 91 |
+
```
|
| 92 |
+
3. Run training with perception: `make train`
|
| 93 |
+
|
| 94 |
+
## Performance
|
| 95 |
+
|
| 96 |
+
- **Cold Start**: 15-30 seconds (ZeroGPU allocation)
|
| 97 |
+
- **Warm Inference**: 0.5-2 seconds per query
|
| 98 |
+
- **Recommended Use**: Batch processing, validation, demos
|
| 99 |
+
|
| 100 |
+
For real-time training loops (30 Hz), use the local dummy estimator instead.
|
| 101 |
+
|
| 102 |
+
## TODO
|
| 103 |
+
|
| 104 |
+
- [ ] Install actual FoundationPose model and weights
|
| 105 |
+
- [ ] Implement real pose estimation (currently returns placeholder results)
|
| 106 |
+
- [ ] Add pose visualization overlay on query images
|
| 107 |
+
- [ ] Support for CAD models in addition to reference images
|
| 108 |
+
- [ ] Batch inference for multiple objects
|
| 109 |
+
|
| 110 |
+
## Citation
|
| 111 |
+
|
| 112 |
+
```bibtex
|
| 113 |
+
@inproceedings{wen2023foundationpose,
|
| 114 |
+
title={FoundationPose: Unified 6D Pose Estimation and Tracking of Novel Objects},
|
| 115 |
+
author={Wen, Bowen and Yang, Wei and Kautz, Jan and Birchfield, Stan},
|
| 116 |
+
booktitle={CVPR},
|
| 117 |
+
year={2024}
|
| 118 |
+
}
|
| 119 |
+
```
|
STATUS.md
ADDED
|
@@ -0,0 +1,383 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# FoundationPose Deployment - Current Status
|
| 2 |
+
|
| 3 |
+
**Last Updated:** 2026-01-28
|
| 4 |
+
**Repository:** `/Users/georgpuschel/repos/robot-ml/foundationpose/`
|
| 5 |
+
**Hugging Face Space:** https://huggingface.co/spaces/gpue/foundationpose
|
| 6 |
+
|
| 7 |
+
---
|
| 8 |
+
|
| 9 |
+
## 📦 What's Been Completed
|
| 10 |
+
|
| 11 |
+
### Core Files Created
|
| 12 |
+
|
| 13 |
+
1. **app.py** (570 lines)
|
| 14 |
+
- Complete Gradio application with ZeroGPU integration
|
| 15 |
+
- Dual mode: Placeholder (default) and Real FoundationPose
|
| 16 |
+
- REST API endpoints: `/api/initialize` and `/api/estimate`
|
| 17 |
+
- Web UI with tabs for initialization, estimation, and API docs
|
| 18 |
+
- Environment variable `USE_REAL_MODEL` controls mode
|
| 19 |
+
|
| 20 |
+
2. **estimator.py** (350+ lines)
|
| 21 |
+
- FoundationPoseEstimator class wrapping the real FoundationPose API
|
| 22 |
+
- Methods: `register_object()`, `estimate_pose()`, `reset_tracking()`
|
| 23 |
+
- Handles camera intrinsics, depth images, and segmentation masks
|
| 24 |
+
- Quaternion and rotation matrix conversions
|
| 25 |
+
- Mesh loading and reconstruction (placeholder for BundleSDF)
|
| 26 |
+
|
| 27 |
+
3. **client.py** (200+ lines)
|
| 28 |
+
- Python client for calling the API from robot-ml
|
| 29 |
+
- FoundationPoseClient class with initialize() and estimate_pose()
|
| 30 |
+
- Image encoding/decoding utilities
|
| 31 |
+
- Example usage code
|
| 32 |
+
|
| 33 |
+
4. **requirements.txt**
|
| 34 |
+
- Core dependencies: gradio, spaces, torch, opencv-python
|
| 35 |
+
- 3D vision: trimesh, pyrender, scikit-image
|
| 36 |
+
- Placeholder for FoundationPose installation
|
| 37 |
+
|
| 38 |
+
5. **Dockerfile**
|
| 39 |
+
- CUDA 12.1 base image
|
| 40 |
+
- System dependencies (eigen3, OpenGL, etc.)
|
| 41 |
+
- FoundationPose repository clone and build
|
| 42 |
+
- NVDiffRast, Kaolin, PyTorch3D installation
|
| 43 |
+
|
| 44 |
+
6. **download_weights.py**
|
| 45 |
+
- Script to check for and download model weights
|
| 46 |
+
- Instructions for manual weight setup
|
| 47 |
+
- Git-LFS integration guide
|
| 48 |
+
|
| 49 |
+
7. **deploy.sh**
|
| 50 |
+
- Interactive deployment script
|
| 51 |
+
- Checks for weights and git status
|
| 52 |
+
- Offers placeholder vs real mode deployment
|
| 53 |
+
- Guides through git-lfs setup
|
| 54 |
+
|
| 55 |
+
8. **Documentation**
|
| 56 |
+
- README.md (updated with full details)
|
| 57 |
+
- DEPLOYMENT.md (step-by-step deployment guide)
|
| 58 |
+
- QUICKSTART.md (quick start for both modes)
|
| 59 |
+
- STATUS.md (this file)
|
| 60 |
+
|
| 61 |
+
9. **.gitignore**
|
| 62 |
+
- Python cache files
|
| 63 |
+
- Virtual environments
|
| 64 |
+
- Model weights (for git-lfs)
|
| 65 |
+
- Test images
|
| 66 |
+
|
| 67 |
+
---
|
| 68 |
+
|
| 69 |
+
## 🎯 How It Works
|
| 70 |
+
|
| 71 |
+
### Placeholder Mode (Default)
|
| 72 |
+
|
| 73 |
+
- **Purpose**: API testing without GPU requirements
|
| 74 |
+
- **Behavior**: Returns empty pose results with success=true
|
| 75 |
+
- **Use Cases**:
|
| 76 |
+
- Developing client integrations
|
| 77 |
+
- Testing API structure
|
| 78 |
+
- Demos without GPU costs
|
| 79 |
+
|
| 80 |
+
### Real Mode (Requires Setup)
|
| 81 |
+
|
| 82 |
+
- **Purpose**: Actual 6D pose estimation
|
| 83 |
+
- **Requirements**:
|
| 84 |
+
- Model weights in `weights/` directory
|
| 85 |
+
- FoundationPose repository cloned
|
| 86 |
+
- Environment variable `USE_REAL_MODEL=true`
|
| 87 |
+
- **Behavior**: Uses actual FoundationPose inference
|
| 88 |
+
- **Use Cases**:
|
| 89 |
+
- Production pose estimation
|
| 90 |
+
- Validation and testing with real data
|
| 91 |
+
- Integration with robot-ml training
|
| 92 |
+
|
| 93 |
+
---
|
| 94 |
+
|
| 95 |
+
## 🚀 Deployment Options
|
| 96 |
+
|
| 97 |
+
### Option 1: Test Locally (Placeholder)
|
| 98 |
+
|
| 99 |
+
```bash
|
| 100 |
+
cd foundationpose
|
| 101 |
+
pip install -r requirements.txt
|
| 102 |
+
python app.py
|
| 103 |
+
# Visit http://localhost:7860
|
| 104 |
+
```
|
| 105 |
+
|
| 106 |
+
### Option 2: Deploy to HF (Placeholder)
|
| 107 |
+
|
| 108 |
+
```bash
|
| 109 |
+
./deploy.sh
|
| 110 |
+
# Or manually:
|
| 111 |
+
git add .
|
| 112 |
+
git commit -m "Deploy FoundationPose Space"
|
| 113 |
+
git push origin main
|
| 114 |
+
```
|
| 115 |
+
|
| 116 |
+
### Option 3: Deploy to HF (Real Mode)
|
| 117 |
+
|
| 118 |
+
**Requirements:**
|
| 119 |
+
1. Download weights from Google Drive
|
| 120 |
+
2. Set up git-lfs
|
| 121 |
+
3. Enable USE_REAL_MODEL=true
|
| 122 |
+
|
| 123 |
+
**Steps:**
|
| 124 |
+
```bash
|
| 125 |
+
# 1. Download weights (manual step)
|
| 126 |
+
# Visit: https://drive.google.com/drive/folders/1GCyGE-LbFGgRC-FuGsF3a1zeBuzsQ1Da
|
| 127 |
+
# Extract to: weights/2023-10-28-18-33-37/ and weights/2024-01-11-20-02-45/
|
| 128 |
+
|
| 129 |
+
# 2. Set up git-lfs
|
| 130 |
+
git lfs install
|
| 131 |
+
git lfs track "weights/**"
|
| 132 |
+
git add .gitattributes
|
| 133 |
+
|
| 134 |
+
# 3. Add weights
|
| 135 |
+
git add weights/
|
| 136 |
+
git commit -m "Add model weights"
|
| 137 |
+
|
| 138 |
+
# 4. Deploy
|
| 139 |
+
git push origin main
|
| 140 |
+
|
| 141 |
+
# 5. Set Space secret: USE_REAL_MODEL=true
|
| 142 |
+
```
|
| 143 |
+
|
| 144 |
+
---
|
| 145 |
+
|
| 146 |
+
## 🔗 Integration with robot-ml
|
| 147 |
+
|
| 148 |
+
### Update FoundationPose Wrapper
|
| 149 |
+
|
| 150 |
+
Edit `/training/nova_sim_trainer/perception/foundation_pose_wrapper.py`:
|
| 151 |
+
|
| 152 |
+
```python
|
| 153 |
+
from foundationpose.client import FoundationPoseClient
|
| 154 |
+
from pathlib import Path
|
| 155 |
+
import cv2
|
| 156 |
+
|
| 157 |
+
class FoundationPoseWrapper(PoseEstimator):
|
| 158 |
+
def __init__(self, api_url: str, tracked_objects: List[Dict], **kwargs):
|
| 159 |
+
super().__init__()
|
| 160 |
+
self.client = FoundationPoseClient(api_url)
|
| 161 |
+
self.object_ids = []
|
| 162 |
+
|
| 163 |
+
# Initialize each tracked object
|
| 164 |
+
for obj_config in tracked_objects:
|
| 165 |
+
if not obj_config.get("enabled", True):
|
| 166 |
+
continue
|
| 167 |
+
|
| 168 |
+
object_id = obj_config["object_id"]
|
| 169 |
+
ref_dir = Path(obj_config["reference_images_dir"])
|
| 170 |
+
|
| 171 |
+
# Load reference images
|
| 172 |
+
ref_images = []
|
| 173 |
+
for img_path in sorted(ref_dir.glob("*.jpg")):
|
| 174 |
+
img = cv2.imread(str(img_path))
|
| 175 |
+
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
| 176 |
+
ref_images.append(img)
|
| 177 |
+
|
| 178 |
+
# Register object with API
|
| 179 |
+
logger.info(f"Registering {object_id} with {len(ref_images)} images...")
|
| 180 |
+
self.client.initialize(object_id, ref_images)
|
| 181 |
+
self.object_ids.append(object_id)
|
| 182 |
+
|
| 183 |
+
def estimate_poses(self, frame, camera_intrinsics, scene_objects):
|
| 184 |
+
detected_poses = []
|
| 185 |
+
|
| 186 |
+
for object_id in self.object_ids:
|
| 187 |
+
poses = self.client.estimate_pose(object_id, frame, camera_intrinsics)
|
| 188 |
+
|
| 189 |
+
for pose in poses:
|
| 190 |
+
detected_poses.append(DetectedPose(
|
| 191 |
+
object_id=pose["object_id"],
|
| 192 |
+
position=pose["position"],
|
| 193 |
+
orientation=pose["orientation"],
|
| 194 |
+
confidence=pose["confidence"],
|
| 195 |
+
timestamp=0.0,
|
| 196 |
+
dimensions=tuple(pose.get("dimensions", [0.1, 0.1, 0.1]))
|
| 197 |
+
))
|
| 198 |
+
|
| 199 |
+
return detected_poses
|
| 200 |
+
```
|
| 201 |
+
|
| 202 |
+
### Update Configuration
|
| 203 |
+
|
| 204 |
+
Edit `/training/observations.yaml`:
|
| 205 |
+
|
| 206 |
+
```yaml
|
| 207 |
+
perception:
|
| 208 |
+
enabled: true
|
| 209 |
+
model: foundation_pose
|
| 210 |
+
api_url: https://gpue-foundationpose.hf.space # Your deployed Space URL
|
| 211 |
+
camera: aux_top
|
| 212 |
+
inference_fps: 5
|
| 213 |
+
|
| 214 |
+
tracked_objects:
|
| 215 |
+
- object_id: target_cube
|
| 216 |
+
scene_object_name: t_object
|
| 217 |
+
enabled: true
|
| 218 |
+
reference_images_dir: ./perception/reference/target_cube
|
| 219 |
+
dimensions: [0.1, 0.1, 0.1]
|
| 220 |
+
```
|
| 221 |
+
|
| 222 |
+
---
|
| 223 |
+
|
| 224 |
+
## ⚠️ What's NOT Done Yet
|
| 225 |
+
|
| 226 |
+
### Missing Pieces
|
| 227 |
+
|
| 228 |
+
1. **Model Weights**
|
| 229 |
+
- Not included in repo (too large)
|
| 230 |
+
- Must be downloaded manually
|
| 231 |
+
- Requires git-lfs setup
|
| 232 |
+
|
| 233 |
+
2. **FoundationPose Repository**
|
| 234 |
+
- Not included (git submodule or clone needed)
|
| 235 |
+
- C++ extensions need to be built
|
| 236 |
+
- Tested only with placeholder code
|
| 237 |
+
|
| 238 |
+
3. **BundleSDF Integration**
|
| 239 |
+
- Mesh reconstruction not implemented
|
| 240 |
+
- Currently uses placeholder cube mesh
|
| 241 |
+
- Needed for model-free mode
|
| 242 |
+
|
| 243 |
+
4. **Segmentation**
|
| 244 |
+
- Object segmentation uses placeholder mask
|
| 245 |
+
- Should integrate SAM (Segment Anything Model)
|
| 246 |
+
- Optional but improves accuracy
|
| 247 |
+
|
| 248 |
+
5. **Pose Visualization**
|
| 249 |
+
- Estimation results don't show overlays yet
|
| 250 |
+
- TODO: Render detected pose on query image
|
| 251 |
+
- Would improve debugging
|
| 252 |
+
|
| 253 |
+
### Testing Status
|
| 254 |
+
|
| 255 |
+
- ✅ Placeholder mode tested locally
|
| 256 |
+
- ⚠️ Real mode NOT tested (no weights)
|
| 257 |
+
- ⚠️ API integration NOT tested end-to-end
|
| 258 |
+
- ⚠️ ZeroGPU behavior unknown (not deployed yet)
|
| 259 |
+
|
| 260 |
+
---
|
| 261 |
+
|
| 262 |
+
## 📊 Performance Expectations
|
| 263 |
+
|
| 264 |
+
### ZeroGPU Characteristics
|
| 265 |
+
|
| 266 |
+
- **Cold Start**: 15-30 seconds (GPU allocation + model loading)
|
| 267 |
+
- **Warm Inference**: 0.5-2 seconds per query
|
| 268 |
+
- **Free Tier**: Limited monthly usage
|
| 269 |
+
- **Timeout**: GPU allocation lasts for duration specified in decorator
|
| 270 |
+
|
| 271 |
+
### robot-ml Training Integration
|
| 272 |
+
|
| 273 |
+
**Not suitable for real-time training loop (30 Hz):**
|
| 274 |
+
- 5 Hz perception requires 200ms per frame
|
| 275 |
+
- ZeroGPU latency: 500ms-2s warm, 15-30s cold
|
| 276 |
+
- ❌ Too slow for synchronous training
|
| 277 |
+
|
| 278 |
+
**Suitable for:**
|
| 279 |
+
- ✅ Batch processing recorded episodes
|
| 280 |
+
- ✅ Validation and testing
|
| 281 |
+
- ✅ Demos and visualization
|
| 282 |
+
- ✅ Reference data collection
|
| 283 |
+
|
| 284 |
+
**Recommendation:**
|
| 285 |
+
- Use dummy estimator during training (reads ground truth from sim)
|
| 286 |
+
- Use FoundationPose API for validation/testing only
|
| 287 |
+
- Consider local GPU deployment for production
|
| 288 |
+
|
| 289 |
+
---
|
| 290 |
+
|
| 291 |
+
## 📝 Next Steps
|
| 292 |
+
|
| 293 |
+
### Immediate (Before Deployment)
|
| 294 |
+
|
| 295 |
+
1. **Test Locally**
|
| 296 |
+
```bash
|
| 297 |
+
cd foundationpose
|
| 298 |
+
python app.py
|
| 299 |
+
# Test UI at http://localhost:7860
|
| 300 |
+
```
|
| 301 |
+
|
| 302 |
+
2. **Deploy Placeholder Mode**
|
| 303 |
+
```bash
|
| 304 |
+
./deploy.sh
|
| 305 |
+
# Choose "N" for real mode
|
| 306 |
+
```
|
| 307 |
+
|
| 308 |
+
3. **Verify Space Works**
|
| 309 |
+
- Visit https://huggingface.co/spaces/gpue/foundationpose
|
| 310 |
+
- Test initialization with test images
|
| 311 |
+
- Check logs for errors
|
| 312 |
+
|
| 313 |
+
### Short Term (With Weights)
|
| 314 |
+
|
| 315 |
+
1. **Download Model Weights**
|
| 316 |
+
- Get from Google Drive (see DEPLOYMENT.md)
|
| 317 |
+
- Extract to `weights/` directory
|
| 318 |
+
|
| 319 |
+
2. **Test Real Mode Locally**
|
| 320 |
+
```bash
|
| 321 |
+
export USE_REAL_MODEL=true
|
| 322 |
+
python app.py
|
| 323 |
+
# Upload real reference images
|
| 324 |
+
# Test pose estimation
|
| 325 |
+
```
|
| 326 |
+
|
| 327 |
+
3. **Deploy Real Mode**
|
| 328 |
+
```bash
|
| 329 |
+
# Set up git-lfs
|
| 330 |
+
git lfs track "weights/**"
|
| 331 |
+
git add .gitattributes weights/
|
| 332 |
+
git commit -m "Add model weights"
|
| 333 |
+
git push
|
| 334 |
+
|
| 335 |
+
# Set Space secret: USE_REAL_MODEL=true
|
| 336 |
+
```
|
| 337 |
+
|
| 338 |
+
### Long Term (Production)
|
| 339 |
+
|
| 340 |
+
1. **Optimize Performance**
|
| 341 |
+
- Implement batch inference
|
| 342 |
+
- Add caching for frequently used objects
|
| 343 |
+
- Tune GPU duration parameters
|
| 344 |
+
|
| 345 |
+
2. **Improve Accuracy**
|
| 346 |
+
- Integrate SAM for segmentation
|
| 347 |
+
- Add depth image support
|
| 348 |
+
- Implement BundleSDF reconstruction
|
| 349 |
+
|
| 350 |
+
3. **Production Deployment**
|
| 351 |
+
- Consider dedicated GPU (RunPod, Modal, etc.)
|
| 352 |
+
- Set up monitoring and logging
|
| 353 |
+
- Implement retry logic and error handling
|
| 354 |
+
|
| 355 |
+
---
|
| 356 |
+
|
| 357 |
+
## 📚 Reference Links
|
| 358 |
+
|
| 359 |
+
- **FoundationPose GitHub**: https://github.com/NVlabs/FoundationPose
|
| 360 |
+
- **Research Paper**: https://arxiv.org/abs/2312.08344
|
| 361 |
+
- **Model Weights**: https://drive.google.com/drive/folders/1GCyGE-LbFGgRC-FuGsF3a1zeBuzsQ1Da
|
| 362 |
+
- **HF Spaces Docs**: https://huggingface.co/docs/hub/spaces
|
| 363 |
+
- **ZeroGPU Docs**: https://huggingface.co/docs/hub/spaces-gpus-zerogpu
|
| 364 |
+
- **Gradio Docs**: https://www.gradio.app/docs
|
| 365 |
+
|
| 366 |
+
---
|
| 367 |
+
|
| 368 |
+
## 🤝 Citation
|
| 369 |
+
|
| 370 |
+
If you use this in your work:
|
| 371 |
+
|
| 372 |
+
```bibtex
|
| 373 |
+
@inproceedings{wen2023foundationpose,
|
| 374 |
+
title={FoundationPose: Unified 6D Pose Estimation and Tracking of Novel Objects},
|
| 375 |
+
author={Wen, Bowen and Yang, Wei and Kautz, Jan and Birchfield, Stan},
|
| 376 |
+
booktitle={CVPR},
|
| 377 |
+
year={2024}
|
| 378 |
+
}
|
| 379 |
+
```
|
| 380 |
+
|
| 381 |
+
---
|
| 382 |
+
|
| 383 |
+
**Summary:** The FoundationPose Space is fully set up and ready to deploy. It defaults to placeholder mode (no GPU needed) for testing the API structure. To enable real pose estimation, you need to manually download the model weights and set USE_REAL_MODEL=true. The integration code for robot-ml is ready but untested without the actual weights.
|
app.py
ADDED
|
@@ -0,0 +1,569 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
FoundationPose Inference Server with ZeroGPU Support
|
| 3 |
+
|
| 4 |
+
This Gradio app provides an API for 6D object pose estimation using FoundationPose.
|
| 5 |
+
It's designed to be called from the robot-ml training pipeline via HTTP requests.
|
| 6 |
+
|
| 7 |
+
API Endpoints:
|
| 8 |
+
- /api/initialize: Set up tracking for an object with reference images
|
| 9 |
+
- /api/estimate: Estimate 6D pose from a query image
|
| 10 |
+
"""
|
| 11 |
+
|
| 12 |
+
import base64
|
| 13 |
+
import io
|
| 14 |
+
import json
|
| 15 |
+
import logging
|
| 16 |
+
import os
|
| 17 |
+
from pathlib import Path
|
| 18 |
+
from typing import Dict, List, Optional
|
| 19 |
+
|
| 20 |
+
import cv2
|
| 21 |
+
import gradio as gr
|
| 22 |
+
import numpy as np
|
| 23 |
+
import spaces
|
| 24 |
+
import torch
|
| 25 |
+
from PIL import Image
|
| 26 |
+
|
| 27 |
+
logging.basicConfig(
|
| 28 |
+
level=logging.INFO,
|
| 29 |
+
format="[%(asctime)s] %(levelname)s: %(message)s"
|
| 30 |
+
)
|
| 31 |
+
logger = logging.getLogger(__name__)
|
| 32 |
+
|
| 33 |
+
# Check if running in real FoundationPose mode or placeholder mode
|
| 34 |
+
USE_REAL_MODEL = os.environ.get("USE_REAL_MODEL", "false").lower() == "true"
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
class FoundationPoseInference:
|
| 38 |
+
"""Wrapper for FoundationPose model inference."""
|
| 39 |
+
|
| 40 |
+
def __init__(self):
|
| 41 |
+
self.model = None
|
| 42 |
+
self.device = None
|
| 43 |
+
self.initialized = False
|
| 44 |
+
self.tracked_objects = {}
|
| 45 |
+
self.use_real_model = USE_REAL_MODEL
|
| 46 |
+
|
| 47 |
+
@spaces.GPU(duration=120) # Allocate GPU for 120 seconds (includes model loading)
|
| 48 |
+
def initialize_model(self):
|
| 49 |
+
"""Initialize the FoundationPose model on GPU."""
|
| 50 |
+
if self.initialized:
|
| 51 |
+
logger.info("Model already initialized")
|
| 52 |
+
return
|
| 53 |
+
|
| 54 |
+
logger.info("Initializing FoundationPose model...")
|
| 55 |
+
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 56 |
+
logger.info(f"Using device: {self.device}")
|
| 57 |
+
|
| 58 |
+
if self.use_real_model:
|
| 59 |
+
try:
|
| 60 |
+
logger.info("Loading real FoundationPose model...")
|
| 61 |
+
from estimator import FoundationPoseEstimator
|
| 62 |
+
|
| 63 |
+
self.model = FoundationPoseEstimator(
|
| 64 |
+
device=str(self.device),
|
| 65 |
+
weights_dir="weights"
|
| 66 |
+
)
|
| 67 |
+
logger.info("✓ Real FoundationPose model initialized successfully")
|
| 68 |
+
|
| 69 |
+
except Exception as e:
|
| 70 |
+
logger.error(f"Failed to initialize real model: {e}", exc_info=True)
|
| 71 |
+
logger.warning("Falling back to placeholder mode")
|
| 72 |
+
self.use_real_model = False
|
| 73 |
+
self.model = None
|
| 74 |
+
else:
|
| 75 |
+
logger.info("Using placeholder mode (set USE_REAL_MODEL=true for real inference)")
|
| 76 |
+
self.model = None
|
| 77 |
+
|
| 78 |
+
self.initialized = True
|
| 79 |
+
logger.info("FoundationPose inference ready")
|
| 80 |
+
|
| 81 |
+
def register_object(
|
| 82 |
+
self,
|
| 83 |
+
object_id: str,
|
| 84 |
+
reference_images: List[np.ndarray],
|
| 85 |
+
camera_intrinsics: Optional[Dict] = None,
|
| 86 |
+
mesh_path: Optional[str] = None
|
| 87 |
+
) -> bool:
|
| 88 |
+
"""Register an object for tracking with reference images.
|
| 89 |
+
|
| 90 |
+
Args:
|
| 91 |
+
object_id: Unique identifier for the object
|
| 92 |
+
reference_images: List of RGB images (numpy arrays) showing the object from different angles
|
| 93 |
+
camera_intrinsics: Camera parameters (fx, fy, cx, cy)
|
| 94 |
+
mesh_path: Optional path to CAD mesh file
|
| 95 |
+
|
| 96 |
+
Returns:
|
| 97 |
+
True if registration successful
|
| 98 |
+
"""
|
| 99 |
+
if not self.initialized:
|
| 100 |
+
self.initialize_model()
|
| 101 |
+
|
| 102 |
+
logger.info(f"Registering object '{object_id}' with {len(reference_images)} reference images")
|
| 103 |
+
|
| 104 |
+
if self.use_real_model and self.model is not None:
|
| 105 |
+
# Use real FoundationPose model
|
| 106 |
+
try:
|
| 107 |
+
success = self.model.register_object(
|
| 108 |
+
object_id=object_id,
|
| 109 |
+
reference_images=reference_images,
|
| 110 |
+
camera_intrinsics=camera_intrinsics,
|
| 111 |
+
mesh_path=mesh_path
|
| 112 |
+
)
|
| 113 |
+
if success:
|
| 114 |
+
self.tracked_objects[object_id] = {
|
| 115 |
+
"num_references": len(reference_images),
|
| 116 |
+
"camera_intrinsics": camera_intrinsics,
|
| 117 |
+
"mesh_path": mesh_path
|
| 118 |
+
}
|
| 119 |
+
return success
|
| 120 |
+
except Exception as e:
|
| 121 |
+
logger.error(f"Registration failed: {e}", exc_info=True)
|
| 122 |
+
return False
|
| 123 |
+
else:
|
| 124 |
+
# Placeholder mode
|
| 125 |
+
self.tracked_objects[object_id] = {
|
| 126 |
+
"num_references": len(reference_images),
|
| 127 |
+
"camera_intrinsics": camera_intrinsics,
|
| 128 |
+
"mesh_path": mesh_path
|
| 129 |
+
}
|
| 130 |
+
logger.info(f"✓ Object '{object_id}' registered (placeholder mode)")
|
| 131 |
+
return True
|
| 132 |
+
|
| 133 |
+
@spaces.GPU(duration=10) # Allocate GPU for 10 seconds per inference
|
| 134 |
+
def estimate_pose(
|
| 135 |
+
self,
|
| 136 |
+
object_id: str,
|
| 137 |
+
query_image: np.ndarray,
|
| 138 |
+
camera_intrinsics: Optional[Dict] = None,
|
| 139 |
+
depth_image: Optional[np.ndarray] = None,
|
| 140 |
+
mask: Optional[np.ndarray] = None
|
| 141 |
+
) -> Dict:
|
| 142 |
+
"""Estimate 6D pose of an object in a query image.
|
| 143 |
+
|
| 144 |
+
Args:
|
| 145 |
+
object_id: ID of object to detect
|
| 146 |
+
query_image: RGB query image as numpy array
|
| 147 |
+
camera_intrinsics: Optional camera parameters
|
| 148 |
+
depth_image: Optional depth map
|
| 149 |
+
mask: Optional object segmentation mask
|
| 150 |
+
|
| 151 |
+
Returns:
|
| 152 |
+
Dictionary with pose estimation results:
|
| 153 |
+
{
|
| 154 |
+
"success": bool,
|
| 155 |
+
"poses": [
|
| 156 |
+
{
|
| 157 |
+
"object_id": str,
|
| 158 |
+
"position": {"x": float, "y": float, "z": float},
|
| 159 |
+
"orientation": {"w": float, "x": float, "y": float, "z": float},
|
| 160 |
+
"confidence": float,
|
| 161 |
+
"dimensions": [float, float, float]
|
| 162 |
+
}
|
| 163 |
+
]
|
| 164 |
+
}
|
| 165 |
+
"""
|
| 166 |
+
if not self.initialized:
|
| 167 |
+
return {"success": False, "error": "Model not initialized"}
|
| 168 |
+
|
| 169 |
+
if object_id not in self.tracked_objects:
|
| 170 |
+
return {"success": False, "error": f"Object '{object_id}' not registered"}
|
| 171 |
+
|
| 172 |
+
logger.info(f"Estimating pose for object '{object_id}'")
|
| 173 |
+
|
| 174 |
+
if self.use_real_model and self.model is not None:
|
| 175 |
+
# Use real FoundationPose model
|
| 176 |
+
try:
|
| 177 |
+
pose_result = self.model.estimate_pose(
|
| 178 |
+
object_id=object_id,
|
| 179 |
+
rgb_image=query_image,
|
| 180 |
+
depth_image=depth_image,
|
| 181 |
+
mask=mask,
|
| 182 |
+
camera_intrinsics=camera_intrinsics
|
| 183 |
+
)
|
| 184 |
+
|
| 185 |
+
if pose_result is None:
|
| 186 |
+
return {
|
| 187 |
+
"success": False,
|
| 188 |
+
"error": "Pose estimation returned None",
|
| 189 |
+
"poses": []
|
| 190 |
+
}
|
| 191 |
+
|
| 192 |
+
return {
|
| 193 |
+
"success": True,
|
| 194 |
+
"poses": [pose_result]
|
| 195 |
+
}
|
| 196 |
+
|
| 197 |
+
except Exception as e:
|
| 198 |
+
logger.error(f"Pose estimation error: {e}", exc_info=True)
|
| 199 |
+
return {"success": False, "error": str(e), "poses": []}
|
| 200 |
+
else:
|
| 201 |
+
# Placeholder mode - return empty poses
|
| 202 |
+
logger.info("Placeholder mode: returning empty pose result")
|
| 203 |
+
return {
|
| 204 |
+
"success": True,
|
| 205 |
+
"poses": [],
|
| 206 |
+
"note": "Placeholder mode - set USE_REAL_MODEL=true for real inference"
|
| 207 |
+
}
|
| 208 |
+
|
| 209 |
+
|
| 210 |
+
# Global model instance
|
| 211 |
+
pose_estimator = FoundationPoseInference()
|
| 212 |
+
|
| 213 |
+
|
| 214 |
+
def initialize_api(request: gr.Request) -> Dict:
|
| 215 |
+
"""API endpoint for initializing object tracking.
|
| 216 |
+
|
| 217 |
+
Request body:
|
| 218 |
+
{
|
| 219 |
+
"object_id": str,
|
| 220 |
+
"reference_images_b64": [str, ...],
|
| 221 |
+
"camera_intrinsics": str (JSON),
|
| 222 |
+
"mesh_path": str (optional)
|
| 223 |
+
}
|
| 224 |
+
|
| 225 |
+
Returns:
|
| 226 |
+
{"success": bool, "message": str}
|
| 227 |
+
"""
|
| 228 |
+
try:
|
| 229 |
+
data = request.json() if hasattr(request, 'json') else {}
|
| 230 |
+
|
| 231 |
+
object_id = data.get("object_id")
|
| 232 |
+
reference_images_b64 = data.get("reference_images_b64", [])
|
| 233 |
+
camera_intrinsics_str = data.get("camera_intrinsics")
|
| 234 |
+
mesh_path = data.get("mesh_path")
|
| 235 |
+
|
| 236 |
+
if not object_id:
|
| 237 |
+
return {"success": False, "error": "Missing object_id"}
|
| 238 |
+
|
| 239 |
+
if not reference_images_b64:
|
| 240 |
+
return {"success": False, "error": "Missing reference_images_b64"}
|
| 241 |
+
|
| 242 |
+
# Decode reference images
|
| 243 |
+
reference_images = []
|
| 244 |
+
for img_b64 in reference_images_b64:
|
| 245 |
+
img_bytes = base64.b64decode(img_b64)
|
| 246 |
+
img_array = np.frombuffer(img_bytes, dtype=np.uint8)
|
| 247 |
+
img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
|
| 248 |
+
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
| 249 |
+
reference_images.append(img)
|
| 250 |
+
|
| 251 |
+
# Parse camera intrinsics
|
| 252 |
+
intrinsics = json.loads(camera_intrinsics_str) if camera_intrinsics_str else None
|
| 253 |
+
|
| 254 |
+
# Register object
|
| 255 |
+
success = pose_estimator.register_object(
|
| 256 |
+
object_id=object_id,
|
| 257 |
+
reference_images=reference_images,
|
| 258 |
+
camera_intrinsics=intrinsics,
|
| 259 |
+
mesh_path=mesh_path
|
| 260 |
+
)
|
| 261 |
+
|
| 262 |
+
return {
|
| 263 |
+
"success": success,
|
| 264 |
+
"message": f"Object '{object_id}' registered with {len(reference_images)} reference images"
|
| 265 |
+
}
|
| 266 |
+
|
| 267 |
+
except Exception as e:
|
| 268 |
+
logger.error(f"Initialization error: {e}", exc_info=True)
|
| 269 |
+
return {"success": False, "error": str(e)}
|
| 270 |
+
|
| 271 |
+
|
| 272 |
+
def estimate_api(request: gr.Request) -> Dict:
|
| 273 |
+
"""API endpoint for pose estimation.
|
| 274 |
+
|
| 275 |
+
Request body:
|
| 276 |
+
{
|
| 277 |
+
"object_id": str,
|
| 278 |
+
"query_image_b64": str,
|
| 279 |
+
"camera_intrinsics": str (JSON),
|
| 280 |
+
"depth_image_b64": str (optional),
|
| 281 |
+
"mask_b64": str (optional)
|
| 282 |
+
}
|
| 283 |
+
|
| 284 |
+
Returns:
|
| 285 |
+
Pose estimation results
|
| 286 |
+
"""
|
| 287 |
+
try:
|
| 288 |
+
data = request.json() if hasattr(request, 'json') else {}
|
| 289 |
+
|
| 290 |
+
object_id = data.get("object_id")
|
| 291 |
+
query_image_b64 = data.get("query_image_b64")
|
| 292 |
+
camera_intrinsics_str = data.get("camera_intrinsics")
|
| 293 |
+
depth_image_b64 = data.get("depth_image_b64")
|
| 294 |
+
mask_b64 = data.get("mask_b64")
|
| 295 |
+
|
| 296 |
+
if not object_id:
|
| 297 |
+
return {"success": False, "error": "Missing object_id"}
|
| 298 |
+
|
| 299 |
+
if not query_image_b64:
|
| 300 |
+
return {"success": False, "error": "Missing query_image_b64"}
|
| 301 |
+
|
| 302 |
+
# Decode query image
|
| 303 |
+
img_bytes = base64.b64decode(query_image_b64)
|
| 304 |
+
img_array = np.frombuffer(img_bytes, dtype=np.uint8)
|
| 305 |
+
img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
|
| 306 |
+
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
| 307 |
+
|
| 308 |
+
# Decode optional depth image
|
| 309 |
+
depth = None
|
| 310 |
+
if depth_image_b64:
|
| 311 |
+
depth_bytes = base64.b64decode(depth_image_b64)
|
| 312 |
+
depth = np.frombuffer(depth_bytes, dtype=np.float32)
|
| 313 |
+
|
| 314 |
+
# Decode optional mask
|
| 315 |
+
mask = None
|
| 316 |
+
if mask_b64:
|
| 317 |
+
mask_bytes = base64.b64decode(mask_b64)
|
| 318 |
+
mask_array = np.frombuffer(mask_bytes, dtype=np.uint8)
|
| 319 |
+
mask = cv2.imdecode(mask_array, cv2.IMREAD_GRAYSCALE)
|
| 320 |
+
|
| 321 |
+
# Parse camera intrinsics
|
| 322 |
+
intrinsics = json.loads(camera_intrinsics_str) if camera_intrinsics_str else None
|
| 323 |
+
|
| 324 |
+
# Estimate pose
|
| 325 |
+
result = pose_estimator.estimate_pose(
|
| 326 |
+
object_id=object_id,
|
| 327 |
+
query_image=img,
|
| 328 |
+
camera_intrinsics=intrinsics,
|
| 329 |
+
depth_image=depth,
|
| 330 |
+
mask=mask
|
| 331 |
+
)
|
| 332 |
+
|
| 333 |
+
return result
|
| 334 |
+
|
| 335 |
+
except Exception as e:
|
| 336 |
+
logger.error(f"Estimation error: {e}", exc_info=True)
|
| 337 |
+
return {"success": False, "error": str(e)}
|
| 338 |
+
|
| 339 |
+
|
| 340 |
+
# Gradio UI for testing
|
| 341 |
+
def test_initialization(object_id: str, reference_images: List):
|
| 342 |
+
"""Test UI for initialization."""
|
| 343 |
+
if not object_id:
|
| 344 |
+
return "❌ Please enter an object ID"
|
| 345 |
+
|
| 346 |
+
if not reference_images:
|
| 347 |
+
return "❌ Please upload reference images"
|
| 348 |
+
|
| 349 |
+
try:
|
| 350 |
+
# Convert PIL images to numpy arrays
|
| 351 |
+
ref_imgs = []
|
| 352 |
+
for img in reference_images:
|
| 353 |
+
ref_imgs.append(np.array(img))
|
| 354 |
+
|
| 355 |
+
success = pose_estimator.register_object(object_id, ref_imgs, None)
|
| 356 |
+
|
| 357 |
+
if success:
|
| 358 |
+
return f"✅ Object '{object_id}' registered with {len(ref_imgs)} images"
|
| 359 |
+
else:
|
| 360 |
+
return "❌ Registration failed"
|
| 361 |
+
|
| 362 |
+
except Exception as e:
|
| 363 |
+
logger.error(f"Test initialization error: {e}", exc_info=True)
|
| 364 |
+
return f"❌ Error: {str(e)}"
|
| 365 |
+
|
| 366 |
+
|
| 367 |
+
def test_estimation(object_id: str, query_image):
|
| 368 |
+
"""Test UI for pose estimation."""
|
| 369 |
+
if not object_id:
|
| 370 |
+
return "❌ Please enter an object ID", None
|
| 371 |
+
|
| 372 |
+
if query_image is None:
|
| 373 |
+
return "❌ Please upload a query image", None
|
| 374 |
+
|
| 375 |
+
try:
|
| 376 |
+
query_img = np.array(query_image)
|
| 377 |
+
result = pose_estimator.estimate_pose(object_id, query_img, None)
|
| 378 |
+
|
| 379 |
+
if result["success"]:
|
| 380 |
+
num_poses = len(result["poses"])
|
| 381 |
+
output_text = f"✅ Detection complete: {num_poses} pose(s) detected\n\n"
|
| 382 |
+
|
| 383 |
+
if num_poses == 0:
|
| 384 |
+
output_text += "Note: " + result.get("note", "No poses detected")
|
| 385 |
+
else:
|
| 386 |
+
for i, pose in enumerate(result["poses"]):
|
| 387 |
+
output_text += f"Pose {i+1}:\n"
|
| 388 |
+
output_text += f" Position: ({pose['position']['x']:.3f}, {pose['position']['y']:.3f}, {pose['position']['z']:.3f})\n"
|
| 389 |
+
output_text += f" Confidence: {pose['confidence']:.3f}\n\n"
|
| 390 |
+
|
| 391 |
+
# TODO: Visualize detected pose on image
|
| 392 |
+
output_image = query_image
|
| 393 |
+
|
| 394 |
+
return output_text, output_image
|
| 395 |
+
else:
|
| 396 |
+
return f"❌ Detection failed: {result.get('error', 'Unknown error')}", None
|
| 397 |
+
|
| 398 |
+
except Exception as e:
|
| 399 |
+
logger.error(f"Test estimation error: {e}", exc_info=True)
|
| 400 |
+
return f"❌ Error: {str(e)}", None
|
| 401 |
+
|
| 402 |
+
|
| 403 |
+
# Build Gradio interface
|
| 404 |
+
with gr.Blocks(title="FoundationPose Inference", theme=gr.themes.Soft()) as demo:
|
| 405 |
+
gr.Markdown("# 🎯 FoundationPose 6D Object Pose Estimation")
|
| 406 |
+
|
| 407 |
+
mode_indicator = gr.Markdown(
|
| 408 |
+
f"**Mode:** {'🟢 Real FoundationPose' if USE_REAL_MODEL else '🟡 Placeholder (set USE_REAL_MODEL=true)'}",
|
| 409 |
+
elem_id="mode"
|
| 410 |
+
)
|
| 411 |
+
|
| 412 |
+
gr.Markdown("""
|
| 413 |
+
This service provides 6D object pose estimation using FoundationPose.
|
| 414 |
+
|
| 415 |
+
**Usage:**
|
| 416 |
+
1. Register an object with reference images using the Initialize tab
|
| 417 |
+
2. Estimate poses in query images using the Estimate tab
|
| 418 |
+
|
| 419 |
+
**API Endpoints:**
|
| 420 |
+
- POST `/api/initialize` - Register object with reference images
|
| 421 |
+
- POST `/api/estimate` - Estimate 6D pose from query image
|
| 422 |
+
""")
|
| 423 |
+
|
| 424 |
+
with gr.Tab("🔧 Initialize Object"):
|
| 425 |
+
gr.Markdown("### Register an object for tracking")
|
| 426 |
+
with gr.Row():
|
| 427 |
+
with gr.Column():
|
| 428 |
+
init_object_id = gr.Textbox(
|
| 429 |
+
label="Object ID",
|
| 430 |
+
placeholder="e.g., target_cube",
|
| 431 |
+
info="Unique identifier for the object"
|
| 432 |
+
)
|
| 433 |
+
init_ref_images = gr.File(
|
| 434 |
+
label="Reference Images (16-20 recommended)",
|
| 435 |
+
file_count="multiple",
|
| 436 |
+
file_types=["image"],
|
| 437 |
+
type="filepath"
|
| 438 |
+
)
|
| 439 |
+
init_button = gr.Button("Register Object", variant="primary", size="lg")
|
| 440 |
+
with gr.Column():
|
| 441 |
+
init_output = gr.Textbox(label="Result", lines=8)
|
| 442 |
+
|
| 443 |
+
gr.Markdown("""
|
| 444 |
+
**Tips:**
|
| 445 |
+
- Capture 16-20 images from different viewpoints
|
| 446 |
+
- Include various angles and distances
|
| 447 |
+
- Ensure good lighting and sharp focus
|
| 448 |
+
""")
|
| 449 |
+
|
| 450 |
+
init_button.click(
|
| 451 |
+
fn=test_initialization,
|
| 452 |
+
inputs=[init_object_id, init_ref_images],
|
| 453 |
+
outputs=init_output
|
| 454 |
+
)
|
| 455 |
+
|
| 456 |
+
with gr.Tab("🔍 Estimate Pose"):
|
| 457 |
+
gr.Markdown("### Detect object pose in a query image")
|
| 458 |
+
with gr.Row():
|
| 459 |
+
with gr.Column():
|
| 460 |
+
est_object_id = gr.Textbox(
|
| 461 |
+
label="Object ID",
|
| 462 |
+
placeholder="e.g., target_cube",
|
| 463 |
+
info="Must match an initialized object"
|
| 464 |
+
)
|
| 465 |
+
est_query_image = gr.Image(
|
| 466 |
+
label="Query Image",
|
| 467 |
+
type="pil",
|
| 468 |
+
sources=["upload", "webcam"]
|
| 469 |
+
)
|
| 470 |
+
est_button = gr.Button("Estimate Pose", variant="primary", size="lg")
|
| 471 |
+
with gr.Column():
|
| 472 |
+
est_output_text = gr.Textbox(label="Detection Results", lines=15)
|
| 473 |
+
est_output_image = gr.Image(label="Visualization (coming soon)")
|
| 474 |
+
|
| 475 |
+
est_button.click(
|
| 476 |
+
fn=test_estimation,
|
| 477 |
+
inputs=[est_object_id, est_query_image],
|
| 478 |
+
outputs=[est_output_text, est_output_image]
|
| 479 |
+
)
|
| 480 |
+
|
| 481 |
+
with gr.Tab("📖 API Documentation"):
|
| 482 |
+
gr.Markdown("""
|
| 483 |
+
### HTTP API
|
| 484 |
+
|
| 485 |
+
#### Initialize Object
|
| 486 |
+
```bash
|
| 487 |
+
curl -X POST https://gpue-foundationpose.hf.space/api/initialize \\
|
| 488 |
+
-H "Content-Type: application/json" \\
|
| 489 |
+
-d '{
|
| 490 |
+
"object_id": "target_cube",
|
| 491 |
+
"reference_images_b64": ["<base64-encoded-jpeg>", ...],
|
| 492 |
+
"camera_intrinsics": "{\\"fx\\": 500, \\"fy\\": 500, \\"cx\\": 320, \\"cy\\": 240}"
|
| 493 |
+
}'
|
| 494 |
+
```
|
| 495 |
+
|
| 496 |
+
#### Estimate Pose
|
| 497 |
+
```bash
|
| 498 |
+
curl -X POST https://gpue-foundationpose.hf.space/api/estimate \\
|
| 499 |
+
-H "Content-Type: application/json" \\
|
| 500 |
+
-d '{
|
| 501 |
+
"object_id": "target_cube",
|
| 502 |
+
"query_image_b64": "<base64-encoded-jpeg>",
|
| 503 |
+
"camera_intrinsics": "{\\"fx\\": 500, \\"fy\\": 500, \\"cx\\": 320, \\"cy\\": 240}"
|
| 504 |
+
}'
|
| 505 |
+
```
|
| 506 |
+
|
| 507 |
+
**Response Format:**
|
| 508 |
+
```json
|
| 509 |
+
{
|
| 510 |
+
"success": true,
|
| 511 |
+
"poses": [
|
| 512 |
+
{
|
| 513 |
+
"object_id": "target_cube",
|
| 514 |
+
"position": {"x": 0.5, "y": 0.3, "z": 0.1},
|
| 515 |
+
"orientation": {"w": 1.0, "x": 0.0, "y": 0.0, "z": 0.0},
|
| 516 |
+
"confidence": 0.95,
|
| 517 |
+
"dimensions": [0.1, 0.1, 0.1]
|
| 518 |
+
}
|
| 519 |
+
]
|
| 520 |
+
}
|
| 521 |
+
```
|
| 522 |
+
|
| 523 |
+
### Integration with robot-ml
|
| 524 |
+
|
| 525 |
+
```python
|
| 526 |
+
from foundationpose.client import FoundationPoseClient
|
| 527 |
+
|
| 528 |
+
client = FoundationPoseClient("https://gpue-foundationpose.hf.space")
|
| 529 |
+
|
| 530 |
+
# Load reference images
|
| 531 |
+
ref_images = load_reference_images("./perception/reference/target_cube")
|
| 532 |
+
|
| 533 |
+
# Initialize object
|
| 534 |
+
client.initialize("target_cube", ref_images)
|
| 535 |
+
|
| 536 |
+
# Estimate pose
|
| 537 |
+
poses = client.estimate_pose("target_cube", query_image)
|
| 538 |
+
```
|
| 539 |
+
""")
|
| 540 |
+
|
| 541 |
+
gr.Markdown("""
|
| 542 |
+
---
|
| 543 |
+
**Citation:**
|
| 544 |
+
```bibtex
|
| 545 |
+
@inproceedings{wen2023foundationpose,
|
| 546 |
+
title={FoundationPose: Unified 6D Pose Estimation and Tracking of Novel Objects},
|
| 547 |
+
author={Wen, Bowen and Yang, Wei and Kautz, Jan and Birchfield, Stan},
|
| 548 |
+
booktitle={CVPR},
|
| 549 |
+
year={2024}
|
| 550 |
+
}
|
| 551 |
+
```
|
| 552 |
+
|
| 553 |
+
[GitHub](https://github.com/NVlabs/FoundationPose) | [Paper](https://arxiv.org/abs/2312.08344)
|
| 554 |
+
""")
|
| 555 |
+
|
| 556 |
+
|
| 557 |
+
# Launch app
|
| 558 |
+
if __name__ == "__main__":
|
| 559 |
+
logger.info("=" * 60)
|
| 560 |
+
logger.info("FoundationPose Inference Server Starting")
|
| 561 |
+
logger.info(f"Mode: {'Real Model' if USE_REAL_MODEL else 'Placeholder'}")
|
| 562 |
+
logger.info("=" * 60)
|
| 563 |
+
|
| 564 |
+
demo.launch(
|
| 565 |
+
server_name="0.0.0.0",
|
| 566 |
+
server_port=7860,
|
| 567 |
+
share=False,
|
| 568 |
+
show_api=True
|
| 569 |
+
)
|
client.py
ADDED
|
@@ -0,0 +1,212 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Client for FoundationPose Hugging Face Space API
|
| 3 |
+
|
| 4 |
+
This client can be used from the robot-ml training pipeline to call the
|
| 5 |
+
FoundationPose inference API hosted on Hugging Face Spaces.
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
import base64
|
| 9 |
+
import json
|
| 10 |
+
import logging
|
| 11 |
+
from io import BytesIO
|
| 12 |
+
from pathlib import Path
|
| 13 |
+
from typing import Dict, List, Optional
|
| 14 |
+
|
| 15 |
+
import cv2
|
| 16 |
+
import numpy as np
|
| 17 |
+
import requests
|
| 18 |
+
|
| 19 |
+
logger = logging.getLogger(__name__)
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
class FoundationPoseClient:
|
| 23 |
+
"""Client for FoundationPose API."""
|
| 24 |
+
|
| 25 |
+
def __init__(self, api_url: str = "https://gpue-foundationpose.hf.space"):
|
| 26 |
+
"""Initialize client.
|
| 27 |
+
|
| 28 |
+
Args:
|
| 29 |
+
api_url: Base URL of the FoundationPose Space
|
| 30 |
+
"""
|
| 31 |
+
self.api_url = api_url.rstrip("/")
|
| 32 |
+
self.session = requests.Session()
|
| 33 |
+
self.session.headers.update({"Content-Type": "application/json"})
|
| 34 |
+
|
| 35 |
+
def _encode_image(self, image: np.ndarray) -> str:
|
| 36 |
+
"""Encode image as base64 JPEG.
|
| 37 |
+
|
| 38 |
+
Args:
|
| 39 |
+
image: RGB image as numpy array
|
| 40 |
+
|
| 41 |
+
Returns:
|
| 42 |
+
Base64-encoded JPEG string
|
| 43 |
+
"""
|
| 44 |
+
# Convert RGB to BGR for OpenCV
|
| 45 |
+
image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
| 46 |
+
|
| 47 |
+
# Encode as JPEG
|
| 48 |
+
_, buffer = cv2.imencode(".jpg", image_bgr, [cv2.IMWRITE_JPEG_QUALITY, 85])
|
| 49 |
+
|
| 50 |
+
# Convert to base64
|
| 51 |
+
image_b64 = base64.b64encode(buffer).decode("utf-8")
|
| 52 |
+
return image_b64
|
| 53 |
+
|
| 54 |
+
def initialize(
|
| 55 |
+
self,
|
| 56 |
+
object_id: str,
|
| 57 |
+
reference_images: List[np.ndarray],
|
| 58 |
+
camera_intrinsics: Optional[Dict] = None
|
| 59 |
+
) -> bool:
|
| 60 |
+
"""Initialize object tracking with reference images.
|
| 61 |
+
|
| 62 |
+
Args:
|
| 63 |
+
object_id: Unique ID for the object
|
| 64 |
+
reference_images: List of RGB images (numpy arrays)
|
| 65 |
+
camera_intrinsics: Optional camera parameters
|
| 66 |
+
|
| 67 |
+
Returns:
|
| 68 |
+
True if successful
|
| 69 |
+
|
| 70 |
+
Raises:
|
| 71 |
+
RuntimeError: If initialization fails
|
| 72 |
+
"""
|
| 73 |
+
logger.info(f"Initializing object '{object_id}' with {len(reference_images)} reference images")
|
| 74 |
+
|
| 75 |
+
# Encode images
|
| 76 |
+
images_b64 = [self._encode_image(img) for img in reference_images]
|
| 77 |
+
|
| 78 |
+
# Prepare request
|
| 79 |
+
payload = {
|
| 80 |
+
"object_id": object_id,
|
| 81 |
+
"reference_images_b64": images_b64,
|
| 82 |
+
}
|
| 83 |
+
|
| 84 |
+
if camera_intrinsics:
|
| 85 |
+
payload["camera_intrinsics"] = json.dumps(camera_intrinsics)
|
| 86 |
+
|
| 87 |
+
# Send request
|
| 88 |
+
try:
|
| 89 |
+
response = self.session.post(
|
| 90 |
+
f"{self.api_url}/api/initialize",
|
| 91 |
+
json=payload,
|
| 92 |
+
timeout=120 # Long timeout for model loading
|
| 93 |
+
)
|
| 94 |
+
response.raise_for_status()
|
| 95 |
+
|
| 96 |
+
result = response.json()
|
| 97 |
+
|
| 98 |
+
if not result.get("success"):
|
| 99 |
+
error = result.get("error", "Unknown error")
|
| 100 |
+
raise RuntimeError(f"Initialization failed: {error}")
|
| 101 |
+
|
| 102 |
+
logger.info(f"Object '{object_id}' initialized successfully")
|
| 103 |
+
return True
|
| 104 |
+
|
| 105 |
+
except requests.exceptions.RequestException as e:
|
| 106 |
+
logger.error(f"API request failed: {e}")
|
| 107 |
+
raise RuntimeError(f"Failed to initialize object: {e}")
|
| 108 |
+
|
| 109 |
+
def estimate_pose(
|
| 110 |
+
self,
|
| 111 |
+
object_id: str,
|
| 112 |
+
query_image: np.ndarray,
|
| 113 |
+
camera_intrinsics: Optional[Dict] = None
|
| 114 |
+
) -> List[Dict]:
|
| 115 |
+
"""Estimate 6D pose of object in query image.
|
| 116 |
+
|
| 117 |
+
Args:
|
| 118 |
+
object_id: ID of object to detect
|
| 119 |
+
query_image: RGB query image as numpy array
|
| 120 |
+
camera_intrinsics: Optional camera parameters
|
| 121 |
+
|
| 122 |
+
Returns:
|
| 123 |
+
List of detected poses:
|
| 124 |
+
[
|
| 125 |
+
{
|
| 126 |
+
"object_id": str,
|
| 127 |
+
"position": {"x": float, "y": float, "z": float},
|
| 128 |
+
"orientation": {"w": float, "x": float, "y": float, "z": float},
|
| 129 |
+
"confidence": float,
|
| 130 |
+
"dimensions": [float, float, float]
|
| 131 |
+
}
|
| 132 |
+
]
|
| 133 |
+
|
| 134 |
+
Raises:
|
| 135 |
+
RuntimeError: If estimation fails
|
| 136 |
+
"""
|
| 137 |
+
# Encode image
|
| 138 |
+
image_b64 = self._encode_image(query_image)
|
| 139 |
+
|
| 140 |
+
# Prepare request
|
| 141 |
+
payload = {
|
| 142 |
+
"object_id": object_id,
|
| 143 |
+
"query_image_b64": image_b64,
|
| 144 |
+
}
|
| 145 |
+
|
| 146 |
+
if camera_intrinsics:
|
| 147 |
+
payload["camera_intrinsics"] = json.dumps(camera_intrinsics)
|
| 148 |
+
|
| 149 |
+
# Send request
|
| 150 |
+
try:
|
| 151 |
+
response = self.session.post(
|
| 152 |
+
f"{self.api_url}/api/estimate",
|
| 153 |
+
json=payload,
|
| 154 |
+
timeout=30
|
| 155 |
+
)
|
| 156 |
+
response.raise_for_status()
|
| 157 |
+
|
| 158 |
+
result = response.json()
|
| 159 |
+
|
| 160 |
+
if not result.get("success"):
|
| 161 |
+
error = result.get("error", "Unknown error")
|
| 162 |
+
raise RuntimeError(f"Pose estimation failed: {error}")
|
| 163 |
+
|
| 164 |
+
return result.get("poses", [])
|
| 165 |
+
|
| 166 |
+
except requests.exceptions.RequestException as e:
|
| 167 |
+
logger.error(f"API request failed: {e}")
|
| 168 |
+
raise RuntimeError(f"Failed to estimate pose: {e}")
|
| 169 |
+
|
| 170 |
+
|
| 171 |
+
def load_reference_images(directory: Path) -> List[np.ndarray]:
|
| 172 |
+
"""Load reference images from directory.
|
| 173 |
+
|
| 174 |
+
Args:
|
| 175 |
+
directory: Path to directory containing images
|
| 176 |
+
|
| 177 |
+
Returns:
|
| 178 |
+
List of RGB images as numpy arrays
|
| 179 |
+
"""
|
| 180 |
+
images = []
|
| 181 |
+
for img_path in sorted(directory.glob("*.jpg")):
|
| 182 |
+
img = cv2.imread(str(img_path))
|
| 183 |
+
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
| 184 |
+
images.append(img)
|
| 185 |
+
|
| 186 |
+
logger.info(f"Loaded {len(images)} reference images from {directory}")
|
| 187 |
+
return images
|
| 188 |
+
|
| 189 |
+
|
| 190 |
+
# Example usage
|
| 191 |
+
if __name__ == "__main__":
|
| 192 |
+
logging.basicConfig(level=logging.INFO)
|
| 193 |
+
|
| 194 |
+
# Initialize client
|
| 195 |
+
client = FoundationPoseClient()
|
| 196 |
+
|
| 197 |
+
# Load reference images
|
| 198 |
+
ref_dir = Path("../training/perception/reference/target_cube")
|
| 199 |
+
if ref_dir.exists():
|
| 200 |
+
ref_images = load_reference_images(ref_dir)
|
| 201 |
+
|
| 202 |
+
# Initialize object
|
| 203 |
+
client.initialize("target_cube", ref_images)
|
| 204 |
+
|
| 205 |
+
# Estimate pose on first reference image (for testing)
|
| 206 |
+
poses = client.estimate_pose("target_cube", ref_images[0])
|
| 207 |
+
print(f"Detected {len(poses)} poses:")
|
| 208 |
+
for pose in poses:
|
| 209 |
+
print(f" {pose}")
|
| 210 |
+
else:
|
| 211 |
+
print(f"Reference directory not found: {ref_dir}")
|
| 212 |
+
print("Run 'make capture-reference' to collect reference images first")
|
deploy.sh
ADDED
|
@@ -0,0 +1,108 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/bin/bash
|
| 2 |
+
# Deploy FoundationPose to Hugging Face Spaces
|
| 3 |
+
|
| 4 |
+
set -e
|
| 5 |
+
|
| 6 |
+
SPACE_URL="https://huggingface.co/spaces/gpue/foundationpose"
|
| 7 |
+
|
| 8 |
+
echo "=========================================="
|
| 9 |
+
echo "FoundationPose Hugging Face Deployment"
|
| 10 |
+
echo "=========================================="
|
| 11 |
+
echo ""
|
| 12 |
+
|
| 13 |
+
# Check if we're in the right directory
|
| 14 |
+
if [ ! -f "app.py" ]; then
|
| 15 |
+
echo "Error: Must run from foundationpose directory"
|
| 16 |
+
exit 1
|
| 17 |
+
fi
|
| 18 |
+
|
| 19 |
+
# Check git remote
|
| 20 |
+
if ! git remote get-url origin | grep -q "huggingface"; then
|
| 21 |
+
echo "Setting up Hugging Face remote..."
|
| 22 |
+
git remote add origin https://huggingface.co/spaces/gpue/foundationpose
|
| 23 |
+
else
|
| 24 |
+
echo "✓ Hugging Face remote configured"
|
| 25 |
+
fi
|
| 26 |
+
|
| 27 |
+
# Check for uncommitted changes
|
| 28 |
+
if [ -n "$(git status --porcelain)" ]; then
|
| 29 |
+
echo ""
|
| 30 |
+
echo "Uncommitted changes found. Commit them?"
|
| 31 |
+
echo ""
|
| 32 |
+
git status --short
|
| 33 |
+
echo ""
|
| 34 |
+
read -p "Commit all changes? (y/N) " -n 1 -r
|
| 35 |
+
echo
|
| 36 |
+
if [[ $REPLY =~ ^[Yy]$ ]]; then
|
| 37 |
+
read -p "Commit message: " commit_msg
|
| 38 |
+
git add .
|
| 39 |
+
git commit -m "$commit_msg"
|
| 40 |
+
else
|
| 41 |
+
echo "Deployment cancelled."
|
| 42 |
+
exit 0
|
| 43 |
+
fi
|
| 44 |
+
fi
|
| 45 |
+
|
| 46 |
+
# Check for model weights
|
| 47 |
+
echo ""
|
| 48 |
+
echo "Checking for model weights..."
|
| 49 |
+
if [ -d "weights/2023-10-28-18-33-37" ] && [ -d "weights/2024-01-11-20-02-45" ]; then
|
| 50 |
+
echo "✓ Model weights found"
|
| 51 |
+
echo ""
|
| 52 |
+
echo "Deploy in REAL mode (with model weights)?"
|
| 53 |
+
echo " - Pro: Actual pose estimation"
|
| 54 |
+
echo " - Con: Large files, GPU costs"
|
| 55 |
+
echo ""
|
| 56 |
+
read -p "Enable real mode? (y/N) " -n 1 -r
|
| 57 |
+
echo
|
| 58 |
+
if [[ $REPLY =~ ^[Yy]$ ]]; then
|
| 59 |
+
USE_REAL="true"
|
| 60 |
+
echo ""
|
| 61 |
+
echo "Note: Make sure git-lfs is set up for weights:"
|
| 62 |
+
echo " git lfs track 'weights/**'"
|
| 63 |
+
echo " git add .gitattributes"
|
| 64 |
+
echo ""
|
| 65 |
+
else
|
| 66 |
+
USE_REAL="false"
|
| 67 |
+
fi
|
| 68 |
+
else
|
| 69 |
+
echo "⚠ Model weights not found in weights/"
|
| 70 |
+
echo "Deploying in PLACEHOLDER mode (empty results)"
|
| 71 |
+
echo ""
|
| 72 |
+
echo "To add weights:"
|
| 73 |
+
echo " 1. Download from: https://drive.google.com/drive/folders/1GCyGE-LbFGgRC-FuGsF3a1zeBuzsQ1Da"
|
| 74 |
+
echo " 2. Extract to weights/ directory"
|
| 75 |
+
echo " 3. Re-run this script"
|
| 76 |
+
echo ""
|
| 77 |
+
USE_REAL="false"
|
| 78 |
+
fi
|
| 79 |
+
|
| 80 |
+
# Push to Hugging Face
|
| 81 |
+
echo "Pushing to Hugging Face Spaces..."
|
| 82 |
+
git push origin main
|
| 83 |
+
|
| 84 |
+
echo ""
|
| 85 |
+
echo "=========================================="
|
| 86 |
+
echo "Deployment Complete!"
|
| 87 |
+
echo "=========================================="
|
| 88 |
+
echo ""
|
| 89 |
+
echo "Your Space is available at:"
|
| 90 |
+
echo " $SPACE_URL"
|
| 91 |
+
echo ""
|
| 92 |
+
echo "Mode: $([ "$USE_REAL" = "true" ] && echo "🟢 Real FoundationPose" || echo "🟡 Placeholder")"
|
| 93 |
+
echo ""
|
| 94 |
+
|
| 95 |
+
if [ "$USE_REAL" = "false" ]; then
|
| 96 |
+
echo "To enable real mode:"
|
| 97 |
+
echo " 1. Add model weights to weights/ directory"
|
| 98 |
+
echo " 2. Set USE_REAL_MODEL=true in Space secrets"
|
| 99 |
+
echo " 3. Push again"
|
| 100 |
+
echo ""
|
| 101 |
+
fi
|
| 102 |
+
|
| 103 |
+
echo "Monitor build progress:"
|
| 104 |
+
echo " https://huggingface.co/spaces/gpue/foundationpose/logs"
|
| 105 |
+
echo ""
|
| 106 |
+
echo "Test the Space:"
|
| 107 |
+
echo " open $SPACE_URL"
|
| 108 |
+
echo ""
|
download_weights.py
ADDED
|
@@ -0,0 +1,73 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""
|
| 3 |
+
Download FoundationPose pre-trained model weights.
|
| 4 |
+
|
| 5 |
+
The official weights are hosted on Google Drive. This script downloads them
|
| 6 |
+
to the weights/ directory.
|
| 7 |
+
"""
|
| 8 |
+
|
| 9 |
+
import os
|
| 10 |
+
import sys
|
| 11 |
+
from pathlib import Path
|
| 12 |
+
|
| 13 |
+
try:
|
| 14 |
+
import gdown
|
| 15 |
+
except ImportError:
|
| 16 |
+
print("Installing gdown...")
|
| 17 |
+
os.system(f"{sys.executable} -m pip install gdown")
|
| 18 |
+
import gdown
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
def download_weights():
|
| 22 |
+
"""Download pre-trained FoundationPose weights."""
|
| 23 |
+
weights_dir = Path("weights")
|
| 24 |
+
weights_dir.mkdir(exist_ok=True)
|
| 25 |
+
|
| 26 |
+
# FoundationPose model weights (from official Google Drive)
|
| 27 |
+
# Note: These are the file IDs from the FoundationPose repo
|
| 28 |
+
weights_files = {
|
| 29 |
+
# Model checkpoint folders
|
| 30 |
+
"2023-10-28-18-33-37": "FOLDER_ID_1", # TODO: Replace with actual folder ID
|
| 31 |
+
"2024-01-11-20-02-45": "FOLDER_ID_2", # TODO: Replace with actual folder ID
|
| 32 |
+
}
|
| 33 |
+
|
| 34 |
+
print("=" * 60)
|
| 35 |
+
print("FoundationPose Weight Download")
|
| 36 |
+
print("=" * 60)
|
| 37 |
+
print()
|
| 38 |
+
print("Note: Official weights must be downloaded manually from:")
|
| 39 |
+
print("https://drive.google.com/drive/folders/1GCyGE-LbFGgRC-FuGsF3a1zeBuzsQ1Da")
|
| 40 |
+
print()
|
| 41 |
+
print("Required files:")
|
| 42 |
+
print(" - 2023-10-28-18-33-37/ (refiner weights)")
|
| 43 |
+
print(" - 2024-01-11-20-02-45/ (scorer weights)")
|
| 44 |
+
print()
|
| 45 |
+
print(f"Extract them to: {weights_dir.absolute()}")
|
| 46 |
+
print()
|
| 47 |
+
print("=" * 60)
|
| 48 |
+
|
| 49 |
+
# Check if weights already exist
|
| 50 |
+
weight_folders = [
|
| 51 |
+
weights_dir / "2023-10-28-18-33-37",
|
| 52 |
+
weights_dir / "2024-01-11-20-02-45"
|
| 53 |
+
]
|
| 54 |
+
|
| 55 |
+
if all(folder.exists() for folder in weight_folders):
|
| 56 |
+
print("✓ Model weights found!")
|
| 57 |
+
return True
|
| 58 |
+
else:
|
| 59 |
+
print("⚠ Model weights not found.")
|
| 60 |
+
print()
|
| 61 |
+
print("For Hugging Face Spaces deployment:")
|
| 62 |
+
print("1. Download weights manually")
|
| 63 |
+
print("2. Use git-lfs to add them to the repository:")
|
| 64 |
+
print(" git lfs track 'weights/**'")
|
| 65 |
+
print(" git add weights/")
|
| 66 |
+
print(" git commit -m 'Add model weights'")
|
| 67 |
+
print()
|
| 68 |
+
return False
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
if __name__ == "__main__":
|
| 72 |
+
success = download_weights()
|
| 73 |
+
sys.exit(0 if success else 1)
|
estimator.py
ADDED
|
@@ -0,0 +1,413 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
FoundationPose Estimator Wrapper
|
| 3 |
+
|
| 4 |
+
This module wraps the FoundationPose API for easy integration with the Gradio app.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
import logging
|
| 8 |
+
import sys
|
| 9 |
+
from pathlib import Path
|
| 10 |
+
from typing import Dict, List, Optional, Tuple
|
| 11 |
+
|
| 12 |
+
import cv2
|
| 13 |
+
import numpy as np
|
| 14 |
+
import torch
|
| 15 |
+
import trimesh
|
| 16 |
+
|
| 17 |
+
logger = logging.getLogger(__name__)
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
class FoundationPoseEstimator:
|
| 21 |
+
"""Wrapper for FoundationPose 6D pose estimation."""
|
| 22 |
+
|
| 23 |
+
def __init__(self, device: str = "cuda", weights_dir: str = "weights"):
|
| 24 |
+
"""Initialize FoundationPose.
|
| 25 |
+
|
| 26 |
+
Args:
|
| 27 |
+
device: Device to run inference on ("cuda" or "cpu")
|
| 28 |
+
weights_dir: Path to model weights directory
|
| 29 |
+
"""
|
| 30 |
+
self.device = device
|
| 31 |
+
self.weights_dir = Path(weights_dir)
|
| 32 |
+
|
| 33 |
+
# Add FoundationPose to Python path
|
| 34 |
+
foundationpose_dir = Path("FoundationPose")
|
| 35 |
+
if foundationpose_dir.exists():
|
| 36 |
+
sys.path.insert(0, str(foundationpose_dir))
|
| 37 |
+
else:
|
| 38 |
+
raise RuntimeError(
|
| 39 |
+
"FoundationPose repository not found. "
|
| 40 |
+
"Clone it with: git clone https://github.com/NVlabs/FoundationPose.git"
|
| 41 |
+
)
|
| 42 |
+
|
| 43 |
+
# Import FoundationPose modules
|
| 44 |
+
try:
|
| 45 |
+
from estimater import FoundationPose
|
| 46 |
+
from datareader import SceneReader
|
| 47 |
+
import pytorch3d.transforms as transforms
|
| 48 |
+
|
| 49 |
+
self.FoundationPose = FoundationPose
|
| 50 |
+
self.SceneReader = SceneReader
|
| 51 |
+
self.transforms = transforms
|
| 52 |
+
|
| 53 |
+
except ImportError as e:
|
| 54 |
+
raise RuntimeError(
|
| 55 |
+
f"Failed to import FoundationPose modules: {e}\n"
|
| 56 |
+
"Make sure FoundationPose is properly installed with all dependencies."
|
| 57 |
+
)
|
| 58 |
+
|
| 59 |
+
# Initialize models
|
| 60 |
+
self._init_models()
|
| 61 |
+
|
| 62 |
+
# Tracking state
|
| 63 |
+
self.tracked_objects = {}
|
| 64 |
+
self.pose_estimators = {}
|
| 65 |
+
|
| 66 |
+
def _init_models(self):
|
| 67 |
+
"""Initialize scorer and refiner models."""
|
| 68 |
+
logger.info("Initializing FoundationPose models...")
|
| 69 |
+
|
| 70 |
+
try:
|
| 71 |
+
# Load scorer model
|
| 72 |
+
scorer_weights = self.weights_dir / "2024-01-11-20-02-45"
|
| 73 |
+
if not scorer_weights.exists():
|
| 74 |
+
raise FileNotFoundError(f"Scorer weights not found at {scorer_weights}")
|
| 75 |
+
|
| 76 |
+
# Load refiner model
|
| 77 |
+
refiner_weights = self.weights_dir / "2023-10-28-18-33-37"
|
| 78 |
+
if not refiner_weights.exists():
|
| 79 |
+
raise FileNotFoundError(f"Refiner weights not found at {refiner_weights}")
|
| 80 |
+
|
| 81 |
+
# Import and initialize models (actual implementation depends on FoundationPose API)
|
| 82 |
+
from model import FoundationPoseModel
|
| 83 |
+
|
| 84 |
+
self.scorer = FoundationPoseModel(
|
| 85 |
+
checkpoint_dir=str(scorer_weights),
|
| 86 |
+
model_type="scorer"
|
| 87 |
+
).to(self.device)
|
| 88 |
+
self.scorer.eval()
|
| 89 |
+
|
| 90 |
+
self.refiner = FoundationPoseModel(
|
| 91 |
+
checkpoint_dir=str(refiner_weights),
|
| 92 |
+
model_type="refiner"
|
| 93 |
+
).to(self.device)
|
| 94 |
+
self.refiner.eval()
|
| 95 |
+
|
| 96 |
+
# Initialize CUDA rasterization context
|
| 97 |
+
import nvdiffrast.torch as dr
|
| 98 |
+
self.glctx = dr.RasterizeCudaContext()
|
| 99 |
+
|
| 100 |
+
logger.info("✓ Models initialized successfully")
|
| 101 |
+
|
| 102 |
+
except Exception as e:
|
| 103 |
+
logger.error(f"Failed to initialize models: {e}")
|
| 104 |
+
raise
|
| 105 |
+
|
| 106 |
+
def register_object(
|
| 107 |
+
self,
|
| 108 |
+
object_id: str,
|
| 109 |
+
reference_images: List[np.ndarray],
|
| 110 |
+
camera_intrinsics: Optional[Dict] = None,
|
| 111 |
+
mesh_path: Optional[str] = None
|
| 112 |
+
) -> bool:
|
| 113 |
+
"""Register an object for tracking.
|
| 114 |
+
|
| 115 |
+
Args:
|
| 116 |
+
object_id: Unique identifier for the object
|
| 117 |
+
reference_images: List of RGB images from different viewpoints
|
| 118 |
+
camera_intrinsics: Camera parameters (fx, fy, cx, cy)
|
| 119 |
+
mesh_path: Optional path to CAD mesh (for model-based mode)
|
| 120 |
+
|
| 121 |
+
Returns:
|
| 122 |
+
True if registration successful
|
| 123 |
+
"""
|
| 124 |
+
logger.info(f"Registering object '{object_id}'...")
|
| 125 |
+
|
| 126 |
+
try:
|
| 127 |
+
# Load or reconstruct mesh
|
| 128 |
+
if mesh_path and Path(mesh_path).exists():
|
| 129 |
+
# Model-based: use CAD mesh
|
| 130 |
+
mesh = trimesh.load(mesh_path)
|
| 131 |
+
logger.info(f"Loaded mesh from {mesh_path}")
|
| 132 |
+
else:
|
| 133 |
+
# Model-free: reconstruct from reference images
|
| 134 |
+
logger.info("Reconstructing mesh from reference images...")
|
| 135 |
+
mesh = self._reconstruct_mesh_from_references(
|
| 136 |
+
reference_images, camera_intrinsics
|
| 137 |
+
)
|
| 138 |
+
|
| 139 |
+
# Create FoundationPose estimator for this object
|
| 140 |
+
estimator = self.FoundationPose(
|
| 141 |
+
model_pts=mesh.vertices,
|
| 142 |
+
model_normals=mesh.vertex_normals,
|
| 143 |
+
mesh=mesh,
|
| 144 |
+
scorer=self.scorer,
|
| 145 |
+
refiner=self.refiner,
|
| 146 |
+
debug_dir=None,
|
| 147 |
+
debug=0,
|
| 148 |
+
glctx=self.glctx
|
| 149 |
+
)
|
| 150 |
+
|
| 151 |
+
# Store object data
|
| 152 |
+
self.tracked_objects[object_id] = {
|
| 153 |
+
"mesh": mesh,
|
| 154 |
+
"camera_intrinsics": camera_intrinsics,
|
| 155 |
+
"registered": True
|
| 156 |
+
}
|
| 157 |
+
self.pose_estimators[object_id] = {
|
| 158 |
+
"estimator": estimator,
|
| 159 |
+
"tracking": False,
|
| 160 |
+
"last_pose": None
|
| 161 |
+
}
|
| 162 |
+
|
| 163 |
+
logger.info(f"✓ Object '{object_id}' registered successfully")
|
| 164 |
+
return True
|
| 165 |
+
|
| 166 |
+
except Exception as e:
|
| 167 |
+
logger.error(f"Failed to register object: {e}", exc_info=True)
|
| 168 |
+
return False
|
| 169 |
+
|
| 170 |
+
def _reconstruct_mesh_from_references(
|
| 171 |
+
self,
|
| 172 |
+
reference_images: List[np.ndarray],
|
| 173 |
+
camera_intrinsics: Optional[Dict]
|
| 174 |
+
) -> trimesh.Trimesh:
|
| 175 |
+
"""Reconstruct 3D mesh from reference images using BundleSDF.
|
| 176 |
+
|
| 177 |
+
Args:
|
| 178 |
+
reference_images: List of RGB images
|
| 179 |
+
camera_intrinsics: Camera parameters
|
| 180 |
+
|
| 181 |
+
Returns:
|
| 182 |
+
Reconstructed mesh
|
| 183 |
+
"""
|
| 184 |
+
# TODO: Implement BundleSDF reconstruction
|
| 185 |
+
# For now, return a simple placeholder mesh
|
| 186 |
+
logger.warning("Mesh reconstruction not fully implemented, using placeholder")
|
| 187 |
+
|
| 188 |
+
# Create a simple cube mesh as placeholder
|
| 189 |
+
mesh = trimesh.creation.box(extents=[0.1, 0.1, 0.1])
|
| 190 |
+
return mesh
|
| 191 |
+
|
| 192 |
+
def estimate_pose(
|
| 193 |
+
self,
|
| 194 |
+
object_id: str,
|
| 195 |
+
rgb_image: np.ndarray,
|
| 196 |
+
depth_image: Optional[np.ndarray] = None,
|
| 197 |
+
mask: Optional[np.ndarray] = None,
|
| 198 |
+
camera_intrinsics: Optional[Dict] = None
|
| 199 |
+
) -> Optional[Dict]:
|
| 200 |
+
"""Estimate 6D pose of object in image.
|
| 201 |
+
|
| 202 |
+
Args:
|
| 203 |
+
object_id: ID of registered object
|
| 204 |
+
rgb_image: RGB image (H, W, 3)
|
| 205 |
+
depth_image: Optional depth map (H, W)
|
| 206 |
+
mask: Optional object segmentation mask (H, W)
|
| 207 |
+
camera_intrinsics: Camera parameters
|
| 208 |
+
|
| 209 |
+
Returns:
|
| 210 |
+
Pose dictionary with position, orientation, and confidence
|
| 211 |
+
"""
|
| 212 |
+
if object_id not in self.pose_estimators:
|
| 213 |
+
logger.error(f"Object '{object_id}' not registered")
|
| 214 |
+
return None
|
| 215 |
+
|
| 216 |
+
try:
|
| 217 |
+
estimator_data = self.pose_estimators[object_id]
|
| 218 |
+
estimator = estimator_data["estimator"]
|
| 219 |
+
|
| 220 |
+
# Get camera intrinsics
|
| 221 |
+
if camera_intrinsics is None:
|
| 222 |
+
camera_intrinsics = self.tracked_objects[object_id]["camera_intrinsics"]
|
| 223 |
+
|
| 224 |
+
K = self._build_intrinsics_matrix(camera_intrinsics, rgb_image.shape)
|
| 225 |
+
|
| 226 |
+
# Generate synthetic depth if not provided
|
| 227 |
+
if depth_image is None:
|
| 228 |
+
depth_image = np.zeros((rgb_image.shape[0], rgb_image.shape[1]), dtype=np.float32)
|
| 229 |
+
|
| 230 |
+
# Auto-segment if mask not provided
|
| 231 |
+
if mask is None:
|
| 232 |
+
mask = self._segment_object(rgb_image)
|
| 233 |
+
|
| 234 |
+
# First frame: register
|
| 235 |
+
if not estimator_data["tracking"]:
|
| 236 |
+
logger.info(f"Initial registration for '{object_id}'")
|
| 237 |
+
pose = estimator.register(
|
| 238 |
+
K=K,
|
| 239 |
+
rgb=rgb_image,
|
| 240 |
+
depth=depth_image,
|
| 241 |
+
ob_mask=mask,
|
| 242 |
+
iteration=5 # Number of refinement iterations
|
| 243 |
+
)
|
| 244 |
+
estimator_data["tracking"] = True
|
| 245 |
+
estimator_data["last_pose"] = pose
|
| 246 |
+
else:
|
| 247 |
+
# Subsequent frames: track
|
| 248 |
+
pose = estimator.track_one(
|
| 249 |
+
rgb=rgb_image,
|
| 250 |
+
depth=depth_image,
|
| 251 |
+
K=K,
|
| 252 |
+
iteration=2
|
| 253 |
+
)
|
| 254 |
+
estimator_data["last_pose"] = pose
|
| 255 |
+
|
| 256 |
+
# Convert pose matrix to position + quaternion
|
| 257 |
+
result = self._pose_matrix_to_dict(pose, object_id)
|
| 258 |
+
|
| 259 |
+
logger.info(f"Estimated pose for '{object_id}': confidence={result['confidence']:.3f}")
|
| 260 |
+
return result
|
| 261 |
+
|
| 262 |
+
except Exception as e:
|
| 263 |
+
logger.error(f"Pose estimation failed: {e}", exc_info=True)
|
| 264 |
+
return None
|
| 265 |
+
|
| 266 |
+
def _build_intrinsics_matrix(
|
| 267 |
+
self,
|
| 268 |
+
intrinsics: Optional[Dict],
|
| 269 |
+
image_shape: Tuple[int, int, int]
|
| 270 |
+
) -> np.ndarray:
|
| 271 |
+
"""Build camera intrinsics matrix.
|
| 272 |
+
|
| 273 |
+
Args:
|
| 274 |
+
intrinsics: Dict with fx, fy, cx, cy
|
| 275 |
+
image_shape: (H, W, C)
|
| 276 |
+
|
| 277 |
+
Returns:
|
| 278 |
+
3x3 intrinsics matrix
|
| 279 |
+
"""
|
| 280 |
+
H, W = image_shape[:2]
|
| 281 |
+
|
| 282 |
+
if intrinsics:
|
| 283 |
+
fx = intrinsics.get("fx", 500.0)
|
| 284 |
+
fy = intrinsics.get("fy", 500.0)
|
| 285 |
+
cx = intrinsics.get("cx", W / 2)
|
| 286 |
+
cy = intrinsics.get("cy", H / 2)
|
| 287 |
+
else:
|
| 288 |
+
# Default intrinsics
|
| 289 |
+
fx = fy = 500.0
|
| 290 |
+
cx = W / 2
|
| 291 |
+
cy = H / 2
|
| 292 |
+
|
| 293 |
+
K = np.array([
|
| 294 |
+
[fx, 0, cx],
|
| 295 |
+
[0, fy, cy],
|
| 296 |
+
[0, 0, 1]
|
| 297 |
+
], dtype=np.float32)
|
| 298 |
+
|
| 299 |
+
return K
|
| 300 |
+
|
| 301 |
+
def _segment_object(self, rgb_image: np.ndarray) -> np.ndarray:
|
| 302 |
+
"""Segment object from background.
|
| 303 |
+
|
| 304 |
+
This is a placeholder - in production, use SAM or similar.
|
| 305 |
+
|
| 306 |
+
Args:
|
| 307 |
+
rgb_image: RGB image
|
| 308 |
+
|
| 309 |
+
Returns:
|
| 310 |
+
Binary mask
|
| 311 |
+
"""
|
| 312 |
+
# Simple color-based segmentation placeholder
|
| 313 |
+
# In production, use Segment Anything Model (SAM)
|
| 314 |
+
H, W = rgb_image.shape[:2]
|
| 315 |
+
mask = np.ones((H, W), dtype=np.uint8) * 255
|
| 316 |
+
|
| 317 |
+
logger.warning("Using placeholder segmentation - implement SAM for production")
|
| 318 |
+
return mask
|
| 319 |
+
|
| 320 |
+
def _pose_matrix_to_dict(self, pose_matrix: np.ndarray, object_id: str) -> Dict:
|
| 321 |
+
"""Convert 4x4 pose matrix to dictionary format.
|
| 322 |
+
|
| 323 |
+
Args:
|
| 324 |
+
pose_matrix: 4x4 transformation matrix
|
| 325 |
+
object_id: Object identifier
|
| 326 |
+
|
| 327 |
+
Returns:
|
| 328 |
+
Dictionary with position, orientation (quaternion), confidence
|
| 329 |
+
"""
|
| 330 |
+
# Extract translation
|
| 331 |
+
position = {
|
| 332 |
+
"x": float(pose_matrix[0, 3]),
|
| 333 |
+
"y": float(pose_matrix[1, 3]),
|
| 334 |
+
"z": float(pose_matrix[2, 3])
|
| 335 |
+
}
|
| 336 |
+
|
| 337 |
+
# Extract rotation matrix and convert to quaternion
|
| 338 |
+
rotation_matrix = pose_matrix[:3, :3]
|
| 339 |
+
quat = self._rotation_matrix_to_quaternion(rotation_matrix)
|
| 340 |
+
|
| 341 |
+
orientation = {
|
| 342 |
+
"w": float(quat[0]),
|
| 343 |
+
"x": float(quat[1]),
|
| 344 |
+
"y": float(quat[2]),
|
| 345 |
+
"z": float(quat[3])
|
| 346 |
+
}
|
| 347 |
+
|
| 348 |
+
# Estimate confidence based on tracking state
|
| 349 |
+
# In production, use actual confidence from the model
|
| 350 |
+
confidence = 0.9 if self.pose_estimators[object_id]["tracking"] else 0.7
|
| 351 |
+
|
| 352 |
+
# Get object dimensions from mesh
|
| 353 |
+
mesh = self.tracked_objects[object_id]["mesh"]
|
| 354 |
+
extents = mesh.bounds[1] - mesh.bounds[0]
|
| 355 |
+
dimensions = [float(extents[0]), float(extents[1]), float(extents[2])]
|
| 356 |
+
|
| 357 |
+
return {
|
| 358 |
+
"object_id": object_id,
|
| 359 |
+
"position": position,
|
| 360 |
+
"orientation": orientation,
|
| 361 |
+
"confidence": confidence,
|
| 362 |
+
"dimensions": dimensions,
|
| 363 |
+
"timestamp": 0.0 # Add timestamp if needed
|
| 364 |
+
}
|
| 365 |
+
|
| 366 |
+
def _rotation_matrix_to_quaternion(self, R: np.ndarray) -> np.ndarray:
|
| 367 |
+
"""Convert 3x3 rotation matrix to quaternion (w, x, y, z).
|
| 368 |
+
|
| 369 |
+
Args:
|
| 370 |
+
R: 3x3 rotation matrix
|
| 371 |
+
|
| 372 |
+
Returns:
|
| 373 |
+
Quaternion as numpy array [w, x, y, z]
|
| 374 |
+
"""
|
| 375 |
+
trace = np.trace(R)
|
| 376 |
+
|
| 377 |
+
if trace > 0:
|
| 378 |
+
s = 0.5 / np.sqrt(trace + 1.0)
|
| 379 |
+
w = 0.25 / s
|
| 380 |
+
x = (R[2, 1] - R[1, 2]) * s
|
| 381 |
+
y = (R[0, 2] - R[2, 0]) * s
|
| 382 |
+
z = (R[1, 0] - R[0, 1]) * s
|
| 383 |
+
elif R[0, 0] > R[1, 1] and R[0, 0] > R[2, 2]:
|
| 384 |
+
s = 2.0 * np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2])
|
| 385 |
+
w = (R[2, 1] - R[1, 2]) / s
|
| 386 |
+
x = 0.25 * s
|
| 387 |
+
y = (R[0, 1] + R[1, 0]) / s
|
| 388 |
+
z = (R[0, 2] + R[2, 0]) / s
|
| 389 |
+
elif R[1, 1] > R[2, 2]:
|
| 390 |
+
s = 2.0 * np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2])
|
| 391 |
+
w = (R[0, 2] - R[2, 0]) / s
|
| 392 |
+
x = (R[0, 1] + R[1, 0]) / s
|
| 393 |
+
y = 0.25 * s
|
| 394 |
+
z = (R[1, 2] + R[2, 1]) / s
|
| 395 |
+
else:
|
| 396 |
+
s = 2.0 * np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1])
|
| 397 |
+
w = (R[1, 0] - R[0, 1]) / s
|
| 398 |
+
x = (R[0, 2] + R[2, 0]) / s
|
| 399 |
+
y = (R[1, 2] + R[2, 1]) / s
|
| 400 |
+
z = 0.25 * s
|
| 401 |
+
|
| 402 |
+
return np.array([w, x, y, z])
|
| 403 |
+
|
| 404 |
+
def reset_tracking(self, object_id: str):
|
| 405 |
+
"""Reset tracking state for an object.
|
| 406 |
+
|
| 407 |
+
Args:
|
| 408 |
+
object_id: Object to reset
|
| 409 |
+
"""
|
| 410 |
+
if object_id in self.pose_estimators:
|
| 411 |
+
self.pose_estimators[object_id]["tracking"] = False
|
| 412 |
+
self.pose_estimators[object_id]["last_pose"] = None
|
| 413 |
+
logger.info(f"Reset tracking for '{object_id}'")
|
requirements.txt
ADDED
|
@@ -0,0 +1,19 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Core dependencies
|
| 2 |
+
gradio>=4.0.0
|
| 3 |
+
spaces
|
| 4 |
+
numpy>=1.24.0
|
| 5 |
+
opencv-python>=4.8.0
|
| 6 |
+
Pillow>=10.0.0
|
| 7 |
+
|
| 8 |
+
# Deep learning
|
| 9 |
+
torch>=2.0.0
|
| 10 |
+
torchvision>=0.15.0
|
| 11 |
+
|
| 12 |
+
# 3D vision dependencies
|
| 13 |
+
trimesh>=4.0.0
|
| 14 |
+
pyrender>=0.1.45
|
| 15 |
+
scikit-image>=0.21.0
|
| 16 |
+
|
| 17 |
+
# FoundationPose specific (will need to install from source)
|
| 18 |
+
# The actual FoundationPose repo needs to be cloned and installed
|
| 19 |
+
# git+https://github.com/NVlabs/FoundationPose.git
|
test_local.py
ADDED
|
@@ -0,0 +1,264 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""
|
| 3 |
+
Test FoundationPose Space locally before deploying to Hugging Face.
|
| 4 |
+
|
| 5 |
+
This script tests both placeholder and real modes (if weights available).
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
import os
|
| 9 |
+
import sys
|
| 10 |
+
import time
|
| 11 |
+
from pathlib import Path
|
| 12 |
+
|
| 13 |
+
import cv2
|
| 14 |
+
import numpy as np
|
| 15 |
+
|
| 16 |
+
# Set to test placeholder mode
|
| 17 |
+
os.environ["USE_REAL_MODEL"] = "false"
|
| 18 |
+
|
| 19 |
+
print("=" * 60)
|
| 20 |
+
print("FoundationPose Local Test")
|
| 21 |
+
print("=" * 60)
|
| 22 |
+
print()
|
| 23 |
+
|
| 24 |
+
# Import after setting environment variable
|
| 25 |
+
try:
|
| 26 |
+
from app import pose_estimator
|
| 27 |
+
print("✓ Successfully imported app.py")
|
| 28 |
+
except Exception as e:
|
| 29 |
+
print(f"✗ Failed to import app.py: {e}")
|
| 30 |
+
sys.exit(1)
|
| 31 |
+
|
| 32 |
+
print(f"Mode: {'Real' if pose_estimator.use_real_model else 'Placeholder'}")
|
| 33 |
+
print()
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
def test_placeholder_mode():
|
| 37 |
+
"""Test the Space in placeholder mode."""
|
| 38 |
+
print("Test 1: Placeholder Mode")
|
| 39 |
+
print("-" * 40)
|
| 40 |
+
|
| 41 |
+
# Create dummy reference images
|
| 42 |
+
ref_images = []
|
| 43 |
+
for i in range(5):
|
| 44 |
+
img = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
|
| 45 |
+
ref_images.append(img)
|
| 46 |
+
|
| 47 |
+
# Test registration
|
| 48 |
+
print("Registering object with 5 reference images...")
|
| 49 |
+
start = time.time()
|
| 50 |
+
success = pose_estimator.register_object(
|
| 51 |
+
object_id="test_object",
|
| 52 |
+
reference_images=ref_images,
|
| 53 |
+
camera_intrinsics={"fx": 500, "fy": 500, "cx": 320, "cy": 240}
|
| 54 |
+
)
|
| 55 |
+
elapsed = time.time() - start
|
| 56 |
+
|
| 57 |
+
if success:
|
| 58 |
+
print(f"✓ Registration successful ({elapsed:.2f}s)")
|
| 59 |
+
else:
|
| 60 |
+
print(f"✗ Registration failed")
|
| 61 |
+
return False
|
| 62 |
+
|
| 63 |
+
# Test pose estimation
|
| 64 |
+
print("Estimating pose from query image...")
|
| 65 |
+
query_img = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
|
| 66 |
+
|
| 67 |
+
start = time.time()
|
| 68 |
+
result = pose_estimator.estimate_pose(
|
| 69 |
+
object_id="test_object",
|
| 70 |
+
query_image=query_img,
|
| 71 |
+
camera_intrinsics={"fx": 500, "fy": 500, "cx": 320, "cy": 240}
|
| 72 |
+
)
|
| 73 |
+
elapsed = time.time() - start
|
| 74 |
+
|
| 75 |
+
if result["success"]:
|
| 76 |
+
num_poses = len(result["poses"])
|
| 77 |
+
print(f"✓ Pose estimation successful ({elapsed:.2f}s)")
|
| 78 |
+
print(f" Detected poses: {num_poses}")
|
| 79 |
+
if num_poses == 0 and "note" in result:
|
| 80 |
+
print(f" Note: {result['note']}")
|
| 81 |
+
return True
|
| 82 |
+
else:
|
| 83 |
+
print(f"✗ Pose estimation failed: {result.get('error', 'Unknown')}")
|
| 84 |
+
return False
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
def test_with_reference_images():
|
| 88 |
+
"""Test with actual reference images if available."""
|
| 89 |
+
print()
|
| 90 |
+
print("Test 2: Real Reference Images")
|
| 91 |
+
print("-" * 40)
|
| 92 |
+
|
| 93 |
+
# Check for reference images
|
| 94 |
+
ref_dir = Path("../training/perception/reference/target_cube")
|
| 95 |
+
if not ref_dir.exists():
|
| 96 |
+
print("⊘ Reference images not found, skipping")
|
| 97 |
+
print(f" Expected at: {ref_dir}")
|
| 98 |
+
return True
|
| 99 |
+
|
| 100 |
+
# Load reference images
|
| 101 |
+
ref_images = []
|
| 102 |
+
for img_path in sorted(ref_dir.glob("*.jpg")):
|
| 103 |
+
img = cv2.imread(str(img_path))
|
| 104 |
+
if img is not None:
|
| 105 |
+
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
| 106 |
+
ref_images.append(img)
|
| 107 |
+
|
| 108 |
+
if not ref_images:
|
| 109 |
+
print("⊘ No .jpg files found in reference directory")
|
| 110 |
+
return True
|
| 111 |
+
|
| 112 |
+
print(f"Found {len(ref_images)} reference images")
|
| 113 |
+
|
| 114 |
+
# Test registration
|
| 115 |
+
print("Registering target_cube...")
|
| 116 |
+
start = time.time()
|
| 117 |
+
success = pose_estimator.register_object(
|
| 118 |
+
object_id="target_cube",
|
| 119 |
+
reference_images=ref_images
|
| 120 |
+
)
|
| 121 |
+
elapsed = time.time() - start
|
| 122 |
+
|
| 123 |
+
if success:
|
| 124 |
+
print(f"✓ Registration successful ({elapsed:.2f}s)")
|
| 125 |
+
else:
|
| 126 |
+
print(f"✗ Registration failed")
|
| 127 |
+
return False
|
| 128 |
+
|
| 129 |
+
# Test pose estimation with first reference image as query
|
| 130 |
+
print("Estimating pose (using first reference image as query)...")
|
| 131 |
+
start = time.time()
|
| 132 |
+
result = pose_estimator.estimate_pose(
|
| 133 |
+
object_id="target_cube",
|
| 134 |
+
query_image=ref_images[0]
|
| 135 |
+
)
|
| 136 |
+
elapsed = time.time() - start
|
| 137 |
+
|
| 138 |
+
if result["success"]:
|
| 139 |
+
num_poses = len(result["poses"])
|
| 140 |
+
print(f"✓ Pose estimation successful ({elapsed:.2f}s)")
|
| 141 |
+
print(f" Detected poses: {num_poses}")
|
| 142 |
+
|
| 143 |
+
if num_poses > 0:
|
| 144 |
+
pose = result["poses"][0]
|
| 145 |
+
print(f" Position: ({pose['position']['x']:.3f}, {pose['position']['y']:.3f}, {pose['position']['z']:.3f})")
|
| 146 |
+
print(f" Confidence: {pose['confidence']:.3f}")
|
| 147 |
+
else:
|
| 148 |
+
print(f" Note: {result.get('note', 'No poses detected')}")
|
| 149 |
+
|
| 150 |
+
return True
|
| 151 |
+
else:
|
| 152 |
+
print(f"✗ Pose estimation failed: {result.get('error', 'Unknown')}")
|
| 153 |
+
return False
|
| 154 |
+
|
| 155 |
+
|
| 156 |
+
def test_api_format():
|
| 157 |
+
"""Test that API format matches expected structure."""
|
| 158 |
+
print()
|
| 159 |
+
print("Test 3: API Format Validation")
|
| 160 |
+
print("-" * 40)
|
| 161 |
+
|
| 162 |
+
# Create test object
|
| 163 |
+
ref_img = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
|
| 164 |
+
pose_estimator.register_object("api_test", [ref_img])
|
| 165 |
+
|
| 166 |
+
# Get result
|
| 167 |
+
result = pose_estimator.estimate_pose("api_test", ref_img)
|
| 168 |
+
|
| 169 |
+
# Check format
|
| 170 |
+
required_keys = ["success", "poses"]
|
| 171 |
+
optional_keys = ["error", "note"]
|
| 172 |
+
|
| 173 |
+
print("Checking response format...")
|
| 174 |
+
|
| 175 |
+
for key in required_keys:
|
| 176 |
+
if key in result:
|
| 177 |
+
print(f" ✓ Has '{key}' field")
|
| 178 |
+
else:
|
| 179 |
+
print(f" ✗ Missing '{key}' field")
|
| 180 |
+
return False
|
| 181 |
+
|
| 182 |
+
if result["success"]:
|
| 183 |
+
if len(result["poses"]) > 0:
|
| 184 |
+
pose = result["poses"][0]
|
| 185 |
+
pose_required = ["object_id", "position", "orientation", "confidence", "dimensions"]
|
| 186 |
+
|
| 187 |
+
for key in pose_required:
|
| 188 |
+
if key in pose:
|
| 189 |
+
print(f" ✓ Pose has '{key}' field")
|
| 190 |
+
else:
|
| 191 |
+
print(f" ✗ Pose missing '{key}' field")
|
| 192 |
+
return False
|
| 193 |
+
|
| 194 |
+
# Check nested structure
|
| 195 |
+
if isinstance(pose["position"], dict) and "x" in pose["position"]:
|
| 196 |
+
print(f" ✓ Position format correct")
|
| 197 |
+
else:
|
| 198 |
+
print(f" ✗ Position format incorrect")
|
| 199 |
+
return False
|
| 200 |
+
|
| 201 |
+
if isinstance(pose["orientation"], dict) and "w" in pose["orientation"]:
|
| 202 |
+
print(f" ✓ Orientation format correct")
|
| 203 |
+
else:
|
| 204 |
+
print(f" ✗ Orientation format incorrect")
|
| 205 |
+
return False
|
| 206 |
+
else:
|
| 207 |
+
print(f" ℹ No poses detected (OK for placeholder mode)")
|
| 208 |
+
|
| 209 |
+
print("✓ API format valid")
|
| 210 |
+
return True
|
| 211 |
+
|
| 212 |
+
|
| 213 |
+
def main():
|
| 214 |
+
"""Run all tests."""
|
| 215 |
+
print("Starting tests...")
|
| 216 |
+
print()
|
| 217 |
+
|
| 218 |
+
tests = [
|
| 219 |
+
("Placeholder Mode", test_placeholder_mode),
|
| 220 |
+
("Reference Images", test_with_reference_images),
|
| 221 |
+
("API Format", test_api_format),
|
| 222 |
+
]
|
| 223 |
+
|
| 224 |
+
results = []
|
| 225 |
+
for name, test_func in tests:
|
| 226 |
+
try:
|
| 227 |
+
success = test_func()
|
| 228 |
+
results.append((name, success))
|
| 229 |
+
except Exception as e:
|
| 230 |
+
print(f"✗ Exception in {name}: {e}")
|
| 231 |
+
results.append((name, False))
|
| 232 |
+
|
| 233 |
+
# Summary
|
| 234 |
+
print()
|
| 235 |
+
print("=" * 60)
|
| 236 |
+
print("Test Summary")
|
| 237 |
+
print("=" * 60)
|
| 238 |
+
|
| 239 |
+
passed = sum(1 for _, success in results if success)
|
| 240 |
+
total = len(results)
|
| 241 |
+
|
| 242 |
+
for name, success in results:
|
| 243 |
+
status = "✓ PASS" if success else "✗ FAIL"
|
| 244 |
+
print(f"{status}: {name}")
|
| 245 |
+
|
| 246 |
+
print()
|
| 247 |
+
print(f"Results: {passed}/{total} tests passed")
|
| 248 |
+
|
| 249 |
+
if passed == total:
|
| 250 |
+
print()
|
| 251 |
+
print("🎉 All tests passed! Ready to deploy.")
|
| 252 |
+
print()
|
| 253 |
+
print("Next steps:")
|
| 254 |
+
print(" 1. Run './deploy.sh' to deploy to Hugging Face")
|
| 255 |
+
print(" 2. Or start locally: python app.py")
|
| 256 |
+
return 0
|
| 257 |
+
else:
|
| 258 |
+
print()
|
| 259 |
+
print("⚠ Some tests failed. Fix issues before deploying.")
|
| 260 |
+
return 1
|
| 261 |
+
|
| 262 |
+
|
| 263 |
+
if __name__ == "__main__":
|
| 264 |
+
sys.exit(main())
|