predict.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""滤光片外观质检 · 推理脚本（单文件，只依赖 onnxruntime + numpy + pillow）

用法
-----
单张图：
    python predict.py -i /path/to/filter_0001.png

整个目录（递归），输出 xlsx：
    python predict.py -i /path/to/images/ -o result.xlsx

同时输出缺陷类型：
    python predict.py -i /path/to/images/ -o result.xlsx --with-type

自定义判废阈值（默认用标定好的零误报阈值）：
    python predict.py -i imgs/ -o r.xlsx --threshold 0.5

输出 CSV 而不是 xlsx（不需要 openpyxl）：
    python predict.py -i imgs/ -o result.csv
"""
import os, sys, json, argparse, time, csv

try:
    import numpy as np
except ImportError:
    sys.exit('缺少 numpy：pip install numpy')
try:
    import onnxruntime as ort
except ImportError:
    sys.exit('缺少 onnxruntime：pip install onnxruntime')
try:
    from PIL import Image
except ImportError:
    sys.exit('缺少 Pillow：pip install pillow')

if getattr(sys, 'frozen', False):
    HERE = os.path.dirname(sys.executable)          # PyInstaller: exe 所在目录
    BUNDLED = getattr(sys, '_MEIPASS', HERE)        # 打进二进制里的资源
else:
    HERE = BUNDLED = os.path.dirname(os.path.abspath(__file__))
IMG_EXT = {'.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff'}


def find_model(name):
    """先找 exe 旁边（方便换模型），再找打进包里的。"""
    for d in (HERE, BUNDLED):
        p = os.path.join(d, name)
        if os.path.exists(p):
            return p
    return os.path.join(BUNDLED, name)


# ---------------------------------------------------------------- 预处理
def preprocess(path, size):
    """等比缩放到 size 内 + 居中补零 + 逐图灰世界白平衡 + 逐图 z-score。
    必须与训练时逐字节一致，否则结果不可信。"""
    im = Image.open(path).convert('RGB')
    w, h = im.size
    s = size / max(w, h)
    nw, nh = max(1, round(w * s)), max(1, round(h * s))
    im = im.resize((nw, nh), Image.BILINEAR)
    canvas = np.zeros((size, size, 3), np.uint8)
    y0, x0 = (size - nh) // 2, (size - nw) // 2
    canvas[y0:y0 + nh, x0:x0 + nw] = np.asarray(im)
    f = canvas.astype(np.float32)
    m = f.reshape(-1, 3).mean(0)                       # 灰世界
    f = f * (m.mean() / np.maximum(m, 1e-6))
    f = (f - f.mean()) / max(f.std(), 1e-6)            # z-score
    return np.transpose(f, (2, 0, 1))                  # HWC -> CHW


def softmax(x, axis=-1):
    e = np.exp(x - x.max(axis=axis, keepdims=True))
    return e / e.sum(axis=axis, keepdims=True)


# ---------------------------------------------------------------- 模型
class Model:
    def __init__(self, onnx_path):
        meta_p = onnx_path.replace('.onnx', '.json')
        self.meta = json.load(open(meta_p, encoding='utf-8')) if os.path.exists(meta_p) else {}
        self.size = int(self.meta.get('size', 320))
        self.classes = self.meta.get('classes', ['良品', '不良'])
        so = ort.SessionOptions()
        so.intra_op_num_threads = min(8, os.cpu_count() or 4)
        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        prov = ['CUDAExecutionProvider', 'CPUExecutionProvider'] \
            if 'CUDAExecutionProvider' in ort.get_available_providers() else ['CPUExecutionProvider']
        self.sess = ort.InferenceSession(onnx_path, so, providers=prov)
        self.provider = self.sess.get_providers()[0]
        self.iname = self.sess.get_inputs()[0].name

    def run(self, paths, bs=16):
        out = []
        for i in range(0, len(paths), bs):
            chunk = paths[i:i + bs]
            x = np.stack([preprocess(p, self.size) for p in chunk]).astype(np.float32)
            out.append(softmax(self.sess.run(None, {self.iname: x})[0], axis=1))
        return np.concatenate(out) if out else np.zeros((0, len(self.classes)))


def collect(inp):
    if os.path.isfile(inp):
        return [inp]
    fs = []
    for r, _, names in os.walk(inp):
        for n in sorted(names):
            if os.path.splitext(n)[1].lower() in IMG_EXT:
                fs.append(os.path.join(r, n))
    return sorted(fs)


def write_table(rows, header, out):
    ext = os.path.splitext(out)[1].lower()
    if ext in ('.xlsx', '.xls'):
        try:
            from openpyxl import Workbook
        except ImportError:
            out = os.path.splitext(out)[0] + '.csv'
            print('  未装 openpyxl，改写 CSV：' + out)
            ext = '.csv'
        else:
            wb = Workbook(); ws = wb.active; ws.title = '检测结果'
            ws.append(header)
            for r in rows:
                ws.append(r)
            for i, w in enumerate([38, 12, 12, 12, 12, 60], 1):
                if i <= len(header):
                    ws.column_dimensions[chr(64 + i)].width = w
            wb.save(out); return out
    with open(out, 'w', newline='', encoding='utf-8-sig') as f:
        w = csv.writer(f); w.writerow(header); w.writerows(rows)
    return out


def main():
    ap = argparse.ArgumentParser(description='滤光片外观质检推理')
    ap.add_argument('-i', '--input', required=True, help='图片文件或目录')
    ap.add_argument('-o', '--output', default=None, help='输出 xlsx/csv（目录模式下必填）')
    ap.add_argument('-m', '--model', default=None)
    ap.add_argument('--with-type', action='store_true', help='同时用四分类模型给出缺陷类型')
    ap.add_argument('--type-model', default=None)
    ap.add_argument('--threshold', type=float, default=None,
                    help='判废阈值（不良概率）。默认读 filter_binary.json 里标定好的值')
    ap.add_argument('--batch', type=int, default=16)
    a = ap.parse_args()

    if a.model is None: a.model = find_model('filter_binary.onnx')
    if a.type_model is None: a.type_model = find_model('filter_4class.onnx')
    files = collect(a.input)
    if not files:
        sys.exit(f'没找到图片：{a.input}')
    mdl = Model(a.model)
    thr = a.threshold if a.threshold is not None else float(mdl.meta.get('threshold', 0.5))
    print(f'模型 {os.path.basename(a.model)} | {mdl.provider} | 输入 {mdl.size}px | 判废阈值 {thr:.4f}')
    print(f'待检 {len(files)} 张')

    t0 = time.time()
    P = mdl.run(files, a.batch)
    score = P[:, 1:].sum(1)                       # 不良概率
    verdict = np.where(score >= thr, '不良', '良品')

    types = None
    if a.with_type and os.path.exists(a.type_model):
        tm = Model(a.type_model)
        T = tm.run(files, a.batch)
        names = tm.classes
        types = [names[i] for i in T.argmax(1)]
        tconf = T.max(1)
    dt = time.time() - t0
    print(f'完成 {len(files)} 张，用时 {dt:.1f}s（{len(files)/max(dt,1e-9):.1f} 张/秒）')

    if len(files) == 1 and not a.output:
        print(f'\n  文件   : {os.path.basename(files[0])}')
        print(f'  判定   : {verdict[0]}')
        print(f'  不良概率: {score[0]:.4f}')
        if types:
            print(f'  缺陷类型: {types[0]}（置信度 {tconf[0]:.3f}）')
        return

    out = a.output or 'result.xlsx'
    header = ['图片名', '判定', '不良概率']
    rows = []
    for i, p in enumerate(files):
        r = [os.path.basename(p), verdict[i], round(float(score[i]), 4)]
        if types:
            r += [types[i] if verdict[i] == '不良' else '', round(float(tconf[i]), 3) if verdict[i] == '不良' else '']
        r.append(os.path.relpath(p, a.input) if os.path.isdir(a.input) else p)
        rows.append(r)
    if types:
        header += ['缺陷类型', '类型置信度']
    header += ['相对路径']
    out = write_table(rows, header, out)
    n_bad = int((verdict == '不良').sum())
    print(f'\n  良品 {len(files)-n_bad} 张 | 不良 {n_bad} 张（{n_bad/len(files):.1%}）')
    print(f'  结果 -> {out}')


if __name__ == '__main__':
    main()