File size: 2,892 Bytes
abf6b4e | 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 | """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: # 传入的是 0~255 时自动归一化
x = x / 255.0
if x.ndim == 3: # 兼容传入灰度单通道 (128,128,1)
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()
|