| """ScoreVision crime detector v27 — YOLOv11s trained on 4-voter consensus |
| (new vmodel0, hermes_sv, pmodel_9, SAM3-refined; >=2 agree, IoU>=0.5). |
| Per-class confidence + per-class rescue (bonus) logic, flip TTA, cross-class NMS.""" |
| from pathlib import Path |
| import math |
|
|
| import cv2 |
| import numpy as np |
| import onnxruntime as ort |
| from numpy import ndarray |
| from pydantic import BaseModel |
|
|
|
|
| class BoundingBox(BaseModel): |
| x1: int |
| y1: int |
| x2: int |
| y2: int |
| cls_id: int |
| conf: float |
|
|
|
|
| class TVFrameResult(BaseModel): |
| frame_id: int |
| boxes: list[BoundingBox] |
| keypoints: list[tuple[int, int]] |
|
|
|
|
| class Miner: |
| """YOLOv11s + per-class threshold + per-class rescue + flip TTA + cross-class dedup.""" |
|
|
| |
| class_names = ["balaclava", "hoodie", "glove", "bat", "spray paint", "graffiti"] |
| input_size = 1280 |
| iou_thres = 0.30 |
| cross_iou_thresh = 0.80 |
| max_aspect_ratio = 10.0 |
| max_det = 150 |
| _conf_thres_array = np.array( |
| [0.30, 0.70, 0.70, 0.50, 0.50, 0.40], dtype=np.float32 |
| ) |
| _bonus_array = np.array( |
| [0.20, 0.25, 0.20, 0.30, 0.20, 0.15], dtype=np.float32 |
| ) |
|
|
| _glove_cls_id = 2 |
| _glove_spurious_max_conf = 0.04 |
| _glove_spurious_aspect = 2.5 |
| _glove_crowd_min_count = 5 |
| _glove_crowd_max_conf_trigger = 0.6 |
| _glove_crowd_keep_min_conf = 0.27 |
| _glove_split_low_max = 0.26 |
| _glove_split_low_min_count = 4 |
| _glove_split_high_min_count = 2 |
| _glove_pair_low_max = 0.1 |
| _glove_pair_high_min = 0.4 |
| _glove_pair_max_center_x_dist = 0.028 |
| _glove_reject_r_min = 100 |
| _glove_reject_r_max = 125 |
| _glove_reject_g_min = 88 |
| _glove_reject_g_max = 115 |
| _glove_reject_b_min = 68 |
| _glove_reject_b_max = 88 |
| _glove_reject_min_pixel_frac = 0.04 |
| _glove_reject_p90_percentile = 90.0 |
|
|
| _spray_paint_cls_id = 4 |
| _spray_orphan_max_conf = 0.3 |
| _spray_glove_min_overlap = 0.3 |
|
|
| def __init__(self, path_hf_repo: Path) -> None: |
| model_path = path_hf_repo / "weights.onnx" |
| print("ORT version:", ort.__version__) |
| try: |
| ort.preload_dlls() |
| print("preload_dlls success") |
| except Exception as e: |
| print(f"preload_dlls failed: {e}") |
| print("ORT providers:", ort.get_available_providers()) |
|
|
| opts = ort.SessionOptions() |
| opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| try: |
| self.session = ort.InferenceSession( |
| str(model_path), sess_options=opts, |
| providers=["CUDAExecutionProvider", "CPUExecutionProvider"], |
| ) |
| print("CUDA session created") |
| except Exception as e: |
| print(f"CUDA failed, CPU fallback: {e}") |
| self.session = ort.InferenceSession( |
| str(model_path), sess_options=opts, providers=["CPUExecutionProvider"] |
| ) |
| print("session providers:", self.session.get_providers()) |
|
|
| for inp in self.session.get_inputs(): |
| print("INPUT:", inp.name, inp.shape, inp.type) |
| for out in self.session.get_outputs(): |
| print("OUTPUT:", out.name, out.shape, out.type) |
|
|
| self.input_name = self.session.get_inputs()[0].name |
| self.output_names = [o.name for o in self.session.get_outputs()] |
| sh = self.session.get_inputs()[0].shape |
| self.input_height = self._safe_dim(sh[2], self.input_size) |
| self.input_width = self._safe_dim(sh[3], self.input_size) |
| print(f"ONNX loaded: {model_path}") |
| print( |
| "per-class conf: " |
| + ", ".join(f"{n}={t:.2f}" for n, t in zip(self.class_names, self._conf_thres_array.tolist())) |
| ) |
| print( |
| "per-class rescue bonus: " |
| + ", ".join(f"{n}={t:.2f}" for n, t in zip(self.class_names, self._bonus_array.tolist())) |
| ) |
|
|
| @staticmethod |
| def _safe_dim(v, d): |
| return v if isinstance(v, int) and v > 0 else d |
|
|
| def _letterbox(self, image, new_shape, color=(114, 114, 114)): |
| h, w = image.shape[:2] |
| nw, nh = new_shape |
| r = min(nw / w, nh / h) |
| rw, rh = int(round(w * r)), int(round(h * r)) |
| if (rw, rh) != (w, h): |
| interp = cv2.INTER_CUBIC if r > 1.0 else cv2.INTER_LINEAR |
| image = cv2.resize(image, (rw, rh), interpolation=interp) |
| dw, dh = (nw - rw) / 2.0, (nh - rh) / 2.0 |
| top, bot = int(round(dh - 0.1)), int(round(dh + 0.1)) |
| lt, rt = int(round(dw - 0.1)), int(round(dw + 0.1)) |
| padded = cv2.copyMakeBorder(image, top, bot, lt, rt, cv2.BORDER_CONSTANT, value=color) |
| return padded, r, (dw, dh) |
|
|
| def _preprocess(self, image): |
| h, w = image.shape[:2] |
| img, ratio, pad = self._letterbox(image, (self.input_width, self.input_height)) |
| img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 |
| img = np.transpose(img, (2, 0, 1))[None, ...] |
| return np.ascontiguousarray(img, dtype=np.float32), ratio, pad, (w, h) |
|
|
| @staticmethod |
| def _clip(boxes, size): |
| w, h = size |
| boxes[:, 0] = np.clip(boxes[:, 0], 0, w - 1) |
| boxes[:, 1] = np.clip(boxes[:, 1], 0, h - 1) |
| boxes[:, 2] = np.clip(boxes[:, 2], 0, w - 1) |
| boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1) |
| return boxes |
|
|
| @staticmethod |
| def _hard_nms(boxes, scores, iou_thr): |
| n = len(boxes) |
| if n == 0: return np.array([], dtype=np.intp) |
| order = np.argsort(-scores) |
| keep = [] |
| while len(order) > 0: |
| i = int(order[0]); keep.append(i) |
| if len(order) == 1: break |
| rest = order[1:] |
| xx1 = np.maximum(boxes[i, 0], boxes[rest, 0]) |
| yy1 = np.maximum(boxes[i, 1], boxes[rest, 1]) |
| xx2 = np.minimum(boxes[i, 2], boxes[rest, 2]) |
| yy2 = np.minimum(boxes[i, 3], boxes[rest, 3]) |
| inter = np.maximum(0.0, xx2-xx1)*np.maximum(0.0, yy2-yy1) |
| a_i = max(0.0, boxes[i, 2]-boxes[i, 0])*max(0.0, boxes[i, 3]-boxes[i, 1]) |
| a_r = (np.maximum(0.0, boxes[rest, 2]-boxes[rest, 0]) |
| * np.maximum(0.0, boxes[rest, 3]-boxes[rest, 1])) |
| iou = inter / (a_i+a_r-inter+1e-7) |
| order = rest[iou <= iou_thr] |
| return np.array(keep, dtype=np.intp) |
|
|
| def _per_class_hard_nms(self, boxes, scores, cls_ids, iou_thr): |
| if len(boxes) == 0: return np.array([], dtype=np.intp) |
| keep_all = [] |
| for c in np.unique(cls_ids): |
| mask = cls_ids == c |
| idx = np.where(mask)[0] |
| keep = self._hard_nms(boxes[mask], scores[mask], iou_thr) |
| keep_all.extend(idx[keep].tolist()) |
| keep_all.sort() |
| return np.array(keep_all, dtype=np.intp) |
|
|
| def _cross_class_dedup(self, boxes, scores, cls_ids, iou_thr): |
| n = len(boxes) |
| if n == 0: return np.array([], dtype=np.intp) |
| order = np.argsort(-scores) |
| keep = []; suppressed = np.zeros(n, dtype=bool) |
| for i in order: |
| if suppressed[i]: continue |
| keep.append(int(i)) |
| ix1 = np.maximum(boxes[i, 0], boxes[:, 0]) |
| iy1 = np.maximum(boxes[i, 1], boxes[:, 1]) |
| ix2 = np.minimum(boxes[i, 2], boxes[:, 2]) |
| iy2 = np.minimum(boxes[i, 3], boxes[:, 3]) |
| inter = np.maximum(0.0, ix2-ix1)*np.maximum(0.0, iy2-iy1) |
| a_i = max(0.0, boxes[i, 2]-boxes[i, 0])*max(0.0, boxes[i, 3]-boxes[i, 1]) |
| a_r = (np.maximum(0.0, boxes[:, 2]-boxes[:, 0]) |
| * np.maximum(0.0, boxes[:, 3]-boxes[:, 1])) |
| iou = inter / (a_i+a_r-inter+1e-7) |
| iou[i] = 0.0 |
| suppressed |= iou >= iou_thr |
| return np.array(keep, dtype=np.intp) |
|
|
| def _filter_sane(self, boxes, scores, cls_ids, orig_size): |
| if len(boxes) == 0: return boxes, scores, cls_ids |
| w, h = orig_size |
| area_img = float(w * h) |
| bw = np.maximum(0.0, boxes[:, 2]-boxes[:, 0]) |
| bh = np.maximum(0.0, boxes[:, 3]-boxes[:, 1]) |
| area = bw * bh |
| ar = np.where((bw > 0) & (bh > 0), |
| np.maximum(bw / np.maximum(bh, 1e-6), bh / np.maximum(bw, 1e-6)), |
| np.inf) |
| keep = (area <= 0.95 * area_img) & (ar <= self.max_aspect_ratio) |
| return boxes[keep], scores[keep], cls_ids[keep] |
|
|
| def _conf_filter_with_rescue(self, boxes, scores, cls_ids): |
| """Per-class threshold + per-class rescue (admit top-1 if score >= thr - bonus |
| and no box of that class passed the normal threshold).""" |
| if len(scores) == 0: |
| return np.zeros(0, dtype=bool) |
| thr = self._conf_thres_array[cls_ids] |
| keep = scores >= thr |
| for c in np.unique(cls_ids): |
| bonus = float(self._bonus_array[c]) |
| if bonus <= 0.0: continue |
| cm = cls_ids == c |
| if keep[cm].any(): continue |
| idx = np.where(cm)[0] |
| top = int(idx[int(np.argmax(scores[idx]))]) |
| if scores[top] >= self._conf_thres_array[c] - bonus: |
| keep[top] = True |
| return keep |
|
|
| def _decode(self, preds, ratio, pad, orig_size): |
| if preds.ndim == 3 and preds.shape[0] == 1: |
| preds = preds[0] |
| if preds.ndim != 2 or preds.shape[1] < 6: |
| return (np.empty((0, 4), dtype=np.float32), |
| np.empty((0,), dtype=np.float32), |
| np.empty((0,), dtype=np.int32)) |
| boxes = preds[:, :4].astype(np.float32) |
| scores = preds[:, 4].astype(np.float32) |
| cls_ids = preds[:, 5].astype(np.int32) |
| keep = (scores > 0) & (cls_ids >= 0) & (cls_ids < len(self.class_names)) |
| boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep] |
| if len(boxes) == 0: return boxes, scores, cls_ids |
| |
| keep = self._conf_filter_with_rescue(boxes, scores, cls_ids) |
| boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep] |
| if len(boxes) == 0: return boxes, scores, cls_ids |
| pad_w, pad_h = pad |
| boxes[:, [0, 2]] -= pad_w |
| boxes[:, [1, 3]] -= pad_h |
| boxes /= ratio |
| boxes = self._clip(boxes, orig_size) |
| return boxes, scores, cls_ids |
|
|
| @classmethod |
| def _is_spurious_glove(cls, cls_id, x1, y1, x2, y2, conf): |
| if cls_id != cls._glove_cls_id: return False |
| w, h = x2 - x1, y2 - y1 |
| if w <= 0: return False |
| return h > cls._glove_spurious_aspect * w and conf < cls._glove_spurious_max_conf |
|
|
| @staticmethod |
| def _box_intersection_area(a, b): |
| ix1, iy1 = max(a.x1, b.x1), max(a.y1, b.y1) |
| ix2, iy2 = min(a.x2, b.x2), min(a.y2, b.y2) |
| return max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1) |
|
|
| @classmethod |
| def _spray_overlaps_glove(cls, spray, glove): |
| inter = cls._box_intersection_area(spray, glove) |
| area = max(0.0, float(spray.x2 - spray.x1) * float(spray.y2 - spray.y1)) |
| if area <= 0.0: return False |
| return inter / area >= cls._spray_glove_min_overlap |
|
|
| @classmethod |
| def _is_glove_rejection_color(cls, cls_id, image, x1, y1, x2, y2): |
| if cls_id != cls._glove_cls_id or image is None: return False |
| H, W = image.shape[:2] |
| xi1 = max(0, int(math.floor(x1))) |
| yi1 = max(0, int(math.floor(y1))) |
| xi2 = min(W, int(math.ceil(x2))) |
| yi2 = min(H, int(math.ceil(y2))) |
| if xi2 <= xi1 or yi2 <= yi1: return False |
| crop = image[yi1:yi2, xi1:xi2] |
| if crop.size == 0: return False |
| b = crop[:, :, 0].astype(np.float32) |
| g = crop[:, :, 1].astype(np.float32) |
| r = crop[:, :, 2].astype(np.float32) |
| in_range = ((r >= cls._glove_reject_r_min) & (r <= cls._glove_reject_r_max) |
| & (g >= cls._glove_reject_g_min) & (g <= cls._glove_reject_g_max) |
| & (b >= cls._glove_reject_b_min) & (b <= cls._glove_reject_b_max)) |
| if float(in_range.mean()) >= cls._glove_reject_min_pixel_frac: |
| return True |
| p90_r = float(np.percentile(r, cls._glove_reject_p90_percentile)) |
| p90_g = float(np.percentile(g, cls._glove_reject_p90_percentile)) |
| p90_b = float(np.percentile(b, cls._glove_reject_p90_percentile)) |
| return (cls._glove_reject_r_min <= p90_r <= cls._glove_reject_r_max |
| and cls._glove_reject_g_min <= p90_g <= cls._glove_reject_g_max |
| and cls._glove_reject_b_min <= p90_b <= cls._glove_reject_b_max) |
|
|
| @classmethod |
| def _filter_glove_by_color(cls, results, image): |
| if image is None: return results |
| return [b for b in results |
| if not cls._is_glove_rejection_color(b.cls_id, image, b.x1, b.y1, b.x2, b.y2)] |
|
|
| @staticmethod |
| def _glove_cx(b): return (b.x1 + b.x2) / 2.0 |
|
|
| @classmethod |
| def _filter_paired_low_glove(cls, results, image_width): |
| if image_width <= 0: return results |
| gloves = [b for b in results if b.cls_id == cls._glove_cls_id] |
| if len(gloves) < 2: return results |
| remove = set() |
| for i, a in enumerate(gloves): |
| for b in gloves[i+1:]: |
| if a.conf < cls._glove_pair_low_max and b.conf > cls._glove_pair_high_min: |
| low, high = a, b |
| elif b.conf < cls._glove_pair_low_max and a.conf > cls._glove_pair_high_min: |
| low, high = b, a |
| else: |
| continue |
| cx_dist = abs(cls._glove_cx(low) - cls._glove_cx(high)) / float(image_width) |
| if cx_dist < cls._glove_pair_max_center_x_dist: |
| remove.add(id(low)) |
| return [b for b in results if id(b) not in remove] |
|
|
| @classmethod |
| def _filter_weak_glove_with_strong_others(cls, results): |
| others = [b for b in results if b.cls_id != cls._glove_cls_id] |
| if not others: return results |
| if not all(b.conf > cls._glove_pair_high_min for b in others): return results |
| return [b for b in results |
| if b.cls_id != cls._glove_cls_id or b.conf >= cls._glove_pair_low_max] |
|
|
| @classmethod |
| def _filter_split_glove_confidence(cls, results): |
| gloves = [b for b in results if b.cls_id == cls._glove_cls_id] |
| low = [b for b in gloves if b.conf < cls._glove_split_low_max] |
| high = [b for b in gloves if b.conf > cls._glove_split_low_max] |
| if len(low) <= cls._glove_split_low_min_count: return results |
| if len(high) < cls._glove_split_high_min_count: return results |
| best_low = max(low, key=lambda b: b.conf) |
| others = [b for b in results if b.cls_id != cls._glove_cls_id] |
| return others + high + [best_low] |
|
|
| @classmethod |
| def _filter_crowded_gloves(cls, results): |
| gloves = [b for b in results if b.cls_id == cls._glove_cls_id] |
| if len(gloves) <= cls._glove_crowd_min_count: return results |
| if max(b.conf for b in gloves) <= cls._glove_crowd_max_conf_trigger: return results |
| return [b for b in results |
| if b.cls_id != cls._glove_cls_id or b.conf > cls._glove_crowd_keep_min_conf] |
|
|
| @classmethod |
| def _filter_orphan_low_conf_spray(cls, results): |
| gloves = [b for b in results if b.cls_id == cls._glove_cls_id] |
| kept = [] |
| for box in results: |
| if box.cls_id != cls._spray_paint_cls_id or box.conf >= cls._spray_orphan_max_conf: |
| kept.append(box); continue |
| if any(cls._spray_overlaps_glove(box, g) for g in gloves): |
| kept.append(box) |
| return kept |
|
|
| @classmethod |
| def _box_iou_pair(cls, a, b): |
| inter = cls._box_intersection_area(a, b) |
| aw = max(0.0, float(a.x2 - a.x1)) * max(0.0, float(a.y2 - a.y1)) |
| bw = max(0.0, float(b.x2 - b.x1)) * max(0.0, float(b.y2 - b.y1)) |
| union = aw + bw - inter |
| if union <= 0: return 0.0 |
| return inter / union |
|
|
| @classmethod |
| def _aggressive_dedup_balaclava_graffiti(cls, results, iou_thr=0.10): |
| """For balaclava (cls 0) and graffiti (cls 5), drop lower-conf box |
| whenever any two boxes of the same class overlap > iou_thr. |
| Rationale: one face = one balaclava; one tag = one graffiti box.""" |
| target = {0, 5} |
| others = [b for b in results if b.cls_id not in target] |
| by_cls = {} |
| for b in results: |
| if b.cls_id in target: |
| by_cls.setdefault(b.cls_id, []).append(b) |
| kept = list(others) |
| for cid, boxes in by_cls.items(): |
| boxes_sorted = sorted(boxes, key=lambda b: -b.conf) |
| survivors = [] |
| for b in boxes_sorted: |
| if all(cls._box_iou_pair(b, k) < iou_thr for k in survivors): |
| survivors.append(b) |
| kept.extend(survivors) |
| return kept |
|
|
| @classmethod |
| def _apply_rules(cls, results, image_width, image=None): |
| results = cls._filter_glove_by_color(results, image) |
| results = cls._filter_paired_low_glove(results, image_width) |
| results = cls._filter_weak_glove_with_strong_others(results) |
| results = cls._filter_split_glove_confidence(results) |
| results = cls._filter_crowded_gloves(results) |
| results = cls._filter_orphan_low_conf_spray(results) |
| return cls._aggressive_dedup_balaclava_graffiti(results) |
|
|
| def _predict_single(self, image): |
| x, ratio, pad, orig_size = self._preprocess(image) |
| out = self.session.run(self.output_names, {self.input_name: x})[0] |
| return self._decode(out, ratio, pad, orig_size) |
|
|
| def _predict_tta(self, image): |
| b0, s0, c0 = self._predict_single(image) |
| flipped = cv2.flip(image, 1) |
| bf, sf, cf = self._predict_single(flipped) |
| if len(bf): |
| w = image.shape[1] |
| bf2 = bf.copy() |
| bf2[:, 0] = w - bf[:, 2] |
| bf2[:, 2] = w - bf[:, 0] |
| bf = bf2 |
| boxes = np.concatenate([b0, bf], axis=0) if len(b0) or len(bf) else b0 |
| scores = np.concatenate([s0, sf], axis=0) if len(b0) or len(bf) else s0 |
| cls_ids = np.concatenate([c0, cf], axis=0) if len(b0) or len(bf) else c0 |
| if len(boxes) == 0: return [] |
| boxes, scores, cls_ids = self._filter_sane(boxes, scores, cls_ids, |
| (image.shape[1], image.shape[0])) |
| if len(boxes) == 0: return [] |
| keep = self._per_class_hard_nms(boxes, scores, cls_ids, self.iou_thres) |
| if len(keep) == 0: return [] |
| boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep] |
| keep = self._cross_class_dedup(boxes, scores, cls_ids, self.cross_iou_thresh) |
| boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep] |
| if len(scores) > self.max_det: |
| top = np.argsort(-scores)[:self.max_det] |
| boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top] |
| results = [] |
| for b, s, c in zip(boxes, scores, cls_ids): |
| if b[2] <= b[0] or b[3] <= b[1]: continue |
| x1, y1, x2, y2 = float(b[0]), float(b[1]), float(b[2]), float(b[3]) |
| if self._is_spurious_glove(int(c), x1, y1, x2, y2, float(s)): continue |
| results.append(BoundingBox( |
| x1=int(math.floor(x1)), y1=int(math.floor(y1)), |
| x2=int(math.ceil(x2)), y2=int(math.ceil(y2)), |
| cls_id=int(c), conf=float(s), |
| )) |
| return self._apply_rules(results, image.shape[1], image=image) |
|
|
| def predict_batch(self, batch_images, offset, n_keypoints): |
| results = [] |
| for j, img in enumerate(batch_images): |
| try: |
| boxes = self._predict_tta(img) |
| except Exception as e: |
| print(f"Inference failed for frame {offset + j}: {e}") |
| boxes = [] |
| results.append( |
| TVFrameResult( |
| frame_id=offset + j, boxes=boxes, |
| keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))], |
| ) |
| ) |
| return results |
|
|
