vagheshpatel commited on
Commit
3e1458f
·
verified ·
1 Parent(s): b4d0390

Sync heatmap-generation from metro-analytics-catalog

Browse files
.gitattributes CHANGED
@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ expected_output_dlstreamer.gif filter=lfs diff=lfs merge=lfs -text
37
+ expected_output_openvino.gif filter=lfs diff=lfs merge=lfs -text
LICENSE ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ MIT License
2
+
3
+ Copyright (c) Intel Corporation.
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE
README.md ADDED
@@ -0,0 +1,407 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: mit
3
+ license_link: LICENSE
4
+ library_name: openvino
5
+ pipeline_tag: object-detection
6
+ tags:
7
+ - openvino
8
+ - intel
9
+ - yolo
10
+ - yolo26
11
+ - heatmap
12
+ - speed
13
+ - traffic
14
+ - tracking
15
+ - edge-ai
16
+ - metro
17
+ - dlstreamer
18
+ language:
19
+ - en
20
+ ---
21
+
22
+ # Heatmap Generation
23
+
24
+ | Property | Value |
25
+ |---|---|
26
+ | **Category** | Object Detection + Speed Heatmap Aggregation |
27
+ | **Base Model** | [YOLO26](https://docs.ultralytics.com/models/yolo26/) (Ultralytics) |
28
+ | **Source Framework** | PyTorch (Ultralytics) |
29
+ | **Supported Precisions** | FP32, FP16, INT8 (mixed-precision) |
30
+ | **Inference Engine** | OpenVINO |
31
+ | **Hardware** | CPU, GPU, NPU |
32
+ | **Detected Class(es)** | All 80 COCO classes (heatmap colored by object speed) |
33
+
34
+ ---
35
+
36
+ ## Overview
37
+
38
+ Heatmap Generation is a Metro Analytics use case that detects objects across
39
+ video frames and colors each region of the scene by how fast traffic moves
40
+ through it. It is built on [YOLO26](https://docs.ultralytics.com/models/yolo26/), a
41
+ state-of-the-art real-time object detector, quantized to INT8 for efficient
42
+ inference on Intel hardware.
43
+
44
+ Each detection's per-frame displacement is used as a speed estimate, deposited
45
+ over the object's footprint and averaged per location with Gaussian smoothing
46
+ into a color-coded overlay.
47
+
48
+ The overlay uses the following color scheme:
49
+
50
+ - **Red** -- fast-moving traffic.
51
+ - **Yellow / green** -- medium speed.
52
+ - **Blue** -- slow-moving or stationary traffic.
53
+
54
+ Typical Metro deployments include:
55
+
56
+ - **Traffic Speed Mapping** -- highlight fast corridors and slow/congested lanes.
57
+ - **Congestion Detection** -- surface persistently slow (blue) areas for safety planning.
58
+ - **Pedestrian Flow Analysis** -- compare fast throughways against lingering areas.
59
+ - **Incident Spotting** -- flag unusually fast or stalled movement.
60
+
61
+ Available variants: `yolo26n`, `yolo26s`, `yolo26m`, `yolo26l`, `yolo26x`.
62
+ Smaller variants (`yolo26n`, `yolo26s`) are recommended for high-FPS edge
63
+ deployment; larger variants improve recall for small or distant objects.
64
+
65
+ ---
66
+
67
+ ## Prerequisites
68
+
69
+ - Python 3.11+
70
+ - [Install OpenVINO](https://docs.openvino.ai/2026/get-started/install-openvino.html) (latest version)
71
+ - [Install Intel DLStreamer](https://docs.openedgeplatform.intel.com/2026.0/edge-ai-libraries/dlstreamer/get_started/install/install_guide_ubuntu.html) (latest version)
72
+
73
+ Create and activate a Python virtual environment before running the scripts:
74
+
75
+ ```bash
76
+ python3 -m venv .venv --system-site-packages
77
+ source .venv/bin/activate
78
+ ```
79
+
80
+ > **Note:** The `--system-site-packages` flag is required so the virtual
81
+ > environment can access the system-installed OpenVINO and DLStreamer Python
82
+ > packages.
83
+
84
+ ---
85
+
86
+ ## Getting Started
87
+
88
+ ### Download and Quantize Model
89
+
90
+ Run the provided script to download, export to OpenVINO IR, and optionally quantize:
91
+
92
+ ```bash
93
+ chmod +x export_and_quantize.sh
94
+ ./export_and_quantize.sh
95
+ ```
96
+
97
+ This exports the default **yolo26n** model in **FP16** precision.
98
+
99
+ #### Optional: Select a Different Variant or Precision
100
+
101
+ ```bash
102
+ ./export_and_quantize.sh yolo26n FP32 # full-precision
103
+ ./export_and_quantize.sh yolo26n INT8 # quantized
104
+ ./export_and_quantize.sh yolo26s # larger variant, default FP16
105
+ ```
106
+
107
+ The script performs the following steps:
108
+
109
+ 1. Installs dependencies (`openvino`, `ultralytics`; adds `nncf` for INT8).
110
+ 2. Downloads a sample test image (`test.jpg`) and a sample test video (`test_video.mp4`).
111
+ 3. Downloads the PyTorch weights and exports to OpenVINO IR.
112
+ 4. *(INT8 only)* Quantizes the model using NNCF post-training quantization.
113
+
114
+ Output files:
115
+
116
+ - `yolo26n_openvino_model/` -- FP32 or FP16 OpenVINO IR model directory.
117
+ - `yolo26n_heatmap_int8.xml` / `.bin` -- INT8 quantized model *(only when `INT8` is selected)*.
118
+
119
+ #### Precision / Device Compatibility
120
+
121
+ | Precision | CPU | GPU | NPU |
122
+ |---|---|---|---|
123
+ | FP32 | Yes | Yes | No |
124
+ | FP16 | Yes | Yes | Yes |
125
+ | INT8 | Yes | Yes | Yes |
126
+
127
+ ### OpenVINO Sample
128
+
129
+ The sample below runs YOLO26 inference on a video, estimates each object's
130
+ speed from its per-frame displacement, and writes a speed-colored heatmap
131
+ overlay (red = fast, blue = slow) to `output_openvino.mp4`.
132
+ It also saves the final speed heatmap as `heatmap.jpg`.
133
+ Change the `device` string to run on CPU, GPU, or NPU.
134
+
135
+ ```python
136
+ import cv2
137
+ import numpy as np
138
+ import openvino as ov
139
+
140
+ CONF_THRESHOLD = 0.4
141
+ INPUT_SIZE = 640
142
+ HEATMAP_ALPHA = 0.55
143
+ MATCH_DIST = 80.0 # max px between frames to treat detections as the same object
144
+ MAX_SPEED = 20.0 # px/frame that maps to full red
145
+
146
+
147
+ def render_speed_heatmap(frame, speed_sum, count, alpha):
148
+ """Color traffic by average speed: blue = slow, yellow = medium,
149
+ red = fast. Only regions where vehicles were seen are tinted, so
150
+ empty background keeps its original color."""
151
+ avg = np.zeros_like(speed_sum)
152
+ seen = count > 0
153
+ avg[seen] = speed_sum[seen] / count[seen]
154
+ avg = cv2.GaussianBlur(avg, (0, 0), sigmaX=15)
155
+ presence = cv2.GaussianBlur(seen.astype(np.float32), (0, 0), sigmaX=15)
156
+ norm = np.clip(avg / MAX_SPEED, 0, 1) # 0 = slow (blue), 1 = fast (red)
157
+ color = cv2.applyColorMap((norm * 255).astype(np.uint8), cv2.COLORMAP_JET)
158
+ weight = (np.clip(presence, 0, 1) * alpha)[..., np.newaxis]
159
+ overlay = frame.astype(np.float32) * (1 - weight) + color.astype(np.float32) * weight
160
+ return overlay.astype(np.uint8), color
161
+
162
+
163
+ core = ov.Core()
164
+ model = core.read_model("yolo26n_openvino_model/yolo26n.xml")
165
+
166
+ # Change device to "GPU" or "NPU" to run on integrated GPU or NPU.
167
+ compiled = core.compile_model(model, "CPU")
168
+
169
+ cap = cv2.VideoCapture("test_video.mp4")
170
+ fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
171
+ width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
172
+ height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
173
+ writer = cv2.VideoWriter(
174
+ "output_openvino.mp4", cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))
175
+
176
+ speed_sum = np.zeros((height, width), dtype=np.float32)
177
+ count = np.zeros((height, width), dtype=np.float32)
178
+ prev_centroids = []
179
+ heatmap_color = None
180
+ frame_idx = 0
181
+ total_dets = 0
182
+
183
+ while True:
184
+ ok, frame = cap.read()
185
+ if not ok:
186
+ break
187
+ frame_idx += 1
188
+ h0, w0 = frame.shape[:2]
189
+ sx, sy = w0 / INPUT_SIZE, h0 / INPUT_SIZE
190
+
191
+ blob = cv2.resize(frame, (INPUT_SIZE, INPUT_SIZE))
192
+ blob = cv2.cvtColor(blob, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
193
+ blob = blob.transpose(2, 0, 1)[np.newaxis, ...]
194
+
195
+ output = compiled([blob])[compiled.output(0)][0]
196
+ dets = output[output[:, 4] >= CONF_THRESHOLD]
197
+ total_dets += len(dets)
198
+
199
+ cur_centroids = []
200
+ for det in dets:
201
+ x1, y1 = int(det[0] * sx), int(det[1] * sy)
202
+ x2, y2 = int(det[2] * sx), int(det[3] * sy)
203
+ x1, x2 = max(0, x1), min(width, x2)
204
+ y1, y2 = max(0, y1), min(height, y2)
205
+ cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
206
+ cur_centroids.append((cx, cy))
207
+
208
+ # Speed = displacement from the nearest detection in the previous frame.
209
+ speed = 0.0
210
+ if prev_centroids:
211
+ d = min(np.hypot(cx - px, cy - py) for px, py in prev_centroids)
212
+ if d <= MATCH_DIST:
213
+ speed = d
214
+ speed_sum[y1:y2, x1:x2] += speed
215
+ count[y1:y2, x1:x2] += 1.0
216
+ prev_centroids = cur_centroids
217
+
218
+ overlay, heatmap_color = render_speed_heatmap(frame, speed_sum, count, HEATMAP_ALPHA)
219
+ cv2.putText(overlay, f"Detections: {len(dets)}", (10, 30),
220
+ cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
221
+ writer.write(overlay)
222
+
223
+ cap.release()
224
+ writer.release()
225
+
226
+ if heatmap_color is not None:
227
+ cv2.imwrite("heatmap.jpg", heatmap_color)
228
+ print("Saved: heatmap.jpg")
229
+
230
+ print(f"Processed {frame_idx} frames, {total_dets} total detections", flush=True)
231
+ ```
232
+
233
+ **Device targets:**
234
+
235
+ - `"CPU"` -- default, works on all Intel platforms.
236
+ - `"GPU"` -- Intel integrated or discrete GPU.
237
+ - `"NPU"` -- Intel NPU (validate with `benchmark_app -d NPU`).
238
+
239
+ #### Expected Output
240
+
241
+ ![OpenVINO expected output](expected_output_openvino.gif)
242
+
243
+ ### DLStreamer Sample
244
+
245
+ The pipeline below runs the FP16 YOLO26 detector via `gvadetect`.
246
+ A buffer probe estimates each object's speed from its per-frame displacement
247
+ and overlays a speed-colored heatmap (red = fast, blue = slow) on each frame
248
+ before encoding to `output_dlstreamer.mp4`.
249
+
250
+ > **Notes on running this sample:**
251
+ >
252
+ > - Use the FP16 IR (`yolo26n_openvino_model/yolo26n.xml`). Class names are
253
+ > read automatically from the model's embedded `metadata.yaml` by
254
+ > DLStreamer 2026.0+ -- no external `labels-file` is required.
255
+ > - Export `PYTHONPATH` so the DLStreamer Python module is importable:
256
+ >
257
+ > ```bash
258
+ > source /opt/intel/openvino_2026/setupvars.sh
259
+ > source /opt/intel/dlstreamer/scripts/setup_dls_env.sh
260
+ > export PYTHONPATH=/opt/intel/dlstreamer/python:\
261
+ > /opt/intel/dlstreamer/gstreamer/lib/python3/dist-packages:${PYTHONPATH:-}
262
+ > ```
263
+
264
+ ```python
265
+ import gi
266
+
267
+ gi.require_version("Gst", "1.0")
268
+ gi.require_version("GstAnalytics", "1.0")
269
+ from gi.repository import Gst, GLib, GstAnalytics
270
+
271
+ import numpy as np
272
+
273
+ Gst.init([])
274
+
275
+ # Import cv2 after Gst.init to avoid GStreamer re-initialization conflicts.
276
+ import cv2
277
+
278
+ INPUT_VIDEO = "test_video.mp4"
279
+ HEATMAP_ALPHA = 0.55
280
+ MATCH_DIST = 80.0 # max px between frames to treat detections as the same object
281
+ MAX_SPEED = 20.0 # px/frame that maps to full red
282
+
283
+
284
+ def render_speed_heatmap(frame, speed_sum, count, alpha):
285
+ """Color traffic by average speed: blue = slow, yellow = medium,
286
+ red = fast. Only regions where vehicles were seen are tinted, so
287
+ empty background keeps its original color."""
288
+ avg = np.zeros_like(speed_sum)
289
+ seen = count > 0
290
+ avg[seen] = speed_sum[seen] / count[seen]
291
+ avg = cv2.GaussianBlur(avg, (0, 0), sigmaX=15)
292
+ presence = cv2.GaussianBlur(seen.astype(np.float32), (0, 0), sigmaX=15)
293
+ norm = np.clip(avg / MAX_SPEED, 0, 1) # 0 = slow (blue), 1 = fast (red)
294
+ color = cv2.applyColorMap((norm * 255).astype(np.uint8), cv2.COLORMAP_JET)
295
+ weight = (np.clip(presence, 0, 1) * alpha)[..., np.newaxis]
296
+ overlay = frame.astype(np.float32) * (1 - weight) + color.astype(np.float32) * weight
297
+ return overlay.astype(np.uint8)
298
+
299
+
300
+ # For CPU: change device=GPU to device=CPU.
301
+ # For NPU: change device=GPU to device=NPU (batch-size=1, nireq=4 recommended).
302
+ pipeline_str = (
303
+ f"filesrc location={INPUT_VIDEO} ! decodebin3 ! "
304
+ "videoconvert ! video/x-raw,format=BGR ! "
305
+ "gvadetect model=yolo26n_openvino_model/yolo26n.xml "
306
+ "device=GPU "
307
+ "threshold=0.4 ! queue ! "
308
+ "appsink name=sink emit-signals=false sync=false"
309
+ )
310
+ pipeline = Gst.parse_launch(pipeline_str)
311
+ sink = pipeline.get_by_name("sink")
312
+ pipeline.set_state(Gst.State.PLAYING)
313
+
314
+ speed_sum = None
315
+ count = None
316
+ prev_centroids = []
317
+ writer = None
318
+ frame_idx = 0
319
+ total_dets = 0
320
+
321
+ while True:
322
+ sample = sink.emit("pull-sample")
323
+ if sample is None:
324
+ break
325
+ buf = sample.get_buffer()
326
+ caps = sample.get_caps().get_structure(0)
327
+ width = caps.get_value("width")
328
+ height = caps.get_value("height")
329
+
330
+ if speed_sum is None:
331
+ speed_sum = np.zeros((height, width), dtype=np.float32)
332
+ count = np.zeros((height, width), dtype=np.float32)
333
+
334
+ ok, mapinfo = buf.map(Gst.MapFlags.READ)
335
+ if not ok:
336
+ continue
337
+ frame = np.ndarray((height, width, 3), dtype=np.uint8,
338
+ buffer=mapinfo.data).copy()
339
+ buf.unmap(mapinfo)
340
+ frame_idx += 1
341
+
342
+ rmeta = GstAnalytics.buffer_get_analytics_relation_meta(buf)
343
+ cur_centroids = []
344
+ det_count = 0
345
+ if rmeta is not None:
346
+ idx = 1
347
+ while True:
348
+ ok_od, od = rmeta.get_od_mtd(idx)
349
+ if not ok_od:
350
+ break
351
+ _, x, y, w, h, _ = od.get_location()
352
+ x1, y1 = max(0, int(x)), max(0, int(y))
353
+ x2, y2 = min(width, int(x + w)), min(height, int(y + h))
354
+ cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
355
+ cur_centroids.append((cx, cy))
356
+
357
+ # Speed = displacement from the nearest detection last frame.
358
+ speed = 0.0
359
+ if prev_centroids:
360
+ d = min(np.hypot(cx - px, cy - py) for px, py in prev_centroids)
361
+ if d <= MATCH_DIST:
362
+ speed = d
363
+ speed_sum[y1:y2, x1:x2] += speed
364
+ count[y1:y2, x1:x2] += 1.0
365
+ det_count += 1
366
+ idx += 1
367
+ prev_centroids = cur_centroids
368
+ total_dets += det_count
369
+
370
+ overlay = render_speed_heatmap(frame, speed_sum, count, HEATMAP_ALPHA)
371
+
372
+ if writer is None:
373
+ writer = cv2.VideoWriter(
374
+ "output_dlstreamer.mp4", cv2.VideoWriter_fourcc(*"mp4v"),
375
+ 30.0, (width, height))
376
+ writer.write(overlay)
377
+ print(f"Frame {frame_idx}: detections={det_count}", flush=True)
378
+
379
+ pipeline.set_state(Gst.State.NULL)
380
+ if writer:
381
+ writer.release()
382
+ print(f"Processed {frame_idx} frames, {total_dets} total detections", flush=True)
383
+ ```
384
+
385
+ **Device targets:**
386
+
387
+ - `device=GPU` -- default in the sample code.
388
+ - `device=CPU` -- change `device=GPU` to `device=CPU`.
389
+ - `device=NPU` -- change `device=GPU` to `device=NPU`; use `batch-size=1` and `nireq=4` for best NPU utilization.
390
+
391
+ #### Expected Output
392
+
393
+ ![DLStreamer expected output](expected_output_dlstreamer.gif)
394
+
395
+ ---
396
+
397
+ ## License
398
+
399
+ Licensed under the MIT License. See [LICENSE](LICENSE) for details.
400
+
401
+ ## References
402
+
403
+ - [YOLO26 Documentation](https://docs.ultralytics.com/models/yolo26/)
404
+ - [Ultralytics Heatmap Guide](https://docs.ultralytics.com/guides/heatmaps/)
405
+ - [OpenCV Color Maps](https://docs.opencv.org/4.x/d3/d50/group__imgproc__colormap.html)
406
+ - [OpenVINO Documentation](https://docs.openvino.ai/)
407
+ - [Intel DLStreamer](https://docs.openedgeplatform.intel.com/2026.0/edge-ai-libraries/dlstreamer/index.html)
expected_output_dlstreamer.gif ADDED

Git LFS Details

  • SHA256: 22e40d0b499e125e3b7803fa18e2431cc2036a6c7b7448a0b6ed7a3aa6b4615f
  • Pointer size: 133 Bytes
  • Size of remote file: 26 MB
expected_output_openvino.gif ADDED

Git LFS Details

  • SHA256: 5cac46824019cad1cbdb53c12a962114da5da420c7230ac432a0a42fd51baca8
  • Pointer size: 133 Bytes
  • Size of remote file: 26.1 MB
export_and_quantize.sh ADDED
@@ -0,0 +1,104 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env bash
2
+ # SPDX-License-Identifier: MIT
3
+ # Copyright (C) Intel Corporation
4
+ #
5
+ # Export a YOLO26 detector to OpenVINO IR for the heatmap-generation use case.
6
+ # Usage: ./export_and_quantize.sh [MODEL_VARIANT] [PRECISION]
7
+ # Example: ./export_and_quantize.sh yolo26n FP16
8
+
9
+ set -euo pipefail
10
+
11
+ MODEL_NAME="${1:-yolo26n}"
12
+ PRECISION="${2:-FP16}"
13
+ PRECISION="$(echo "${PRECISION}" | tr '[:lower:]' '[:upper:]')"
14
+
15
+ if [[ "${PRECISION}" != "FP32" && "${PRECISION}" != "FP16" && "${PRECISION}" != "INT8" ]]; then
16
+ echo "ERROR: unsupported precision '${PRECISION}'. Choose FP32, FP16, or INT8." >&2
17
+ exit 1
18
+ fi
19
+
20
+ echo "--- Installing dependencies ---"
21
+ if [[ "${PRECISION}" == "INT8" ]]; then
22
+ pip install -qU openvino nncf ultralytics
23
+ else
24
+ pip install -qU openvino ultralytics
25
+ fi
26
+
27
+ # Ask for approval before downloading models and sample files
28
+ echo ""
29
+ echo "This script will download:"
30
+ echo " - Model weights and/or sample files"
31
+ echo ""
32
+ read -p "Continue with downloads? (yes/no): " APPROVAL
33
+ if [[ "${APPROVAL}" != "yes" ]]; then
34
+ echo "Download cancelled by user."
35
+ exit 0
36
+ fi
37
+ echo ""
38
+ echo "--- Downloading sample test image ---"
39
+ if [[ ! -f test.jpg ]]; then
40
+ wget -q -O test.jpg https://ultralytics.com/images/bus.jpg
41
+ echo "Downloaded: test.jpg"
42
+ else
43
+ echo "Already present: test.jpg"
44
+ fi
45
+ echo ""
46
+ echo "--- Downloading sample test video ---"
47
+ if [[ ! -f test_video.mp4 ]]; then
48
+ wget -q -O test_video.mp4 \
49
+ "https://www.pexels.com/download/video/12399570/"
50
+ echo "Downloaded: test_video.mp4"
51
+ else
52
+ echo "Already present: test_video.mp4"
53
+ fi
54
+
55
+ if [[ "${PRECISION}" == "FP32" ]]; then
56
+ HALF_FLAG="False"
57
+ EXPORT_LABEL="FP32"
58
+ else
59
+ HALF_FLAG="True"
60
+ EXPORT_LABEL="FP16"
61
+ fi
62
+
63
+ echo "--- Exporting ${MODEL_NAME} to OpenVINO IR (${EXPORT_LABEL}) ---"
64
+ python3 -c "
65
+ from ultralytics import YOLO
66
+
67
+ model = YOLO('${MODEL_NAME}.pt')
68
+ model.export(format='openvino', half=${HALF_FLAG}, dynamic=False, imgsz=640)
69
+ print('Export complete: ${MODEL_NAME}_openvino_model/')
70
+ "
71
+
72
+ if [[ "${PRECISION}" == "INT8" ]]; then
73
+ echo "--- Quantizing to INT8 with NNCF ---"
74
+ python3 -c "
75
+ import nncf
76
+ import openvino as ov
77
+ import numpy as np
78
+ import cv2
79
+
80
+ core = ov.Core()
81
+ model = core.read_model('${MODEL_NAME}_openvino_model/${MODEL_NAME}.xml')
82
+
83
+ img = cv2.imread('test.jpg')
84
+ img = cv2.resize(img, (640, 640))
85
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
86
+ img = img.transpose(2, 0, 1)[np.newaxis, ...]
87
+
88
+ def transform_fn(data_item):
89
+ return img
90
+
91
+ calibration_dataset = nncf.Dataset(list(range(300)), transform_fn)
92
+
93
+ quantized = nncf.quantize(
94
+ model,
95
+ calibration_dataset,
96
+ preset=nncf.QuantizationPreset.MIXED,
97
+ subset_size=300,
98
+ )
99
+
100
+ ov.save_model(quantized, '${MODEL_NAME}_heatmap_int8.xml')
101
+ print('Quantization complete: ${MODEL_NAME}_heatmap_int8.xml')
102
+ "
103
+ fi
104
+ echo "--- Done ---"