| """Easy4K SFDP · Siamese NET AI 推理示例。 |
| |
| 用法： |
| python example_infer.py [模型文件] [可选：视频文件] |
| |
| 不传视频时用合成数据自检；传视频时对首两帧取同位置块做判重演示。 |
| """ |
| import sys |
|
|
| import numpy as np |
| import onnxruntime as ort |
|
|
| SIZE = 128 |
| THRESHOLD = 0.5 |
|
|
|
|
| def make_session(model_path="model.onnx"): |
| """DirectML 可用就用 GPU，否则自动回退 CPU。""" |
| avail = ort.get_available_providers() |
| providers = [p for p in ("DmlExecutionProvider", "CPUExecutionProvider") if p in avail] |
| sess = ort.InferenceSession(model_path, providers=providers) |
| print("providers:", sess.get_providers()) |
| return sess |
|
|
|
|
| def prep(patch): |
| """灰度块 -> (1, 1, 128, 128) float32，值域 [0, 1]。""" |
| x = np.asarray(patch, dtype=np.float32) |
| if x.max() > 1.5: |
| x = x / 255.0 |
| if x.ndim == 3: |
| x = x[..., 0] |
| return x.reshape(1, 1, SIZE, SIZE) |
|
|
|
|
| def score(sess, gray_a, gray_b): |
| a, b = prep(gray_a), prep(gray_b) |
| return float(sess.run(["score"], {"frame_a": a, "frame_b": b})[0][0, 0]) |
|
|
|
|
| def score_batch(sess, pairs_a, pairs_b): |
| """批量推理：pairs 形状 (N, 128, 128)，返回 (N,) 分数。""" |
| a = np.stack([prep(p).reshape(SIZE, SIZE) for p in pairs_a], axis=0).astype(np.float32) |
| b = np.stack([prep(p).reshape(SIZE, SIZE) for p in pairs_b], axis=0).astype(np.float32) |
| a = a[:, None, :, :] |
| b = b[:, None, :, :] |
| return sess.run(["score"], {"frame_a": a, "frame_b": b})[0].ravel() |
|
|
|
|
| def crop(gray, x, y): |
| p = gray[y:y + SIZE, x:x + SIZE] |
| return p if p.shape[:2] == (SIZE, SIZE) else None |
|
|
|
|
| def main(): |
| model = sys.argv[1] if len(sys.argv) > 1 else "model.onnx" |
| video = sys.argv[2] if len(sys.argv) > 2 else None |
| sess = make_session(model) |
|
|
| if not video: |
| base = np.random.randint(0, 255, (SIZE, SIZE), np.uint8) |
| print("same ->", round(score(sess, base, base), 4)) |
| print("noise ->", round(score(sess, base, np.random.randint(0, 255, (SIZE, SIZE), np.uint8)), 4)) |
| pairs = [base] * 4 |
| print("batch ->", np.round(score_batch(sess, pairs, pairs), 4)) |
| return |
|
|
| import cv2 |
|
|
| cap = cv2.VideoCapture(video) |
| ok1, f1 = cap.read() |
| ok2, f2 = cap.read() |
| cap.release() |
| if not (ok1 and ok2): |
| raise SystemExit("视频读取失败") |
|
|
| g1 = cv2.cvtColor(f1, cv2.COLOR_BGR2GRAY) |
| g2 = cv2.cvtColor(f2, cv2.COLOR_BGR2GRAY) |
|
|
| h, w = g1.shape |
| a, b = crop(g1, (w - SIZE) // 2, (h - SIZE) // 2), crop(g2, (w - SIZE) // 2, (h - SIZE) // 2) |
| if a is None: |
| raise SystemExit("视频分辨率小于 128x128") |
|
|
| s = score(sess, a, b) |
| print("score =", round(s, 4), "| duplicate" if s >= THRESHOLD else "| different") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
