| |
| """Shared ECSeg preprocessing + output post-processing, byte-faithful to the app. |
| |
| The preprocessing here reproduces exactly what the browser app feeds the ONNX graph, so that |
| a Python correctness comparison feeds baseline and candidate the SAME `images` tensor the |
| product would: |
| |
| * OpenCV path mirrors `packages/smart-tools/src/segment-anything/pre-processing.ts`: |
| - decode → RGB (drop alpha) |
| - stretch resize to 640×640 with cv2.INTER_LINEAR (bilinear), aspect NOT preserved |
| - scale by 1/255 (convertTo CV_32F) |
| - ImageNet normalize: (x - mean) / std, mean=[0.485,0.456,0.406] std=[0.229,0.224,0.225] |
| - NCHW float32, shape [1,3,640,640] |
| |
| The post-processing mirrors the `edgecrafter-seg` parser + mask decode: |
| * score filter at confidenceThreshold (0.4) |
| * per-instance 160×160 logit map → bilinear upscale to image size → threshold at maskThreshold |
| (raw LOGIT cut, default 0.0) → binary mask |
| """ |
|
|
| from __future__ import annotations |
|
|
| import glob |
| import os |
| from typing import List, Tuple |
|
|
| import cv2 |
| import numpy as np |
|
|
| SIZE = 640 |
| MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) |
| STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) |
|
|
| CONF_THRESHOLD = 0.4 |
| MASK_THRESHOLD = 0.0 |
|
|
|
|
| def preprocess_bgr(bgr: np.ndarray) -> np.ndarray: |
| """cv2-decoded BGR image (H,W,3 uint8) -> float32 NCHW [1,3,640,640] tensor.""" |
| rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) |
| resized = cv2.resize(rgb, (SIZE, SIZE), interpolation=cv2.INTER_LINEAR) |
| x = resized.astype(np.float32) / 255.0 |
| x = (x - MEAN) / STD |
| |
| x = np.transpose(x, (2, 0, 1))[None, ...] |
| return np.ascontiguousarray(x, dtype=np.float32) |
|
|
|
|
| def preprocess_file(path: str) -> np.ndarray: |
| bgr = cv2.imread(path, cv2.IMREAD_COLOR) |
| if bgr is None: |
| raise ValueError(f"cv2 failed to read {path}") |
| return preprocess_bgr(bgr) |
|
|
|
|
| def list_images(dirs: List[str], limit: int | None = None) -> List[str]: |
| """Deterministically enumerate images across directories (sorted, dedup by basename order).""" |
| exts = ("*.jpg", "*.jpeg", "*.png", "*.JPG", "*.JPEG", "*.PNG") |
| files: List[str] = [] |
| for d in dirs: |
| for ext in exts: |
| files.extend(glob.glob(os.path.join(d, "**", ext), recursive=True)) |
| files = sorted(set(files)) |
| if limit is not None: |
| files = files[:limit] |
| return files |
|
|
|
|
| |
|
|
| def parse_instances( |
| labels: np.ndarray, |
| boxes: np.ndarray, |
| scores: np.ndarray, |
| conf_threshold: float = CONF_THRESHOLD, |
| num_classes: int = 80, |
| ) -> List[dict]: |
| """Mirror parseEdgecrafterSeg: score-filter, validate class id, clamp+order box corners. |
| |
| Returns list of {q, classId, score, box(xyxy in [0,1])}, in raw query order (no NMS). |
| """ |
| labels = np.asarray(labels).reshape(-1) |
| scores = np.asarray(scores).reshape(-1) |
| boxes = np.asarray(boxes).reshape(-1, 4) |
| out = [] |
| for q in range(scores.shape[0]): |
| s = float(scores[q]) |
| if not (s >= conf_threshold): |
| continue |
| cls = int(labels[q]) |
| if cls < 0 or cls >= num_classes: |
| continue |
| x0, y0, x1, y1 = [float(v) for v in boxes[q]] |
| x0, x1 = sorted((min(max(x0, 0.0), 1.0), min(max(x1, 0.0), 1.0))) |
| y0, y1 = sorted((min(max(y0, 0.0), 1.0), min(max(y1, 0.0), 1.0))) |
| out.append({"q": q, "classId": cls, "score": s, "box": (x0, y0, x1, y1)}) |
| return out |
|
|
|
|
| def decode_mask( |
| mask_logits_160: np.ndarray, |
| out_w: int, |
| out_h: int, |
| mask_threshold: float = MASK_THRESHOLD, |
| ) -> np.ndarray: |
| """Mirror decodeEdgeSegmentation: bilinear upscale 160×160 logits to (out_h,out_w), threshold. |
| |
| Returns a boolean mask (out_h, out_w). |
| """ |
| m = np.asarray(mask_logits_160, dtype=np.float32) |
| up = cv2.resize(m, (out_w, out_h), interpolation=cv2.INTER_LINEAR) |
| return up > mask_threshold |
|
|
|
|
| def mask_iou(a: np.ndarray, b: np.ndarray) -> float: |
| """IoU of two boolean masks of identical shape.""" |
| a = a.astype(bool) |
| b = b.astype(bool) |
| inter = np.logical_and(a, b).sum(dtype=np.int64) |
| union = np.logical_or(a, b).sum(dtype=np.int64) |
| if union == 0: |
| return 1.0 |
| return float(inter) / float(union) |
|
|
|
|
| def box_iou(a: Tuple[float, float, float, float], b: Tuple[float, float, float, float]) -> float: |
| ax0, ay0, ax1, ay1 = a |
| bx0, by0, bx1, by1 = b |
| ix0, iy0 = max(ax0, bx0), max(ay0, by0) |
| ix1, iy1 = min(ax1, bx1), min(ay1, by1) |
| iw, ih = max(0.0, ix1 - ix0), max(0.0, iy1 - iy0) |
| inter = iw * ih |
| area_a = max(0.0, ax1 - ax0) * max(0.0, ay1 - ay0) |
| area_b = max(0.0, bx1 - bx0) * max(0.0, by1 - by0) |
| union = area_a + area_b - inter |
| if union <= 0: |
| return 1.0 if inter == 0 else 0.0 |
| return inter / union |
|
|
