|
| 1 | +# Train and deploy YOLO Object Detection with ZenML, Ultralytics, and FiftyOne |
| 2 | + |
| 3 | +Learn how to build a production-ready computer vision pipeline with YOLOv8, FiftyOne dataset management, and deploy it as a warm HTTP service with an interactive web interface. This example showcases the complete **FiftyOne annotation workflow loop**: export → train → predict → import → analyze → visualize. |
| 4 | + |
| 5 | +## 🎯 What You'll Build |
| 6 | + |
| 7 | + |
| 8 | + |
| 9 | +- **Complete FiftyOne Annotation Workflow**: Export COCO data → Train YOLO → Run inference on original FiftyOne dataset → Import predictions back → Analyze performance → Interactive dashboard visualization |
| 10 | +- **Production-Ready Training**: YOLOv8 model training with automatic artifact versioning and performance tracking via ZenML |
| 11 | +- **Real-Time Inference Service**: Deploy as warm HTTP service with sub-second latency using ZenML's deployment system |
| 12 | +- **Interactive Web UI**: Upload images or use URLs for instant object detection testing with visual results |
| 13 | +- **Dataset-Model Lineage**: Full traceability linking FiftyOne datasets to ZenML model artifacts and predictions |
| 14 | +- **Visual Performance Analysis**: Side-by-side comparison of predictions vs ground truth in FiftyOne's interactive dashboard |
| 15 | + |
| 16 | +## 🏃 Quickstart |
| 17 | + |
| 18 | +```bash |
| 19 | +pip install -r requirements.txt |
| 20 | +zenml init |
| 21 | +zenml login |
| 22 | +``` |
| 23 | + |
| 24 | +**Train with FiftyOne analysis** ([see code](pipelines/training_pipeline.py)): |
| 25 | + |
| 26 | +```bash |
| 27 | +# Full workflow: training + FiftyOne analysis |
| 28 | +python run.py --train --samples 50 --epochs 3 |
| 29 | + |
| 30 | +# Fast training (skip FiftyOne analysis) |
| 31 | +python run.py --train --samples 50 --epochs 3 --disable-fiftyone-analysis |
| 32 | +``` |
| 33 | + |
| 34 | +This downloads 50 COCO validation images via FiftyOne, trains a YOLOv8 nano model for 3 epochs, runs inference on the original FiftyOne dataset, and provides interactive analysis capabilities. |
| 35 | + |
| 36 | +**Deploy as a real-time service** ([see code](pipelines/inference_pipeline.py)): |
| 37 | + |
| 38 | +```bash |
| 39 | +zenml pipeline deploy pipelines.inference_pipeline.object_detection_inference_pipeline |
| 40 | +``` |
| 41 | + |
| 42 | +Visit `http://localhost:8000` for the interactive UI ([see code](ui/index.html)). |
| 43 | + |
| 44 | +**Test batch inference locally**: |
| 45 | + |
| 46 | +The inference step supports multiple image input formats: |
| 47 | + |
| 48 | +```bash |
| 49 | +# Using image URL |
| 50 | +python run.py --predict --image https://ultralytics.com/images/bus.jpg |
| 51 | + |
| 52 | +# Using local file path |
| 53 | +python run.py --predict --image /path/to/local/image.jpg |
| 54 | + |
| 55 | +# Using base64 data URI (useful for programmatic usage) |
| 56 | +python run.py --predict --image "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQEA..." |
| 57 | +``` |
| 58 | + |
| 59 | +**Make predictions via API**: |
| 60 | + |
| 61 | +```bash |
| 62 | +# Using image URL |
| 63 | +curl -X POST http://localhost:8000/invoke \ |
| 64 | + -H "Content-Type: application/json" \ |
| 65 | + -d '{ |
| 66 | + "parameters": { |
| 67 | + "image_path": "https://ultralytics.com/images/bus.jpg", |
| 68 | + "confidence_threshold": 0.25 |
| 69 | + } |
| 70 | + }' |
| 71 | +``` |
| 72 | + |
| 73 | +**Explore with FiftyOne Dashboard**: |
| 74 | + |
| 75 | + |
| 76 | + |
| 77 | +After training, the pipeline creates persistent FiftyOne datasets linked to your ZenML model artifacts. Launch the dashboard to analyze results: |
| 78 | + |
| 79 | +```bash |
| 80 | +# Easy way - automatically finds datasets with predictions |
| 81 | +python launch_fiftyone.py |
| 82 | + |
| 83 | +# Launch specific dataset with custom port |
| 84 | +python launch_fiftyone.py coco-2017-validation-50samples --port 8080 |
| 85 | + |
| 86 | +# Manual way |
| 87 | +fiftyone app launch |
| 88 | +# Then in Python: fo.load_dataset('coco-2017-validation-50samples') |
| 89 | +``` |
| 90 | + |
| 91 | +**Dataset-Artifact Connection**: Each FiftyOne dataset (named by parameters like `coco-2017-validation-50samples`) stores predictions from your latest ZenML run. Re-running pipelines overwrites predictions while preserving ground truth, so the dashboard always shows your current model's performance. |
| 92 | + |
| 93 | +In the FiftyOne dashboard you can: |
| 94 | +- Compare predictions vs ground truth side-by-side |
| 95 | +- Filter by confidence levels and object classes |
| 96 | +- Analyze per-class performance metrics |
| 97 | +- Identify false positives and negatives |
| 98 | +- Export problematic samples for retraining |
| 99 | + |
| 100 | +**Use the ZenML Deployment Playground** |
| 101 | + |
| 102 | +The ZenML dashboard includes a built-in playground for deployed pipelines, allowing you to test your service directly from the UI. Navigate to your deployment in the dashboard, fill in the image URL and confidence threshold, and see real-time detection results with visualizations. |
| 103 | + |
| 104 | +## 🏗️ What's Inside |
| 105 | + |
| 106 | +``` |
| 107 | +computer_vision/ |
| 108 | +├── pipelines/ |
| 109 | +│ ├── training_pipeline.py - Train YOLO + optional FiftyOne analysis |
| 110 | +│ ├── inference_pipeline.py - Real-time detection service |
| 111 | +│ └── hooks.py - Warm model loading at startup/shutdown |
| 112 | +├── steps/ |
| 113 | +│ ├── data_loader.py - COCO dataset loading via FiftyOne |
| 114 | +│ ├── model_trainer.py - YOLO model training |
| 115 | +│ ├── evaluate.py - Model evaluation metrics |
| 116 | +│ ├── inference.py - Fast object detection (supports base64 uploads) |
| 117 | +│ └── fiftyone_analysis.py - Complete annotation workflow loop |
| 118 | +├── annotators/ |
| 119 | +│ ├── __init__.py - Annotator package initialization |
| 120 | +│ └── fiftyone_annotator.py - FiftyOne annotator class (ZenML-style) |
| 121 | +├── materializers/ |
| 122 | +│ └── ultralytics_materializer.py - Custom YOLO model serialization |
| 123 | +├── ui/ |
| 124 | +│ └── index.html - Interactive web interface with image upload |
| 125 | +├── run.py - CLI for training and testing |
| 126 | +├── launch_fiftyone.py - Easy FiftyOne dashboard launcher (uses annotator) |
| 127 | +└── requirements.txt - Dependencies |
| 128 | +``` |
| 129 | + |
| 130 | +## 🔑 Important Notes |
| 131 | + |
| 132 | +### **FiftyOne Annotator Class** |
| 133 | + |
| 134 | +This example includes a `FiftyOneAnnotator` class that encapsulates all FiftyOne functionality in a clean, ZenML-style interface. |
| 135 | + |
| 136 | +- 📄 [View FiftyOne annotator documentation](./annotators/README.md) |
| 137 | +- 🔧 [View annotator implementation](./annotators/fiftyone_annotator.py) |
| 138 | +- 🎯 **Key fix**: Proper COCO 80-class mapping ensures successful training (mAP@50: 0.799) vs broken training (0 mAP) |
| 139 | + |
| 140 | +### **Custom Materializers** |
| 141 | + |
| 142 | +This example features a [custom ZenML materializer](https://docs.zenml.io/how-to/types-and-materializers/materializers) for YOLO models that handles model weight serialization and artifact versioning, ensuring seamless model tracking across pipeline runs. ZenML automatically manages and version-controls model artifacts, making them accessible throughout your pipeline lifecycle. |
| 143 | + |
| 144 | +- 📖 [ZenML Materializers Documentation](https://docs.zenml.io/concepts/artifacts/materializers) |
| 145 | +- 📄 [View YOLO model materializer code](./materializers/ultralytics_materializer.py) |
| 146 | + |
| 147 | +### 🎨 Customization |
| 148 | + |
| 149 | +**Use a different dataset**: To use your own dataset (in YOLO format), modify the dataset loading logic in [`run.py`](./run.py) and/or the relevant pipeline step (e.g., `load_coco_dataset` in [`steps/data_loader.py`](./steps/data_loader.py)) to point to your images, labels, and `data.yaml`. |
| 150 | + |
| 151 | +**Use a larger model**: For better accuracy, use `yolov8s.pt`, `yolov8m.pt`, `yolov8l.pt`, or `yolov8x.pt`: |
| 152 | + |
| 153 | +```bash |
| 154 | +python run.py --train --model yolov8m.pt --epochs 10 |
| 155 | +``` |
| 156 | + |
| 157 | +**Run training on cloud**: Use ZenML's remote orchestrators for scalable training: |
| 158 | + |
| 159 | +```bash |
| 160 | +# Kubernetes orchestrator |
| 161 | +zenml orchestrator register k8s --flavor=kubernetes |
| 162 | +zenml stack update -o k8s |
| 163 | + |
| 164 | +# Then run training remotely |
| 165 | +python run.py --train --samples 500 --epochs 20 --model yolov8m.pt |
| 166 | +``` |
| 167 | + |
| 168 | +**Deploy inference to cloud**: Use ZenML's cloud deployers: |
| 169 | + |
| 170 | +```bash |
| 171 | +# AWS App Runner |
| 172 | +zenml deployer register aws --flavor=aws-app-runner --region=us-east-1 |
| 173 | +zenml stack update -d aws |
| 174 | +zenml pipeline deploy pipelines.inference_pipeline.object_detection_inference_pipeline |
| 175 | +``` |
| 176 | + |
| 177 | +## 📚 Learn More |
| 178 | + |
| 179 | +- [Pipeline Deployments Guide](https://docs.zenml.io/how-to/deployment/deployment) |
| 180 | +- [Deployment Settings](https://docs.zenml.io/how-to/deployment/deployment_settings) |
| 181 | +- [Pipeline Hooks](https://docs.zenml.io/how-to/steps-pipelines/advanced_features#pipeline-and-step-hooks) |
| 182 | +- [Ultralytics Documentation](https://docs.ultralytics.com/) |
| 183 | +- [FiftyOne Documentation](https://docs.voxel51.com/) |
| 184 | +- [Related Example: Deploying ML Models](../deploying_ml_model/README.md) |
0 commit comments