licensy commited on
Commit
2f8bd3c
·
verified ·
1 Parent(s): 659aee3

YOLO26s + SAHI 2x2 tiling, c=0.20, shutdown=14400

Browse files
Files changed (2) hide show
  1. miner.py +102 -83
  2. weights.onnx +2 -2
miner.py CHANGED
@@ -36,11 +36,6 @@ class TVFrameResult(BaseModel):
36
 
37
 
38
  class Miner:
39
- """
40
- Auto-generated by subnet_bridge from a Manako element repo.
41
- This miner is intentionally self-contained for chute import restrictions.
42
- """
43
-
44
  def __init__(self, path_hf_repo: Path) -> None:
45
  self.path_hf_repo = path_hf_repo
46
  self.class_names = ['numberplate']
@@ -50,99 +45,123 @@ class Miner:
50
  )
51
  self.input_name = self.session.get_inputs()[0].name
52
  input_shape = self.session.get_inputs()[0].shape
53
- self.input_h = int(input_shape[2])
54
- self.input_w = int(input_shape[3])
55
- self.conf_threshold = 0.15
56
- self.iou_threshold = 0.3
57
 
58
  def __repr__(self) -> str:
59
  return f"ONNX Miner session={type(self.session).__name__} classes={len(self.class_names)}"
60
 
61
- def _preprocess(self, image_bgr: ndarray) -> tuple[np.ndarray, tuple[int, int]]:
62
- h, w = image_bgr.shape[:2]
63
- rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
64
- resized = cv2.resize(rgb, (self.input_w, self.input_h))
65
- x = resized.astype(np.float32) / 255.0
66
- x = np.transpose(x, (2, 0, 1))[None, ...]
67
- return x, (h, w)
68
-
69
- def _normalize_predictions(self, raw: np.ndarray) -> np.ndarray:
70
- pred = raw[0]
71
- if pred.ndim != 2:
72
- raise ValueError(f"Unexpected prediction shape: {raw.shape}")
73
- if pred.shape[0] < pred.shape[1]:
74
- pred = pred.transpose(1, 0)
75
- return pred
76
-
77
- def _nms(self, dets: list[tuple[float, float, float, float, float, int]]) -> list[tuple[float, float, float, float, float, int]]:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
78
  if not dets:
79
  return []
80
- boxes = np.array([[d[0], d[1], d[2], d[3]] for d in dets], dtype=np.float32)
81
- scores = np.array([d[4] for d in dets], dtype=np.float32)
82
- order = scores.argsort()[::-1]
83
  keep = []
84
- while order.size > 0:
85
- i = order[0]
86
- keep.append(i)
87
- xx1 = np.maximum(boxes[i, 0], boxes[order[1:], 0])
88
- yy1 = np.maximum(boxes[i, 1], boxes[order[1:], 1])
89
- xx2 = np.minimum(boxes[i, 2], boxes[order[1:], 2])
90
- yy2 = np.minimum(boxes[i, 3], boxes[order[1:], 3])
91
- w = np.maximum(0.0, xx2 - xx1)
92
- h = np.maximum(0.0, yy2 - yy1)
93
- inter = w * h
94
- area_i = (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
95
- area_rest = (boxes[order[1:], 2] - boxes[order[1:], 0]) * (boxes[order[1:], 3] - boxes[order[1:], 1])
96
- union = np.maximum(area_i + area_rest - inter, 1e-6)
97
- iou = inter / union
98
- remaining = np.where(iou <= self.iou_threshold)[0]
99
- order = order[remaining + 1]
100
- return [dets[idx] for idx in keep]
 
 
101
 
102
  def _infer_single(self, image_bgr: ndarray) -> list[BoundingBox]:
103
- inp, (orig_h, orig_w) = self._preprocess(image_bgr)
104
- out = self.session.run(None, {self.input_name: inp})[0]
105
- pred = self._normalize_predictions(out)
106
- if pred.shape[1] < 5:
107
- return []
108
- boxes = pred[:, :4]
109
- cls_scores = pred[:, 4:]
110
- if cls_scores.shape[1] == 0:
111
- return []
112
- cls_ids = np.argmax(cls_scores, axis=1)
113
- confs = np.max(cls_scores, axis=1)
114
- keep = confs >= self.conf_threshold
115
- boxes = boxes[keep]
116
- confs = confs[keep]
117
- cls_ids = cls_ids[keep]
118
- if boxes.shape[0] == 0:
119
- return []
120
- sx = orig_w / float(self.input_w)
121
- sy = orig_h / float(self.input_h)
122
- dets: list[tuple[float, float, float, float, float, int]] = []
123
- for i in range(boxes.shape[0]):
124
- cx, cy, bw, bh = boxes[i].tolist()
125
- x1 = (cx - bw / 2.0) * sx
126
- y1 = (cy - bh / 2.0) * sy
127
- x2 = (cx + bw / 2.0) * sx
128
- y2 = (cy + bh / 2.0) * sy
129
- dets.append((x1, y1, x2, y2, float(confs[i]), int(cls_ids[i])))
130
- dets = self._nms(dets)
131
- out_boxes: list[BoundingBox] = []
132
- for x1, y1, x2, y2, conf, cls_id in dets:
133
- ix1 = max(0, min(orig_w, math.floor(x1)))
134
- iy1 = max(0, min(orig_h, math.floor(y1)))
135
- ix2 = max(0, min(orig_w, math.ceil(x2)))
136
- iy2 = max(0, min(orig_h, math.ceil(y2)))
 
 
 
 
 
 
137
  out_boxes.append(
138
- BoundingBox(x1=ix1, y1=iy1, x2=ix2, y2=iy2,
139
- cls_id=cls_id, conf=max(0.0, min(1.0, conf))))
 
 
 
 
140
  return out_boxes
141
 
142
  def predict_batch(
143
  self, batch_images: list[ndarray], offset: int, n_keypoints: int,
144
  ) -> list[TVFrameResult]:
145
- results: list[TVFrameResult] = []
146
  for idx, image in enumerate(batch_images):
147
  boxes = self._infer_single(image)
148
  keypoints = [(0, 0) for _ in range(max(0, int(n_keypoints)))]
 
36
 
37
 
38
  class Miner:
 
 
 
 
 
39
  def __init__(self, path_hf_repo: Path) -> None:
40
  self.path_hf_repo = path_hf_repo
41
  self.class_names = ['numberplate']
 
45
  )
46
  self.input_name = self.session.get_inputs()[0].name
47
  input_shape = self.session.get_inputs()[0].shape
48
+ self.input_size = int(input_shape[2]) # 1280
49
+ self.conf_threshold = 0.20
50
+ self.iou_threshold = 0.5
51
+ self.tile_overlap = 0.2
52
 
53
  def __repr__(self) -> str:
54
  return f"ONNX Miner session={type(self.session).__name__} classes={len(self.class_names)}"
55
 
56
+ def _letterbox(self, img: ndarray) -> tuple[ndarray, float, int, int]:
57
+ h, w = img.shape[:2]
58
+ r = min(self.input_size / h, self.input_size / w)
59
+ nw, nh = int(w * r), int(h * r)
60
+ resized = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_LINEAR)
61
+ canvas = np.full((self.input_size, self.input_size, 3), 114, dtype=np.uint8)
62
+ dx = (self.input_size - nw) // 2
63
+ dy = (self.input_size - nh) // 2
64
+ canvas[dy:dy+nh, dx:dx+nw] = resized
65
+ return canvas, r, dx, dy
66
+
67
+ def _run_single(self, img: ndarray) -> list[tuple[float, float, float, float, float]]:
68
+ h, w = img.shape[:2]
69
+ canvas, r, dx, dy = self._letterbox(img)
70
+ blob = (canvas.astype(np.float32) / 255.0).transpose(2, 0, 1)[np.newaxis]
71
+ out = self.session.run(None, {self.input_name: blob})[0][0]
72
+ dets = []
73
+ for row in out:
74
+ x1, y1, x2, y2, conf, cls = row
75
+ if conf < self.conf_threshold:
76
+ continue
77
+ dets.append((
78
+ float(conf),
79
+ (x1 - dx) / r,
80
+ (y1 - dy) / r,
81
+ (x2 - dx) / r,
82
+ (y2 - dy) / r,
83
+ ))
84
+ return dets
85
+
86
+ def _nms(self, dets: list[tuple[float, float, float, float, float]]) -> list[tuple[float, float, float, float, float]]:
87
  if not dets:
88
  return []
89
+ dets.sort(key=lambda x: -x[0])
 
 
90
  keep = []
91
+ used = [False] * len(dets)
92
+ for i in range(len(dets)):
93
+ if used[i]:
94
+ continue
95
+ keep.append(dets[i])
96
+ for j in range(i + 1, len(dets)):
97
+ if used[j]:
98
+ continue
99
+ # compute IoU
100
+ ax1, ay1, ax2, ay2 = dets[i][1], dets[i][2], dets[i][3], dets[i][4]
101
+ bx1, by1, bx2, by2 = dets[j][1], dets[j][2], dets[j][3], dets[j][4]
102
+ ix1 = max(ax1, bx1); iy1 = max(ay1, by1)
103
+ ix2 = min(ax2, bx2); iy2 = min(ay2, by2)
104
+ inter = max(0, ix2-ix1) * max(0, iy2-iy1)
105
+ aa = (ax2-ax1)*(ay2-ay1); bb = (bx2-bx1)*(by2-by1)
106
+ iou = inter / (aa + bb - inter + 1e-6)
107
+ if iou > self.iou_threshold:
108
+ used[j] = True
109
+ return keep
110
 
111
  def _infer_single(self, image_bgr: ndarray) -> list[BoundingBox]:
112
+ orig_h, orig_w = image_bgr.shape[:2]
113
+ all_dets = []
114
+
115
+ # Full image pass
116
+ all_dets.extend(self._run_single(image_bgr))
117
+
118
+ # 2x2 tile passes
119
+ tw = orig_w // 2
120
+ th = orig_h // 2
121
+ ox = int(tw * self.tile_overlap)
122
+ oy = int(th * self.tile_overlap)
123
+ tiles = [
124
+ (0, 0, tw + ox, th + oy),
125
+ (tw - ox, 0, orig_w, th + oy),
126
+ (0, th - oy, tw + ox, orig_h),
127
+ (tw - ox, th - oy, orig_w, orig_h),
128
+ ]
129
+ for tx1, ty1, tx2, ty2 in tiles:
130
+ tx1 = max(0, tx1); ty1 = max(0, ty1)
131
+ tx2 = min(orig_w, tx2); ty2 = min(orig_h, ty2)
132
+ crop = image_bgr[ty1:ty2, tx1:tx2]
133
+ tile_dets = self._run_single(crop)
134
+ for conf, x1, y1, x2, y2 in tile_dets:
135
+ all_dets.append((conf, x1 + tx1, y1 + ty1, x2 + tx1, y2 + ty1))
136
+
137
+ # NMS to merge overlapping detections
138
+ all_dets = self._nms(all_dets)
139
+
140
+ out_boxes = []
141
+ for conf, x1, y1, x2, y2 in all_dets:
142
+ bx1 = max(0, min(orig_w, math.floor(x1)))
143
+ by1 = max(0, min(orig_h, math.floor(y1)))
144
+ bx2 = max(0, min(orig_w, math.ceil(x2)))
145
+ by2 = max(0, min(orig_h, math.ceil(y2)))
146
+ bw = bx2 - bx1
147
+ bh = by2 - by1
148
+ if bw < 6 or bh < 6 or bw * bh < 80:
149
+ continue
150
+ if max(bw / max(bh, 1), bh / max(bw, 1)) > 10:
151
+ continue
152
  out_boxes.append(
153
+ BoundingBox(
154
+ x1=bx1, y1=by1, x2=bx2, y2=by2,
155
+ cls_id=0,
156
+ conf=max(0.0, min(1.0, conf)),
157
+ )
158
+ )
159
  return out_boxes
160
 
161
  def predict_batch(
162
  self, batch_images: list[ndarray], offset: int, n_keypoints: int,
163
  ) -> list[TVFrameResult]:
164
+ results = []
165
  for idx, image in enumerate(batch_images):
166
  boxes = self._infer_single(image)
167
  keypoints = [(0, 0) for _ in range(max(0, int(n_keypoints)))]
weights.onnx CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:97d257baf273fe85d6b51ea547b59dd36d2a67a60fa64e30c9c15ce7d7949995
3
- size 11966914
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:16b083d01cf9fff412433a8d3bb1bd9253aa29cefabb05a2afc83c5fe55c520d
3
+ size 19405465