wechat_ai/wechat_vision/app.py

from pathlib import Path

import numpy as np
import onnxruntime as ort
from PIL import Image, ImageDraw


MODEL_PATH = Path(__file__).with_name("best.onnx")
TEST_DIR = Path(__file__).with_name("test")
OUTPUT_DIR = Path(__file__).with_name("ouptsw")
IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}
CONFIDENCE_THRESHOLD = 0.1
IOU_THRESHOLD = 0.5


def _dtype(onnx_type):
    if onnx_type == "tensor(float)":
        return np.float32
    if onnx_type == "tensor(float16)":
        return np.float16
    if onnx_type == "tensor(double)":
        return np.float64
    if onnx_type in {"tensor(int64)", "tensor(uint64)"}:
        return np.int64
    if onnx_type in {"tensor(int32)", "tensor(uint32)"}:
        return np.int32
    raise ValueError(f"Unsupported input type: {onnx_type}")


def _shape(input_meta):
    return [dim if isinstance(dim, int) else 1 for dim in input_meta.shape]


def _iou(box, boxes):
    x1 = np.maximum(box[0], boxes[:, 0])
    y1 = np.maximum(box[1], boxes[:, 1])
    x2 = np.minimum(box[2], boxes[:, 2])
    y2 = np.minimum(box[3], boxes[:, 3])

    intersection = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1)
    box_area = np.maximum(0, box[2] - box[0]) * np.maximum(0, box[3] - box[1])
    boxes_area = np.maximum(0, boxes[:, 2] - boxes[:, 0]) * np.maximum(
        0,
        boxes[:, 3] - boxes[:, 1],
    )
    return intersection / np.maximum(box_area + boxes_area - intersection, 1e-6)


def _nms(detections):
    detections = detections[detections[:, 4] >= CONFIDENCE_THRESHOLD]
    if len(detections) == 0:
        return detections

    detections = detections[np.argsort(detections[:, 4])[::-1]]
    kept = []
    while len(detections) > 0:
        best = detections[0]
        kept.append(best)
        if len(detections) == 1:
            break
        detections = detections[1:][_iou(best[:4], detections[1:, :4]) < IOU_THRESHOLD]

    return np.array(kept)


def _preprocess(image, input_shape):
    _, _, height, width = input_shape
    resized = image.resize((width, height))
    array = np.asarray(resized, dtype=np.float32) / 255.0
    return np.transpose(array, (2, 0, 1))[None]


def _scale_detections(detections, original_size, input_shape):
    _, _, input_height, input_width = input_shape
    original_width, original_height = original_size
    scaled = detections.copy()
    scaled[:, [0, 2]] *= original_width / input_width
    scaled[:, [1, 3]] *= original_height / input_height
    scaled[:, [0, 2]] = np.clip(scaled[:, [0, 2]], 0, original_width)
    scaled[:, [1, 3]] = np.clip(scaled[:, [1, 3]], 0, original_height)
    return scaled


def _draw_detections(image, detections):
    draw = ImageDraw.Draw(image)
    for detection in detections:
        x1, y1, x2, y2, score, class_id = detection
        cx = (x1 + x2) / 2
        cy = (y1 + y2) / 2
        label = f"({cx:.1f}, {cy:.1f}) {score:.2f}"

        draw.rectangle((x1, y1, x2, y2), outline="red", width=3)
        draw.ellipse((cx - 5, cy - 5, cx + 5, cy + 5), fill="yellow", outline="red")
        draw.line((cx - 10, cy, cx + 10, cy), fill="red", width=2)
        draw.line((cx, cy - 10, cx, cy + 10), fill="red", width=2)
        draw.text((x1, max(0, y1 - 14)), label, fill="yellow")


def main():
    session = ort.InferenceSession(str(MODEL_PATH))
    input_meta = session.get_inputs()[0]
    input_shape = _shape(input_meta)

    print(f"model: {MODEL_PATH}")
    print(f"providers: {session.get_providers()}")

    print("inputs:")
    feeds = {}
    for input_meta in session.get_inputs():
        print(
            f"- name={input_meta.name}, type={input_meta.type}, "
            f"shape={input_meta.shape}"
        )
        feeds[input_meta.name] = np.zeros(
            _shape(input_meta),
            dtype=_dtype(input_meta.type),
        )

    print("outputs:")
    for output_meta in session.get_outputs():
        print(
            f"- name={output_meta.name}, type={output_meta.type}, "
            f"shape={output_meta.shape}"
        )

    outputs = session.run(None, feeds)
    print("dummy inference: ok")
    for index, output in enumerate(outputs):
        print(f"- output[{index}]: shape={output.shape}, dtype={output.dtype}")

    OUTPUT_DIR.mkdir(exist_ok=True)
    result_lines = ["image,class_id,score,center_x,center_y,x1,y1,x2,y2"]

    for image_path in sorted(TEST_DIR.iterdir()):
        if image_path.suffix.lower() not in IMAGE_SUFFIXES:
            continue

        image = Image.open(image_path).convert("RGB")
        tensor = _preprocess(image, input_shape)
        detections = session.run(None, {input_meta.name: tensor})[0][0]
        detections = _scale_detections(_nms(detections), image.size, input_shape)

        annotated = image.copy()
        _draw_detections(annotated, detections)
        output_path = OUTPUT_DIR / image_path.name
        annotated.save(output_path)

        print(f"{image_path.name}: {len(detections)} target(s) -> {output_path}")
        for detection in detections:
            x1, y1, x2, y2, score, class_id = detection
            cx = (x1 + x2) / 2
            cy = (y1 + y2) / 2
            print(f"  center=({cx:.1f}, {cy:.1f}), score={score:.3f}")
            result_lines.append(
                f"{image_path.name},{int(class_id)},{score:.6f},{cx:.2f},{cy:.2f},"
                f"{x1:.2f},{y1:.2f},{x2:.2f},{y2:.2f}"
            )

    (OUTPUT_DIR / "coordinates.csv").write_text("\n".join(result_lines) + "\n")


if __name__ == "__main__":
    main()