orient_express 2.1.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- orient_express-2.1.1/PKG-INFO +773 -0
- orient_express-2.1.1/README.md +746 -0
- orient_express-2.1.1/orient_express/__init__.py +1 -0
- orient_express-2.1.1/orient_express/deployment.py +97 -0
- orient_express-2.1.1/orient_express/model_wrapper.py +255 -0
- orient_express-2.1.1/orient_express/predictors/__init__.py +76 -0
- orient_express-2.1.1/orient_express/predictors/classification.py +67 -0
- orient_express-2.1.1/orient_express/predictors/instance_segmentation.py +185 -0
- orient_express-2.1.1/orient_express/predictors/multi_label_classification.py +68 -0
- orient_express-2.1.1/orient_express/predictors/object_detection.py +137 -0
- orient_express-2.1.1/orient_express/predictors/predictor.py +94 -0
- orient_express-2.1.1/orient_express/predictors/semantic_segmentation.py +97 -0
- orient_express-2.1.1/orient_express/sklearn_pipeline.py +81 -0
- orient_express-2.1.1/orient_express/utils/colors.py +29 -0
- orient_express-2.1.1/orient_express/utils/gs.py +116 -0
- orient_express-2.1.1/orient_express/utils/image_processor.py +91 -0
- orient_express-2.1.1/orient_express/utils/paths.py +6 -0
- orient_express-2.1.1/orient_express/utils/retry.py +47 -0
- orient_express-2.1.1/orient_express/vertex.py +312 -0
- orient_express-2.1.1/pyproject.toml +36 -0
|
@@ -0,0 +1,773 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: orient_express
|
|
3
|
+
Version: 2.1.1
|
|
4
|
+
Summary: A library to simplify model deployment to Vertex AI
|
|
5
|
+
Author: Alexey Zankevich
|
|
6
|
+
Author-email: alex.zankevich@shiftsmart.com
|
|
7
|
+
Requires-Python: >=3.10,<3.13
|
|
8
|
+
Classifier: Programming Language :: Python :: 3
|
|
9
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
10
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
12
|
+
Requires-Dist: Pillow (>=10.0.1)
|
|
13
|
+
Requires-Dist: gcsfs (>=2024.10.0,<2025.0.0)
|
|
14
|
+
Requires-Dist: google-cloud-aiplatform (>=1.0,<2.0)
|
|
15
|
+
Requires-Dist: google-cloud-storage (>=3.0,<4.0)
|
|
16
|
+
Requires-Dist: onnxruntime (>=1.20,<2.0) ; sys_platform == "darwin" or platform_machine == "aarch64" or platform_machine == "arm64" or sys_platform == "win32" and platform_machine == "ARM64"
|
|
17
|
+
Requires-Dist: onnxruntime-gpu (>=1.20,<2.0) ; sys_platform == "linux" and platform_machine == "x86_64" or sys_platform == "win32" and platform_machine == "AMD64"
|
|
18
|
+
Requires-Dist: opencv-python-headless (>=4.11.0,<5.0.0)
|
|
19
|
+
Requires-Dist: pandas (>=2.0,<3.0)
|
|
20
|
+
Requires-Dist: pyyaml (>=6.0,<7.0)
|
|
21
|
+
Requires-Dist: scikit-learn (>=1.5.0,<2.0.0)
|
|
22
|
+
Requires-Dist: torch (>=2.5.0,<3.0.0)
|
|
23
|
+
Requires-Dist: torchvision (>=0.20.0,<0.21.0)
|
|
24
|
+
Requires-Dist: xgboost (>=3.0.0,<4.0.0)
|
|
25
|
+
Description-Content-Type: text/markdown
|
|
26
|
+
|
|
27
|
+
# Orient Express
|
|
28
|
+
|
|
29
|
+
A library to accelerate model deployments to Vertex AI directly from colab notebooks
|
|
30
|
+
|
|
31
|
+

|
|
32
|
+
|
|
33
|
+
Orient Express provides two main capabilities:
|
|
34
|
+
|
|
35
|
+
1. **Vertex Model Deployment and Retrieval**: Capabilities for uploading, downloading, or deploying models to Vertex AI Model Registry.
|
|
36
|
+
|
|
37
|
+
1. **ONNX Image Model Deployment**: Built-in predictor classes for easily running image classification, object detection, instance segmentation, and semantic segmentation models exported to ONNX format.
|
|
38
|
+
|
|
39
|
+
Both workflows handle versioning, artifact storage in GCS, and integration with Vertex AI Model Registry.
|
|
40
|
+
|
|
41
|
+
## Installation
|
|
42
|
+
|
|
43
|
+
```bash
|
|
44
|
+
pip install orient_express
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
For local development:
|
|
48
|
+
```bash
|
|
49
|
+
pip install -e .
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
Or with Poetry:
|
|
53
|
+
```bash
|
|
54
|
+
poetry install
|
|
55
|
+
```
|
|
56
|
+
|
|
57
|
+
## Workflows
|
|
58
|
+
|
|
59
|
+
### ONNX Image Model Workflow
|
|
60
|
+
|
|
61
|
+
This workflow is for deploying image models (classification, detection, segmentation) exported to ONNX format.
|
|
62
|
+
|
|
63
|
+
```python
|
|
64
|
+
from orient_express.predictors import ClassificationPredictor
|
|
65
|
+
from orient_express.vertex import upload_model, get_vertex_model
|
|
66
|
+
|
|
67
|
+
# 1. Create predictor from your exported ONNX model
|
|
68
|
+
predictor = ClassificationPredictor(
|
|
69
|
+
onnx_path="model.onnx",
|
|
70
|
+
classes={1: "cat", 2: "dog", 3: "bird"}
|
|
71
|
+
)
|
|
72
|
+
|
|
73
|
+
# 2. Upload to Vertex AI Model Registry
|
|
74
|
+
vertex_model = upload_model(
|
|
75
|
+
model=predictor,
|
|
76
|
+
model_name="my-classifier",
|
|
77
|
+
project_name="my-project",
|
|
78
|
+
region="us-central1",
|
|
79
|
+
bucket_name="my-artifacts-bucket",
|
|
80
|
+
)
|
|
81
|
+
|
|
82
|
+
# 3. Later, retrieve and run locally
|
|
83
|
+
vertex_model = get_vertex_model(
|
|
84
|
+
model_name="my-classifier",
|
|
85
|
+
project_name="my-project",
|
|
86
|
+
region="us-central1",
|
|
87
|
+
)
|
|
88
|
+
local_predictor = vertex_model.get_local_predictor()
|
|
89
|
+
|
|
90
|
+
from PIL import Image
|
|
91
|
+
images = [Image.open("test.jpg")]
|
|
92
|
+
predictions = local_predictor.predict(images)
|
|
93
|
+
|
|
94
|
+
# 4. Or deploy to an endpoint for remote inference
|
|
95
|
+
vertex_model.deploy_to_endpoint(
|
|
96
|
+
endpoint_name="my-classifier-endpoint",
|
|
97
|
+
machine_type="n1-standard-4",
|
|
98
|
+
min_replica_count=1,
|
|
99
|
+
max_replica_count=3,
|
|
100
|
+
)
|
|
101
|
+
|
|
102
|
+
# remote prediction API depends on the endpoint container deployed with the model
|
|
103
|
+
predictions = vertex_model.remote_predict(
|
|
104
|
+
[{"image": "https://storage.googleapis.com/ssm-media-uploads/example.jpg"}],
|
|
105
|
+
endpoint_name="my-classifier-endpoint"
|
|
106
|
+
)
|
|
107
|
+
```
|
|
108
|
+
|
|
109
|
+
### Joblib Model Workflow
|
|
110
|
+
|
|
111
|
+
This workflow is for deploying models that can be serialized with joblib, such as scikit-learn pipelines or XGBoost models.
|
|
112
|
+
|
|
113
|
+
```python
|
|
114
|
+
from sklearn.pipeline import Pipeline
|
|
115
|
+
from sklearn.compose import ColumnTransformer
|
|
116
|
+
from sklearn.preprocessing import StandardScaler, OneHotEncoder
|
|
117
|
+
from sklearn.impute import SimpleImputer
|
|
118
|
+
import xgboost as xgb
|
|
119
|
+
import seaborn as sns
|
|
120
|
+
|
|
121
|
+
from orient_express.vertex import upload_model_joblib, get_vertex_model
|
|
122
|
+
|
|
123
|
+
# 1. Train your model
|
|
124
|
+
data = sns.load_dataset('titanic').dropna(subset=['survived'])
|
|
125
|
+
X = data[['pclass', 'sex', 'age', 'sibsp', 'parch', 'fare', 'embarked']]
|
|
126
|
+
y = data['survived']
|
|
127
|
+
|
|
128
|
+
numeric_features = ['age', 'fare', 'sibsp', 'parch']
|
|
129
|
+
numeric_transformer = Pipeline(steps=[
|
|
130
|
+
('imputer', SimpleImputer(strategy='median')),
|
|
131
|
+
('scaler', StandardScaler())
|
|
132
|
+
])
|
|
133
|
+
|
|
134
|
+
categorical_features = ['pclass', 'sex', 'embarked']
|
|
135
|
+
categorical_transformer = Pipeline(steps=[
|
|
136
|
+
('imputer', SimpleImputer(strategy='most_frequent')),
|
|
137
|
+
('onehot', OneHotEncoder(handle_unknown='ignore'))
|
|
138
|
+
])
|
|
139
|
+
|
|
140
|
+
preprocessor = ColumnTransformer(transformers=[
|
|
141
|
+
('num', numeric_transformer, numeric_features),
|
|
142
|
+
('cat', categorical_transformer, categorical_features)
|
|
143
|
+
])
|
|
144
|
+
|
|
145
|
+
model = Pipeline(steps=[
|
|
146
|
+
('preprocessor', preprocessor),
|
|
147
|
+
('classifier', xgb.XGBClassifier(use_label_encoder=False, eval_metric='logloss'))
|
|
148
|
+
])
|
|
149
|
+
|
|
150
|
+
model.fit(X, y)
|
|
151
|
+
|
|
152
|
+
# 2. Upload to Vertex AI Model Registry
|
|
153
|
+
vertex_model = upload_model_joblib(
|
|
154
|
+
model=model,
|
|
155
|
+
model_name="titanic-classifier",
|
|
156
|
+
project_name="my-project",
|
|
157
|
+
region="us-central1",
|
|
158
|
+
bucket_name="my-artifacts-bucket",
|
|
159
|
+
serving_container_image_uri="your-serving-container:latest",
|
|
160
|
+
serving_container_health_route="/health",
|
|
161
|
+
serving_container_predict_route="/predict",
|
|
162
|
+
)
|
|
163
|
+
|
|
164
|
+
# 3. Later, retrieve the model
|
|
165
|
+
vertex_model = get_vertex_model(
|
|
166
|
+
model_name="titanic-classifier",
|
|
167
|
+
project_name="my-project",
|
|
168
|
+
region="us-central1",
|
|
169
|
+
)
|
|
170
|
+
|
|
171
|
+
# 4. Run locally
|
|
172
|
+
local_predictor = vertex_model.get_local_predictor()
|
|
173
|
+
predictions = local_predictor.predict(X_test)
|
|
174
|
+
```
|
|
175
|
+
|
|
176
|
+
## ONNX Runtime and Device Support
|
|
177
|
+
|
|
178
|
+
### Platform Support Matrix
|
|
179
|
+
|
|
180
|
+
| Platform | Architecture | ONNX Runtime Package | CUDA Available |
|
|
181
|
+
|----------|--------------|---------------------|----------------|
|
|
182
|
+
| Linux | x86_64 | onnxruntime-gpu | Yes |
|
|
183
|
+
| Linux | aarch64 | onnxruntime | No |
|
|
184
|
+
| Windows | x64 (AMD64) | onnxruntime-gpu | Yes |
|
|
185
|
+
| Windows | ARM64 | onnxruntime | No |
|
|
186
|
+
| macOS | x86_64 | onnxruntime | No |
|
|
187
|
+
| macOS | arm64 | onnxruntime | No |
|
|
188
|
+
|
|
189
|
+
The appropriate package is installed automatically based on your platform.
|
|
190
|
+
|
|
191
|
+
### Selecting CPU vs CUDA Execution
|
|
192
|
+
|
|
193
|
+
When loading a predictor, use the `device` parameter to specify the execution provider:
|
|
194
|
+
```python
|
|
195
|
+
from orient_express.predictors import ObjectDetectionPredictor
|
|
196
|
+
|
|
197
|
+
# CPU inference (works on all platforms)
|
|
198
|
+
predictor = ObjectDetectionPredictor("/path/to/model", classes, device="cpu")
|
|
199
|
+
|
|
200
|
+
# CUDA inference (requires Linux x64 or Windows x64 with CUDA drivers)
|
|
201
|
+
predictor = ObjectDetectionPredictor("/path/to/model", classes, device="cuda")
|
|
202
|
+
```
|
|
203
|
+
|
|
204
|
+
When using a Vertex AI model:
|
|
205
|
+
```python
|
|
206
|
+
# CPU inference
|
|
207
|
+
predictor = model.get_local_predictor(device="cpu")
|
|
208
|
+
|
|
209
|
+
# CUDA inference
|
|
210
|
+
predictor = model.get_local_predictor(device="cuda")
|
|
211
|
+
```
|
|
212
|
+
|
|
213
|
+
### Pinning Model Versions
|
|
214
|
+
|
|
215
|
+
By default, `get_vertex_model` returns the most recently updated version. To pin to a specific version:
|
|
216
|
+
|
|
217
|
+
```python
|
|
218
|
+
vertex_model = get_vertex_model(
|
|
219
|
+
model_name="my-classifier",
|
|
220
|
+
project_name="my-project",
|
|
221
|
+
region="us-central1",
|
|
222
|
+
version=3, # Pin to version 3
|
|
223
|
+
)
|
|
224
|
+
```
|
|
225
|
+
|
|
226
|
+
---
|
|
227
|
+
|
|
228
|
+
## Built-in Predictor Types
|
|
229
|
+
|
|
230
|
+
Orient Express provides four built-in predictor classes for ONNX image models. Each has specific requirements for the ONNX graph structure.
|
|
231
|
+
|
|
232
|
+
### General ONNX Requirements
|
|
233
|
+
|
|
234
|
+
All ONNX image models share these requirements:
|
|
235
|
+
|
|
236
|
+
- **Input images are resized using simple stretch** (no letterboxing/padding) to the model's expected resolution before inference.
|
|
237
|
+
- **Normalization must be baked into the ONNX graph.** The library passes uint8 RGB images directly to the model; any normalization (e.g., ImageNet mean/std) must be handled inside the graph.
|
|
238
|
+
- **Batch dimension**: Models receive batched inputs with shape `[batch, height, width, 3]`.
|
|
239
|
+
|
|
240
|
+
### ClassificationPredictor
|
|
241
|
+
|
|
242
|
+
<details>
|
|
243
|
+
<summary>Click to expand</summary>
|
|
244
|
+
|
|
245
|
+
For image classification models that output class probabilities.
|
|
246
|
+
|
|
247
|
+
#### ONNX Graph Requirements
|
|
248
|
+
|
|
249
|
+
| | |
|
|
250
|
+
| ----------- | ------------------------------------------------------------ |
|
|
251
|
+
| **Inputs** | `images`: `[batch, height, width, 3]` uint8 RGB |
|
|
252
|
+
| **Outputs** | `scores`: `[batch, num_classes]` float32 class probabilities |
|
|
253
|
+
|
|
254
|
+
The graph must handle normalization internally. No target_sizes input is needed.
|
|
255
|
+
|
|
256
|
+
#### Usage
|
|
257
|
+
|
|
258
|
+
```python
|
|
259
|
+
from orient_express.predictors import ClassificationPredictor
|
|
260
|
+
|
|
261
|
+
predictor = ClassificationPredictor(
|
|
262
|
+
onnx_path="classifier.onnx",
|
|
263
|
+
classes={1: "cat", 2: "dog", 3: "bird"}
|
|
264
|
+
)
|
|
265
|
+
|
|
266
|
+
predictions = predictor.predict(images)
|
|
267
|
+
# Returns: list[ClassificationPrediction]
|
|
268
|
+
```
|
|
269
|
+
|
|
270
|
+
#### Output Structure
|
|
271
|
+
|
|
272
|
+
```python
|
|
273
|
+
@dataclass
|
|
274
|
+
class ClassificationPrediction:
|
|
275
|
+
clss: str # Predicted class name
|
|
276
|
+
score: float # Confidence score for predicted class
|
|
277
|
+
class_scores: dict[str, float] # Scores for all classes
|
|
278
|
+
|
|
279
|
+
# to_dict() output:
|
|
280
|
+
{
|
|
281
|
+
"class": "cat",
|
|
282
|
+
"score": 0.95,
|
|
283
|
+
"class_scores": {"cat": 0.95, "dog": 0.03, "bird": 0.02}
|
|
284
|
+
}
|
|
285
|
+
```
|
|
286
|
+
|
|
287
|
+
</details>
|
|
288
|
+
|
|
289
|
+
### MultiLabelClassificationPredictor
|
|
290
|
+
|
|
291
|
+
<details>
|
|
292
|
+
<summary>Click to expand</summary>
|
|
293
|
+
|
|
294
|
+
For image multi-label classification models that output a set of binary class probabilities.
|
|
295
|
+
|
|
296
|
+
#### ONNX Graph Requirements
|
|
297
|
+
|
|
298
|
+
| | |
|
|
299
|
+
| ----------- | ------------------------------------------------------------ |
|
|
300
|
+
| **Inputs** | `images`: `[batch, height, width, 3]` uint8 RGB |
|
|
301
|
+
| **Outputs** | `scores`: `[batch, num_classes]` float32 class probabilities |
|
|
302
|
+
|
|
303
|
+
The graph must handle normalization internally. No target_sizes input is needed.
|
|
304
|
+
|
|
305
|
+
#### Usage
|
|
306
|
+
|
|
307
|
+
```python
|
|
308
|
+
from orient_express.predictors import MultiLabelClassificationPredictor
|
|
309
|
+
|
|
310
|
+
predictor = MultiLabelClassificationPredictor(
|
|
311
|
+
onnx_path="classifier.onnx",
|
|
312
|
+
classes={1: "contains_cat", 2: "contains_dog", 3: "contains_bird"}
|
|
313
|
+
)
|
|
314
|
+
|
|
315
|
+
predictions = predictor.predict(images, confidence=0.5)
|
|
316
|
+
# Returns: list[MultiLabelClassificationPrediction]
|
|
317
|
+
```
|
|
318
|
+
|
|
319
|
+
#### Output Structure
|
|
320
|
+
|
|
321
|
+
```python
|
|
322
|
+
@dataclass
|
|
323
|
+
class MultiLabelClassificationPrediction:
|
|
324
|
+
classes: list[str] # Predicted class names based on confidence threshold
|
|
325
|
+
class_scores: dict[str, float] # Scores for all classes
|
|
326
|
+
|
|
327
|
+
# to_dict() output:
|
|
328
|
+
{
|
|
329
|
+
"classes": ["contains_cat", "contains_bird"],
|
|
330
|
+
"class_scores": {"contains_cat": 0.95, "contains_dog": 0.03, "contains_bird": 0.82}
|
|
331
|
+
}
|
|
332
|
+
```
|
|
333
|
+
|
|
334
|
+
</details>
|
|
335
|
+
|
|
336
|
+
### BoundingBoxPredictor
|
|
337
|
+
|
|
338
|
+
<details>
|
|
339
|
+
<summary>Click to expand</summary>
|
|
340
|
+
|
|
341
|
+
For object detection models that output bounding boxes.
|
|
342
|
+
|
|
343
|
+
#### ONNX Graph Requirements
|
|
344
|
+
|
|
345
|
+
| | |
|
|
346
|
+
| ----------- | ------------------------------------------------------------------------------------------------- |
|
|
347
|
+
| **Inputs** | `images`: `[batch, height, width, 3]` uint8 RGB |
|
|
348
|
+
| | `target_sizes`: `[batch, 2]` float32 containing `[height, width]` of original images |
|
|
349
|
+
| **Outputs** | `boxes`: `[batch, num_detections, 4]` float32 as `[x1, y1, x2, y2]` in original image coordinates |
|
|
350
|
+
| | `scores`: `[batch, num_detections]` float32 confidence scores |
|
|
351
|
+
| | `labels`: `[batch, num_detections]` int64 class indices |
|
|
352
|
+
|
|
353
|
+
The ONNX graph must rescale bounding boxes to the original image dimensions using `target_sizes`. The library does not perform any box coordinate transformation.
|
|
354
|
+
|
|
355
|
+
#### Usage
|
|
356
|
+
|
|
357
|
+
```python
|
|
358
|
+
from orient_express.predictors import BoundingBoxPredictor
|
|
359
|
+
|
|
360
|
+
predictor = BoundingBoxPredictor(
|
|
361
|
+
onnx_path="detector.onnx",
|
|
362
|
+
classes={1: "person", 2: "car", 3: "bicycle"}
|
|
363
|
+
)
|
|
364
|
+
|
|
365
|
+
predictions = predictor.predict(images, confidence=0.5, nms_threshold=0.4)
|
|
366
|
+
# Returns: list[list[BoundingBoxPrediction]]
|
|
367
|
+
# Outer list: per image, inner list: detections for that image
|
|
368
|
+
```
|
|
369
|
+
|
|
370
|
+
#### Output Structure
|
|
371
|
+
|
|
372
|
+
```python
|
|
373
|
+
@dataclass
|
|
374
|
+
class BoundingBoxPrediction:
|
|
375
|
+
clss: str # Class name
|
|
376
|
+
score: float # Confidence score
|
|
377
|
+
bbox: np.ndarray # [x1, y1, x2, y2] in original image coordinates
|
|
378
|
+
|
|
379
|
+
# to_dict() output:
|
|
380
|
+
{
|
|
381
|
+
"class": "person",
|
|
382
|
+
"score": 0.92,
|
|
383
|
+
"bbox": {"x1": 100.5, "y1": 50.2, "x2": 300.8, "y2": 400.1}
|
|
384
|
+
}
|
|
385
|
+
```
|
|
386
|
+
|
|
387
|
+
#### Annotation
|
|
388
|
+
|
|
389
|
+
```python
|
|
390
|
+
annotated_image = predictor.get_annotated_image(image, predictions[0])
|
|
391
|
+
# Returns PIL.Image with bounding boxes drawn
|
|
392
|
+
```
|
|
393
|
+
|
|
394
|
+
</details>
|
|
395
|
+
|
|
396
|
+
### InstanceSegmentationPredictor
|
|
397
|
+
|
|
398
|
+
<details>
|
|
399
|
+
<summary>Click to expand</summary>
|
|
400
|
+
|
|
401
|
+
For instance segmentation models that output bounding boxes and per-instance masks.
|
|
402
|
+
|
|
403
|
+
#### ONNX Graph Requirements
|
|
404
|
+
|
|
405
|
+
| | |
|
|
406
|
+
| ----------- | ------------------------------------------------------------------------------------------------- |
|
|
407
|
+
| **Inputs** | `images`: `[batch, height, width, 3]` uint8 RGB |
|
|
408
|
+
| | `target_sizes`: `[batch, 2]` float32 containing `[height, width]` of original images |
|
|
409
|
+
| **Outputs** | `boxes`: `[batch, num_detections, 4]` float32 as `[x1, y1, x2, y2]` in original image coordinates |
|
|
410
|
+
| | `scores`: `[batch, num_detections]` float32 confidence scores |
|
|
411
|
+
| | `labels`: `[batch, num_detections]` int64 class indices |
|
|
412
|
+
| | `masks`: `[batch, num_detections, mask_height, mask_width]` float32 mask logits |
|
|
413
|
+
|
|
414
|
+
The ONNX graph must rescale bounding boxes to original image dimensions using `target_sizes`. Masks can be any resolution—they are resized to original image dimensions in Python postprocessing using bilinear interpolation.
|
|
415
|
+
|
|
416
|
+
#### Usage
|
|
417
|
+
|
|
418
|
+
```python
|
|
419
|
+
from orient_express.predictors import InstanceSegmentationPredictor
|
|
420
|
+
|
|
421
|
+
predictor = InstanceSegmentationPredictor(
|
|
422
|
+
onnx_path="instance_seg.onnx",
|
|
423
|
+
classes={1: "person", 2: "car", 3: "bicycle"}
|
|
424
|
+
)
|
|
425
|
+
|
|
426
|
+
predictions = predictor.predict(images, confidence=0.5)
|
|
427
|
+
# Returns: list[list[InstanceSegmentationPrediction]]
|
|
428
|
+
```
|
|
429
|
+
|
|
430
|
+
#### Output Structure
|
|
431
|
+
|
|
432
|
+
```python
|
|
433
|
+
@dataclass
|
|
434
|
+
class InstanceSegmentationPrediction:
|
|
435
|
+
clss: str # Class name
|
|
436
|
+
score: float # Confidence score
|
|
437
|
+
bbox: np.ndarray # [x1, y1, x2, y2] in original image coordinates
|
|
438
|
+
mask: np.ndarray # Boolean mask at original image resolution
|
|
439
|
+
|
|
440
|
+
# to_dict(include_mask=False) output:
|
|
441
|
+
{
|
|
442
|
+
"class": "person",
|
|
443
|
+
"score": 0.89,
|
|
444
|
+
"bbox": {"x1": 100.5, "y1": 50.2, "x2": 300.8, "y2": 400.1}
|
|
445
|
+
}
|
|
446
|
+
|
|
447
|
+
# to_dict(include_mask=True) adds:
|
|
448
|
+
{
|
|
449
|
+
...
|
|
450
|
+
"mask": [[True, True, False, ...], ...] # 2D boolean list
|
|
451
|
+
}
|
|
452
|
+
```
|
|
453
|
+
|
|
454
|
+
#### Annotation
|
|
455
|
+
|
|
456
|
+
```python
|
|
457
|
+
annotated_image = predictor.get_annotated_image(image, predictions[0])
|
|
458
|
+
# Returns PIL.Image with mask overlays and instance indices
|
|
459
|
+
```
|
|
460
|
+
|
|
461
|
+
</details>
|
|
462
|
+
|
|
463
|
+
### SemanticSegmentationPredictor
|
|
464
|
+
|
|
465
|
+
<details>
|
|
466
|
+
<summary>Click to expand</summary>
|
|
467
|
+
|
|
468
|
+
For semantic segmentation models that output per-pixel class predictions.
|
|
469
|
+
|
|
470
|
+
#### ONNX Graph Requirements
|
|
471
|
+
|
|
472
|
+
| | |
|
|
473
|
+
| ----------- | ----------------------------------------------------------------------------- |
|
|
474
|
+
| **Inputs** | `images`: `[batch, height, width, 3]` uint8 RGB |
|
|
475
|
+
| **Outputs** | `masks`: `[batch, num_classes, mask_height, mask_width]` float32 class logits |
|
|
476
|
+
|
|
477
|
+
Masks can be any resolution—they are resized to original image dimensions in Python postprocessing. The class dimension is reduced via argmax to produce a single class ID per pixel.
|
|
478
|
+
|
|
479
|
+
#### Usage
|
|
480
|
+
|
|
481
|
+
```python
|
|
482
|
+
from orient_express.predictors import SemanticSegmentationPredictor
|
|
483
|
+
|
|
484
|
+
predictor = SemanticSegmentationPredictor(
|
|
485
|
+
onnx_path="semantic_seg.onnx",
|
|
486
|
+
classes={0: "background", 1: "road", 2: "building", 3: "vegetation"}
|
|
487
|
+
)
|
|
488
|
+
|
|
489
|
+
predictions = predictor.predict(images)
|
|
490
|
+
# Returns: list[SemanticSegmentationPrediction]
|
|
491
|
+
```
|
|
492
|
+
|
|
493
|
+
#### Output Structure
|
|
494
|
+
|
|
495
|
+
```python
|
|
496
|
+
@dataclass
|
|
497
|
+
class SemanticSegmentationPrediction:
|
|
498
|
+
class_mask: np.ndarray # [height, width] int array of class indices
|
|
499
|
+
conf_masks: np.ndarray # [num_classes, height, width] float confidence per class
|
|
500
|
+
|
|
501
|
+
# to_dict(include_conf_masks=False) output:
|
|
502
|
+
{
|
|
503
|
+
"class_mask": [[0, 0, 1, 2, ...], ...] # 2D int array
|
|
504
|
+
}
|
|
505
|
+
|
|
506
|
+
# to_dict(include_conf_masks=True) adds:
|
|
507
|
+
{
|
|
508
|
+
...
|
|
509
|
+
"conf_masks": [[[0.1, 0.2, ...], ...], ...] # 3D float array
|
|
510
|
+
}
|
|
511
|
+
```
|
|
512
|
+
|
|
513
|
+
#### Annotation
|
|
514
|
+
|
|
515
|
+
```python
|
|
516
|
+
annotated_image = predictor.get_annotated_image(image, predictions[0].class_mask)
|
|
517
|
+
# Returns PIL.Image with color-coded segmentation overlay
|
|
518
|
+
```
|
|
519
|
+
|
|
520
|
+
</details>
|
|
521
|
+
|
|
522
|
+
---
|
|
523
|
+
|
|
524
|
+
## Color Schemes
|
|
525
|
+
|
|
526
|
+
For predictors that support annotation (`BoundingBoxPredictor`, `InstanceSegmentationPredictor`, `SemanticSegmentationPredictor`), you can set a custom color scheme:
|
|
527
|
+
|
|
528
|
+
```python
|
|
529
|
+
predictor.color_scheme = {
|
|
530
|
+
"person": (255, 0, 0), # Red (RGB)
|
|
531
|
+
"car": (0, 255, 0), # Green
|
|
532
|
+
"bicycle": (0, 0, 255), # Blue
|
|
533
|
+
}
|
|
534
|
+
```
|
|
535
|
+
|
|
536
|
+
Colors are specified as RGB tuples.
|
|
537
|
+
|
|
538
|
+
## Legacy API [Still Maintained]
|
|
539
|
+
|
|
540
|
+
<details>
|
|
541
|
+
<summary>Click to expand</summary>
|
|
542
|
+
|
|
543
|
+
## Example
|
|
544
|
+
|
|
545
|
+
### Train Model
|
|
546
|
+
|
|
547
|
+
Train a regular model. In the example below, it's xgboost model, trained on the Titanic dataset.
|
|
548
|
+
|
|
549
|
+
```python
|
|
550
|
+
|
|
551
|
+
# Import necessary libraries
|
|
552
|
+
import xgboost as xgb
|
|
553
|
+
from sklearn.model_selection import train_test_split
|
|
554
|
+
from sklearn.metrics import accuracy_score, confusion_matrix
|
|
555
|
+
from sklearn.pipeline import Pipeline
|
|
556
|
+
from sklearn.compose import ColumnTransformer
|
|
557
|
+
from sklearn.preprocessing import StandardScaler, OneHotEncoder
|
|
558
|
+
from sklearn.impute import SimpleImputer
|
|
559
|
+
|
|
560
|
+
# Load the Titanic dataset
|
|
561
|
+
data = sns.load_dataset('titanic').dropna(subset=['survived']) # Dropping rows with missing target labels
|
|
562
|
+
|
|
563
|
+
# Select features and target
|
|
564
|
+
X = data[['pclass', 'sex', 'age', 'sibsp', 'parch', 'fare', 'embarked']]
|
|
565
|
+
y = data['survived']
|
|
566
|
+
|
|
567
|
+
# Define preprocessing for numeric columns (impute missing values and scale features)
|
|
568
|
+
numeric_features = ['age', 'fare', 'sibsp', 'parch']
|
|
569
|
+
numeric_transformer = Pipeline(steps=[
|
|
570
|
+
('imputer', SimpleImputer(strategy='median')),
|
|
571
|
+
('scaler', StandardScaler())
|
|
572
|
+
])
|
|
573
|
+
|
|
574
|
+
# Define preprocessing for categorical columns (impute missing values and one-hot encode)
|
|
575
|
+
categorical_features = ['pclass', 'sex', 'embarked']
|
|
576
|
+
categorical_transformer = Pipeline(steps=[
|
|
577
|
+
('imputer', SimpleImputer(strategy='most_frequent')),
|
|
578
|
+
('onehot', OneHotEncoder(handle_unknown='ignore'))
|
|
579
|
+
])
|
|
580
|
+
|
|
581
|
+
# Combine preprocessing steps
|
|
582
|
+
preprocessor = ColumnTransformer(
|
|
583
|
+
transformers=[
|
|
584
|
+
('num', numeric_transformer, numeric_features),
|
|
585
|
+
('cat', categorical_transformer, categorical_features)
|
|
586
|
+
])
|
|
587
|
+
|
|
588
|
+
# Create a pipeline that first transforms the data, then trains an XGBoost model
|
|
589
|
+
model = Pipeline(steps=[
|
|
590
|
+
('preprocessor', preprocessor),
|
|
591
|
+
('classifier', xgb.XGBClassifier(use_label_encoder=False, eval_metric='logloss'))
|
|
592
|
+
])
|
|
593
|
+
|
|
594
|
+
# Split the dataset into training and test sets
|
|
595
|
+
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
|
|
596
|
+
|
|
597
|
+
# Train the model
|
|
598
|
+
model.fit(X_train, y_train)
|
|
599
|
+
```
|
|
600
|
+
|
|
601
|
+
## Upload Model To Model Registry
|
|
602
|
+
|
|
603
|
+
```python
|
|
604
|
+
|
|
605
|
+
model_wrapper = ModelExpress(model=model,
|
|
606
|
+
project_name='my-project-name',
|
|
607
|
+
region='us-central1',
|
|
608
|
+
bucket_name='my-artifacts-bucket',
|
|
609
|
+
model_name='titanic')
|
|
610
|
+
model_wrapper.upload()
|
|
611
|
+
```
|
|
612
|
+
|
|
613
|
+
## Local Inference (Without Online Prediction Endpoint)
|
|
614
|
+
|
|
615
|
+
The following code will download the last model from the model registry and run the inference locally.
|
|
616
|
+
|
|
617
|
+
```python
|
|
618
|
+
|
|
619
|
+
# create input dataframe
|
|
620
|
+
titanic_data = {
|
|
621
|
+
"pclass": [1], # Passenger class (1st, 2nd, 3rd)
|
|
622
|
+
"sex": ["female"], # Gender
|
|
623
|
+
"age": [29], # Age
|
|
624
|
+
"sibsp": [0], # Number of siblings/spouses aboard
|
|
625
|
+
"parch": [0], # Number of parents/children aboard
|
|
626
|
+
"fare": [100.0], # Ticket fare
|
|
627
|
+
"embarked": ["S"] # Port of Embarkation (C = Cherbourg, Q = Queenstown, S = Southampton)
|
|
628
|
+
}
|
|
629
|
+
input_df = pd.DataFrame(titanic_data)
|
|
630
|
+
|
|
631
|
+
# init the model wrapper
|
|
632
|
+
model_wrapper = ModelExpress(project_name='my-project-name',
|
|
633
|
+
region='us-central1',
|
|
634
|
+
model_name='titanic')
|
|
635
|
+
|
|
636
|
+
# Run inference locally
|
|
637
|
+
# It will download the most recent version from the model registry automatically
|
|
638
|
+
model_wrapper.local_predict(input_df)
|
|
639
|
+
```
|
|
640
|
+
|
|
641
|
+
## Pin Model Version
|
|
642
|
+
|
|
643
|
+
In many cases, the pipeline should be pinned to a specific model version so the model can only
|
|
644
|
+
be updated explicitly. Just pass a `model_version` parameter when instantiating the ModelExpress wrapper.
|
|
645
|
+
|
|
646
|
+
```python
|
|
647
|
+
|
|
648
|
+
# init the model wrapper
|
|
649
|
+
model_wrapper = ModelExpress(project_name='my-project-name',
|
|
650
|
+
region='us-central1',
|
|
651
|
+
model_name='titanic',
|
|
652
|
+
model_version=11)
|
|
653
|
+
```
|
|
654
|
+
|
|
655
|
+
## Remote Inference (With Online Prediction Endpoint)
|
|
656
|
+
|
|
657
|
+
Make sure the model is deployed:
|
|
658
|
+
|
|
659
|
+
```python
|
|
660
|
+
|
|
661
|
+
model_wrapper = ModelExpress(model=model,
|
|
662
|
+
project_name='my-project-name',
|
|
663
|
+
region='us-central1',
|
|
664
|
+
bucket_name='my-artifacts-bucket',
|
|
665
|
+
model_name='titanic')
|
|
666
|
+
|
|
667
|
+
# upload the version to the registry and deploy it to the endpoint
|
|
668
|
+
model_wrapper.deploy()
|
|
669
|
+
```
|
|
670
|
+
|
|
671
|
+
Run inference with `remote_predict` method. It will make a remote call to the endpoint without fetching the model locally.
|
|
672
|
+
|
|
673
|
+
```python
|
|
674
|
+
|
|
675
|
+
titanic_data = {
|
|
676
|
+
"pclass": [1], # Passenger class (1st, 2nd, 3rd)
|
|
677
|
+
"sex": ["female"], # Gender
|
|
678
|
+
"age": [29], # Age
|
|
679
|
+
"sibsp": [0], # Number of siblings/spouses aboard
|
|
680
|
+
"parch": [0], # Number of parents/children aboard
|
|
681
|
+
"fare": [100.0], # Ticket fare
|
|
682
|
+
"embarked": ["S"] # Port of Embarkation (C = Cherbourg, Q = Queenstown, S = Southampton)
|
|
683
|
+
}
|
|
684
|
+
df = pd.DataFrame(titanic_data)
|
|
685
|
+
|
|
686
|
+
model_wrapper.remote_predict(df)
|
|
687
|
+
```
|
|
688
|
+
|
|
689
|
+
## Pipeline Deployment Function
|
|
690
|
+
|
|
691
|
+
Orient express library also have a helper function to simplify Vertex AI pipeline deployment.
|
|
692
|
+
|
|
693
|
+
Create `deploy.py` script
|
|
694
|
+
|
|
695
|
+
```python
|
|
696
|
+
|
|
697
|
+
from orient_express.deployment import deploy_pipeline
|
|
698
|
+
|
|
699
|
+
import argparse
|
|
700
|
+
import conf
|
|
701
|
+
|
|
702
|
+
from pipeline import pipeline
|
|
703
|
+
from orient_express.deployment import deploy_pipeline
|
|
704
|
+
|
|
705
|
+
|
|
706
|
+
if __name__ == "__main__":
|
|
707
|
+
parser = argparse.ArgumentParser()
|
|
708
|
+
parser.add_argument("--run-type", required=True)
|
|
709
|
+
|
|
710
|
+
args = parser.parse_args()
|
|
711
|
+
deploy_pipeline(run_type=args.run_type,
|
|
712
|
+
pipeline_dsl=pipeline,
|
|
713
|
+
pipeline_root=conf.PIPELINE_ROOT,
|
|
714
|
+
pipeline_name=conf.PIPELINE_NAME,
|
|
715
|
+
pipeline_display_name=conf.PIPELINE_DISPLAY_NAME,
|
|
716
|
+
pipeline_schedule_name=conf.SCHEDULE_NAME,
|
|
717
|
+
gcp_project=conf.PROJECT_ID,
|
|
718
|
+
gcp_location='us-central1',
|
|
719
|
+
gcp_service_account=conf.SERVICE_ACCOUNT,
|
|
720
|
+
gcp_network=conf.NETWORK_NAME,
|
|
721
|
+
gcp_labels={"team": "ml"})
|
|
722
|
+
```
|
|
723
|
+
|
|
724
|
+
And conf.py, make sure to replace the sample values with yours.
|
|
725
|
+
|
|
726
|
+
```python
|
|
727
|
+
|
|
728
|
+
import os
|
|
729
|
+
|
|
730
|
+
BASE_PATH = "gs://pipelines-bucket/vertex-ai/pipelines"
|
|
731
|
+
|
|
732
|
+
PIPELINE_NAME = "my-pipeline"
|
|
733
|
+
PIPELINE_ROOT = f"{BASE_PATH}/{PIPELINE_NAME}"
|
|
734
|
+
PIPELINE_TEMP_ROOT = f"{BASE_PATH}/{PIPELINE_NAME}-temp"
|
|
735
|
+
|
|
736
|
+
PIPELINE_DISPLAY_NAME = "My Pipeline"
|
|
737
|
+
PIPELINE_DESCRIPTION = "My example pipeline"
|
|
738
|
+
|
|
739
|
+
NETWORK_NAME = "project network id"
|
|
740
|
+
|
|
741
|
+
DOCKER_IMAGE = "us-docker.pkg.dev/my-project/my-artifactory/my-pipeline:latest
|
|
742
|
+
BASE_IMAGE = "python:3.11"
|
|
743
|
+
PROJECT_ID = "my-project"
|
|
744
|
+
PROJECT_REGION = "us-central1"
|
|
745
|
+
|
|
746
|
+
SERVICE_ACCOUNT = "my-service-account@my-project.iam.gserviceaccount.com"
|
|
747
|
+
SCHEDULE_NAME = "My Pipeline"
|
|
748
|
+
```
|
|
749
|
+
|
|
750
|
+
For testing it on a local machine, make sure to authorize to GCP first
|
|
751
|
+
|
|
752
|
+
```shell
|
|
753
|
+
|
|
754
|
+
gcloud auth application-default login
|
|
755
|
+
|
|
756
|
+
```
|
|
757
|
+
|
|
758
|
+
Finally, run the pipeline (it will execute once)
|
|
759
|
+
|
|
760
|
+
```shell
|
|
761
|
+
|
|
762
|
+
python deploy.py --run-type single-run
|
|
763
|
+
```
|
|
764
|
+
|
|
765
|
+
Or, create a scheduler to run continuously
|
|
766
|
+
|
|
767
|
+
```shell
|
|
768
|
+
|
|
769
|
+
python deploy.py --run-type scheduled
|
|
770
|
+
```
|
|
771
|
+
|
|
772
|
+
</details>
|
|
773
|
+
|