meaculpitt commited on
Commit
438565c
Β·
verified Β·
1 Parent(s): 95e602a

scorevision: push petrol v2 model

Browse files
Files changed (5) hide show
  1. chute_config.yml +22 -0
  2. class_names.txt +4 -0
  3. miner.py +370 -0
  4. model_type.json +4 -0
  5. weights.onnx +3 -0
chute_config.yml ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Image:
2
+ from_base: parachutes/python:3.12
3
+ run_command:
4
+ - pip install --upgrade setuptools wheel
5
+ - pip install 'numpy>=1.23' 'onnxruntime-gpu>=1.16' 'nvidia-cudnn-cu12' 'nvidia-cublas-cu12'
6
+ 'opencv-python-headless>=4.7' 'pillow>=9.5' 'huggingface_hub>=0.19.4' 'pydantic>=2.0'
7
+ 'pyyaml>=6.0' 'aiohttp>=3.9' 'ensemble-boxes>=1.0' 'torch>=2.6,<3.0'
8
+ NodeSelector:
9
+ gpu_count: 1
10
+ min_vram_gb_per_gpu: 16
11
+ max_hourly_price_per_gpu: 0.50
12
+ exclude:
13
+ - '5090'
14
+ - b200
15
+ - h200
16
+ - mi300x
17
+ Chute:
18
+ timeout_seconds: 300
19
+ concurrency: 4
20
+ max_instances: 5
21
+ scaling_threshold: 0.5
22
+ shutdown_after_seconds: 288000
class_names.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ petrol hose
2
+ petrol pump
3
+ price board
4
+ roof canopy
miner.py ADDED
@@ -0,0 +1,370 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from pathlib import Path
2
+ import math
3
+ import logging
4
+
5
+ import cv2
6
+ import numpy as np
7
+ import onnxruntime as ort
8
+ from numpy import ndarray
9
+ from pydantic import BaseModel
10
+
11
+ logger = logging.getLogger(__name__)
12
+
13
+ # ─── Petrol miner v1.1 ───────────────────────────────────────────────
14
+ # Improvements over auto-generated baseline:
15
+ # 1. Fix end-to-end ONNX decode (model outputs [1,300,6] post-NMS)
16
+ # 2. Spatial co-occurrence scoring (pump+canopy boost, isolated suppress)
17
+ # 3. Geometric validation (aspect ratio + size checks per class)
18
+ # ──────────────────────────────────────────────────────────────────────
19
+
20
+ # Class IDs
21
+ CLS_HOSE = 0
22
+ CLS_PUMP = 1
23
+ CLS_PRICEBOARD = 2
24
+ CLS_CANOPY = 3
25
+
26
+ # ── Geometric validation thresholds (derived from 2000-label analysis) ──
27
+ # Canopy: wide/flat, aspect(w/h) mean=2.96. Suppress if aspect < 0.8 (too tall)
28
+ CANOPY_MIN_ASPECT = 0.8
29
+ # Pump: roughly square/tall, aspect mean=0.91. Suppress if aspect > 4.0 (too wide)
30
+ PUMP_MAX_ASPECT = 4.0
31
+ # Price board: small. Suppress if area > 15% of image
32
+ PRICEBOARD_MAX_AREA_FRAC = 0.15
33
+ # Hose: variable. Suppress if area < 0.05% of image (tiny FP)
34
+ HOSE_MIN_AREA_FRAC = 0.0005
35
+
36
+ # ── Spatial co-occurrence boost/suppress amounts ──
37
+ COOCCUR_BOOST_PUMP_CANOPY = 0.05
38
+ COOCCUR_BOOST_PUMP_HOSE = 0.08
39
+ COOCCUR_BOOST_CANOPY_HOSE = 0.05
40
+ COOCCUR_SUPPRESS_ISOLATED = 0.03 # per missing expected neighbor
41
+ # Proximity threshold: normalized distance between box centers
42
+ COOCCUR_PROXIMITY = 0.5 # half of image dimension
43
+
44
+ # ── Geometric suppress penalty ──
45
+ GEOMETRIC_SUPPRESS_PENALTY = 0.10
46
+
47
+
48
+ class BoundingBox(BaseModel):
49
+ x1: int
50
+ y1: int
51
+ x2: int
52
+ y2: int
53
+ cls_id: int
54
+ conf: float
55
+
56
+
57
+ class TVFrameResult(BaseModel):
58
+ frame_id: int
59
+ boxes: list[BoundingBox]
60
+ keypoints: list[tuple[int, int]]
61
+
62
+
63
+ class Miner:
64
+ VERSION = "petrol-v1.1"
65
+
66
+ def __init__(self, path_hf_repo: Path) -> None:
67
+ self.path_hf_repo = path_hf_repo
68
+ self.class_names = ['petrol hose', 'petrol pump', 'price board', 'roof canopy']
69
+ self.session = ort.InferenceSession(
70
+ str(path_hf_repo / "weights.onnx"),
71
+ providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
72
+ )
73
+ self.input_name = self.session.get_inputs()[0].name
74
+ input_shape = self.session.get_inputs()[0].shape
75
+ self.input_h = int(input_shape[2])
76
+ self.input_w = int(input_shape[3])
77
+ self.conf_threshold = 0.25
78
+ self.iou_threshold = 0.45
79
+
80
+ # Detect output format: end-to-end [1,N,6] vs raw [1,C,N]
81
+ out_shape = self.session.get_outputs()[0].shape
82
+ # End-to-end: [1, max_dets, 6] where max_dets is small (100-300)
83
+ # Raw: [1, 4+nc, N] where N is large (8400+)
84
+ if len(out_shape) == 3 and out_shape[2] == 6 and (out_shape[1] or 0) <= 1000:
85
+ self._end2end = True
86
+ logger.info("[init] End-to-end ONNX output detected")
87
+ else:
88
+ self._end2end = False
89
+ logger.info("[init] Raw ONNX output detected")
90
+
91
+ logger.info(f"[init] {self.VERSION} loaded, input={self.input_w}x{self.input_h}, "
92
+ f"end2end={self._end2end}")
93
+
94
+ def __repr__(self) -> str:
95
+ return f"Petrol Miner {self.VERSION} end2end={self._end2end}"
96
+
97
+ # ─── Preprocessing ────────────────────────────────────────────────
98
+
99
+ def _preprocess(self, image_bgr: ndarray) -> tuple[np.ndarray, tuple[int, int]]:
100
+ h, w = image_bgr.shape[:2]
101
+ rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
102
+ resized = cv2.resize(rgb, (self.input_w, self.input_h))
103
+ x = resized.astype(np.float32) / 255.0
104
+ x = np.transpose(x, (2, 0, 1))[None, ...]
105
+ return x, (h, w)
106
+
107
+ # ─── NMS (only needed for raw output format) ─────────────────────
108
+
109
+ def _nms(self, dets):
110
+ if not dets:
111
+ return []
112
+ boxes = np.array([[d[0], d[1], d[2], d[3]] for d in dets], dtype=np.float32)
113
+ scores = np.array([d[4] for d in dets], dtype=np.float32)
114
+ order = scores.argsort()[::-1]
115
+ keep = []
116
+ while order.size > 0:
117
+ i = order[0]
118
+ keep.append(i)
119
+ xx1 = np.maximum(boxes[i, 0], boxes[order[1:], 0])
120
+ yy1 = np.maximum(boxes[i, 1], boxes[order[1:], 1])
121
+ xx2 = np.minimum(boxes[i, 2], boxes[order[1:], 2])
122
+ yy2 = np.minimum(boxes[i, 3], boxes[order[1:], 3])
123
+ w = np.maximum(0.0, xx2 - xx1)
124
+ h = np.maximum(0.0, yy2 - yy1)
125
+ inter = w * h
126
+ area_i = (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
127
+ area_rest = (boxes[order[1:], 2] - boxes[order[1:], 0]) * (boxes[order[1:], 3] - boxes[order[1:], 1])
128
+ union = np.maximum(area_i + area_rest - inter, 1e-6)
129
+ iou = inter / union
130
+ remaining = np.where(iou <= self.iou_threshold)[0]
131
+ order = order[remaining + 1]
132
+ return [dets[idx] for idx in keep]
133
+
134
+ # ─── Decode: handles both end-to-end and raw formats ─────────────
135
+
136
+ def _decode_end2end(self, out, orig_h, orig_w):
137
+ """Decode end-to-end [1, N, 6] output: [x1,y1,x2,y2,conf,cls_id] in input coords."""
138
+ pred = out[0] # [N, 6]
139
+ if pred.ndim != 2 or pred.shape[1] != 6:
140
+ return []
141
+
142
+ confs = pred[:, 4]
143
+ keep = confs >= self.conf_threshold
144
+ pred = pred[keep]
145
+ if pred.shape[0] == 0:
146
+ return []
147
+
148
+ sx = orig_w / float(self.input_w)
149
+ sy = orig_h / float(self.input_h)
150
+
151
+ results = []
152
+ for i in range(pred.shape[0]):
153
+ x1 = pred[i, 0] * sx
154
+ y1 = pred[i, 1] * sy
155
+ x2 = pred[i, 2] * sx
156
+ y2 = pred[i, 3] * sy
157
+ conf = float(pred[i, 4])
158
+ cls_id = int(pred[i, 5])
159
+ results.append((x1, y1, x2, y2, conf, cls_id))
160
+ return results
161
+
162
+ def _decode_raw(self, out, orig_h, orig_w):
163
+ """Decode raw [1, 4+nc, N] or [1, N, 4+nc] output."""
164
+ pred = out[0]
165
+ if pred.ndim != 2:
166
+ return []
167
+ if pred.shape[0] < pred.shape[1]:
168
+ pred = pred.T
169
+ if pred.shape[1] < 5:
170
+ return []
171
+
172
+ boxes = pred[:, :4]
173
+ cls_scores = pred[:, 4:]
174
+ if cls_scores.shape[1] == 0:
175
+ return []
176
+
177
+ cls_ids = np.argmax(cls_scores, axis=1)
178
+ confs = np.max(cls_scores, axis=1)
179
+ keep = confs >= self.conf_threshold
180
+ boxes, confs, cls_ids = boxes[keep], confs[keep], cls_ids[keep]
181
+ if boxes.shape[0] == 0:
182
+ return []
183
+
184
+ sx = orig_w / float(self.input_w)
185
+ sy = orig_h / float(self.input_h)
186
+
187
+ dets = []
188
+ for i in range(boxes.shape[0]):
189
+ cx, cy, bw, bh = boxes[i].tolist()
190
+ x1 = (cx - bw / 2.0) * sx
191
+ y1 = (cy - bh / 2.0) * sy
192
+ x2 = (cx + bw / 2.0) * sx
193
+ y2 = (cy + bh / 2.0) * sy
194
+ dets.append((x1, y1, x2, y2, float(confs[i]), int(cls_ids[i])))
195
+ return self._nms(dets)
196
+
197
+ # ─── Geometric validation ────────────────────────────────────────
198
+
199
+ def _geometric_validate(self, dets, orig_h, orig_w):
200
+ """Suppress detections that fail basic geometric expectations.
201
+
202
+ Returns list with penalties applied to conf.
203
+ - Canopy: must be wide (aspect w/h >= 0.8)
204
+ - Pump: must not be extremely wide (aspect w/h <= 4.0)
205
+ - Price board: must be small (area <= 15% of image)
206
+ - Hose: must not be tiny (area >= 0.05% of image)
207
+ """
208
+ img_area = max(orig_h * orig_w, 1)
209
+ result = []
210
+ for x1, y1, x2, y2, conf, cls_id in dets:
211
+ bw = max(x2 - x1, 1)
212
+ bh = max(y2 - y1, 1)
213
+ aspect = bw / bh
214
+ box_area = bw * bh
215
+ area_frac = box_area / img_area
216
+ penalty = 0.0
217
+
218
+ if cls_id == CLS_CANOPY:
219
+ if aspect < CANOPY_MIN_ASPECT:
220
+ penalty = GEOMETRIC_SUPPRESS_PENALTY
221
+ elif cls_id == CLS_PUMP:
222
+ if aspect > PUMP_MAX_ASPECT:
223
+ penalty = GEOMETRIC_SUPPRESS_PENALTY
224
+ elif cls_id == CLS_PRICEBOARD:
225
+ if area_frac > PRICEBOARD_MAX_AREA_FRAC:
226
+ penalty = GEOMETRIC_SUPPRESS_PENALTY
227
+ elif cls_id == CLS_HOSE:
228
+ if area_frac < HOSE_MIN_AREA_FRAC:
229
+ penalty = GEOMETRIC_SUPPRESS_PENALTY
230
+
231
+ new_conf = max(0.0, conf - penalty)
232
+ if new_conf >= self.conf_threshold:
233
+ result.append((x1, y1, x2, y2, new_conf, cls_id))
234
+ return result
235
+
236
+ # ─── Spatial co-occurrence scoring ───────────────────────────────
237
+
238
+ def _spatial_cooccurrence(self, dets, orig_h, orig_w):
239
+ """Adjust confidences based on spatial co-occurrence patterns.
240
+
241
+ Boosts:
242
+ - Pump near canopy: both get +0.05
243
+ - Pump near hose: hose gets +0.08
244
+ - Canopy near hose: hose gets +0.05
245
+
246
+ Suppresses:
247
+ - Low-conf detection with no neighbors of expected class: -0.03
248
+ (except price boards, which are 91% solo in training data)
249
+ """
250
+ if not dets:
251
+ return dets
252
+
253
+ n = len(dets)
254
+ adjustments = [0.0] * n
255
+ diag = math.sqrt(orig_h ** 2 + orig_w ** 2)
256
+ prox = COOCCUR_PROXIMITY * diag # absolute pixel distance
257
+
258
+ # Precompute centers
259
+ centers = []
260
+ for x1, y1, x2, y2, conf, cls_id in dets:
261
+ centers.append(((x1 + x2) / 2, (y1 + y2) / 2))
262
+
263
+ # Build per-class index
264
+ cls_map = {}
265
+ for i, (_, _, _, _, _, cls_id) in enumerate(dets):
266
+ cls_map.setdefault(cls_id, []).append(i)
267
+
268
+ def near(i, j):
269
+ dx = centers[i][0] - centers[j][0]
270
+ dy = centers[i][1] - centers[j][1]
271
+ return math.sqrt(dx * dx + dy * dy) < prox
272
+
273
+ # Pump + Canopy boost
274
+ for pi in cls_map.get(CLS_PUMP, []):
275
+ for ci in cls_map.get(CLS_CANOPY, []):
276
+ if near(pi, ci):
277
+ adjustments[pi] = max(adjustments[pi], COOCCUR_BOOST_PUMP_CANOPY)
278
+ adjustments[ci] = max(adjustments[ci], COOCCUR_BOOST_PUMP_CANOPY)
279
+
280
+ # Pump + Hose boost (hose gets larger boost)
281
+ for pi in cls_map.get(CLS_PUMP, []):
282
+ for hi in cls_map.get(CLS_HOSE, []):
283
+ if near(pi, hi):
284
+ adjustments[hi] = max(adjustments[hi], COOCCUR_BOOST_PUMP_HOSE)
285
+
286
+ # Canopy + Hose boost
287
+ for ci in cls_map.get(CLS_CANOPY, []):
288
+ for hi in cls_map.get(CLS_HOSE, []):
289
+ if near(ci, hi):
290
+ adjustments[hi] = max(adjustments[hi], COOCCUR_BOOST_CANOPY_HOSE)
291
+
292
+ # Suppress isolated low-confidence detections (not price boards)
293
+ for i, (x1, y1, x2, y2, conf, cls_id) in enumerate(dets):
294
+ if cls_id == CLS_PRICEBOARD:
295
+ continue # price boards are often solo (91% in training)
296
+ if conf > 0.60:
297
+ continue # high confidence β€” don't suppress
298
+
299
+ has_neighbor = False
300
+ for j in range(n):
301
+ if i == j:
302
+ continue
303
+ if near(i, j):
304
+ has_neighbor = True
305
+ break
306
+ if not has_neighbor:
307
+ adjustments[i] = min(adjustments[i],
308
+ adjustments[i] - COOCCUR_SUPPRESS_ISOLATED)
309
+
310
+ # Apply adjustments
311
+ result = []
312
+ for i, (x1, y1, x2, y2, conf, cls_id) in enumerate(dets):
313
+ new_conf = min(1.0, max(0.0, conf + adjustments[i]))
314
+ if new_conf >= self.conf_threshold:
315
+ result.append((x1, y1, x2, y2, new_conf, cls_id))
316
+ return result
317
+
318
+ # ─── Main inference ──────────────────────────────────────────────
319
+
320
+ def _infer_single(self, image_bgr: ndarray) -> list[BoundingBox]:
321
+ inp, (orig_h, orig_w) = self._preprocess(image_bgr)
322
+ out = self.session.run(None, {self.input_name: inp})[0]
323
+
324
+ # Decode based on detected output format
325
+ if self._end2end:
326
+ dets = self._decode_end2end(out, orig_h, orig_w)
327
+ else:
328
+ dets = self._decode_raw(out, orig_h, orig_w)
329
+
330
+ if not dets:
331
+ return []
332
+
333
+ # Post-processing pipeline
334
+ dets = self._geometric_validate(dets, orig_h, orig_w)
335
+ dets = self._spatial_cooccurrence(dets, orig_h, orig_w)
336
+
337
+ # Convert to BoundingBox
338
+ out_boxes = []
339
+ for x1, y1, x2, y2, conf, cls_id in dets:
340
+ ix1 = max(0, min(orig_w, math.floor(x1)))
341
+ iy1 = max(0, min(orig_h, math.floor(y1)))
342
+ ix2 = max(0, min(orig_w, math.ceil(x2)))
343
+ iy2 = max(0, min(orig_h, math.ceil(y2)))
344
+ out_boxes.append(
345
+ BoundingBox(
346
+ x1=ix1, y1=iy1, x2=ix2, y2=iy2,
347
+ cls_id=cls_id,
348
+ conf=max(0.0, min(1.0, conf)),
349
+ )
350
+ )
351
+ return out_boxes
352
+
353
+ def predict_batch(
354
+ self,
355
+ batch_images: list[ndarray],
356
+ offset: int,
357
+ n_keypoints: int,
358
+ ) -> list[TVFrameResult]:
359
+ results = []
360
+ for idx, image in enumerate(batch_images):
361
+ boxes = self._infer_single(image)
362
+ keypoints = [(0, 0) for _ in range(max(0, int(n_keypoints)))]
363
+ results.append(
364
+ TVFrameResult(
365
+ frame_id=offset + idx,
366
+ boxes=boxes,
367
+ keypoints=keypoints,
368
+ )
369
+ )
370
+ return results
model_type.json ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ {
2
+ "task_type": "object-detection",
3
+ "model_type": "yolov11-nano"
4
+ }
weights.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:68cbfb0b6c19ae57ba3b724eb4380cdfb61b4558f2aa8722b6f3555ef0b267a2
3
+ size 19155617