File size: 4,993 Bytes
9496f98 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | #!/usr/bin/env python3
"""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 # raw logit cut (logit > 0 ⇔ sigmoid > 0.5)
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 # broadcast over channels
# HWC -> CHW -> NCHW
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
# ------------------------------------------------------------------ output post-processing ----
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): # NaN-safe
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 # both empty -> identical
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
|