#!/usr/bin/env python3
"""
2025年07月09日 = Test completed
inference_full_onnx.py

Supports inference for a single image or batch from a folder, using an exported ONNX model (yolov9_full.onnx)
that includes complete post-processing (decode + NMS).
Automatically detects whether to use GPU (CUDAExecutionProvider).
Falls back to CPUExecutionProvider if CUDA is unavailable.
When displaying, the window will be resized to 1/2 of the original image dimensions.
"""

import cv2
import numpy as np
import onnxruntime as ort
from PIL import Image
from pathlib import Path
from datetime import datetime

---------- Configuration ----------

#IMAGE_PATH = "D:/Work/Study/Yolo/Image_Source/ColorSticker/images/val/CamerA_C_20250616_155413.jpg"
IMAGE_PATH = "D:/Work/Study/Yolo/Image_Source/ColorSticker/images/val"
ONNX_PATH = "yolov9_full.onnx"
IMG_SIZE = 640
CLASS_LIST = ["Sticker"] # Single class label

---------- Utility Functions ----------

def letterbox(pil_img, size=IMG_SIZE, color=(114,114,114)):
"""
Resize the image with aspect ratio preserved and pad to square (same as training).
Returns:
- padded PIL.Image (size×ばつsize)
- r float: scaling ratio
- pad_w int : width padding (left-right)
- pad_h int : height padding (top-bottom)
"""
w0, h0 = pil_img.size
r = min(size / w0, size / h0)
nw, nh = int(w0 * r), int(h0 * r)
resized = pil_img.resize((nw, nh), Image.LANCZOS)
pad_w, pad_h = (size - nw) // 2, (size - nh) // 2
canvas = Image.new("RGB", (size, size), color)
canvas.paste(resized, (pad_w, pad_h))
return canvas, r, pad_w, pad_h

def get_image_list(path_str):
"""
If path_str is a file, return [path_str];
If it's a folder, return a list of all common image file paths.
"""
p = Path(path_str)
if p.is_file():
return [p]
elif p.is_dir():
exts = {".jpg", ".jpeg", ".png", ".bmp"}
return sorted([x for x in p.iterdir() if x.suffix.lower() in exts])
else:
raise FileNotFoundError(f"Path not found: {path_str}")

---------- Main Flow ----------

def main():
# 1) Prepare image list
img_paths = get_image_list(IMAGE_PATH)

# 2) Create ONNXRuntime session; check for CUDAExecutionProvider
providers = ["CUDAExecutionProvider"] if "CUDAExecutionProvider" in ort.get_available_providers() \
 else ["CPUExecutionProvider"]
print(f"[{datetime.now():%H:%M:%S}] Using ONNXRuntime Provider: {providers[0]}")
sess = ort.InferenceSession(ONNX_PATH, providers=providers)
inp_name = sess.get_inputs()[0].name
for img_path in img_paths:
 # 3) Load and preprocess image
 pil = Image.open(img_path).convert("RGB")
 img_pad, r, pad_w, pad_h = letterbox(pil, IMG_SIZE)
 inp = np.array(img_pad, dtype=np.float32) / 255.0
 inp = inp.transpose(2, 0, 1)[None, ...] # shape = [1,3,IMG_SIZE,IMG_SIZE]
 # 4) ONNX inference
 dets = sess.run(None, {inp_name: inp})[0] # [num_detections, 6]
 # 5) Map results back to original image and draw boxes
 img_bgr = cv2.cvtColor(np.array(pil), cv2.COLOR_RGB2BGR)
 for cls, x1, y1, x2, y2, conf in dets:
 xx1 = (x1 - pad_w) / r
 yy1 = (y1 - pad_h) / r
 xx2 = (x2 - pad_w) / r
 yy2 = (y2 - pad_h) / r
 W0, H0 = pil.size
 xx1 = max(min(xx1, W0 - 1), 0)
 yy1 = max(min(yy1, H0 - 1), 0)
 xx2 = max(min(xx2, W0 - 1), 0)
 yy2 = max(min(yy2, H0 - 1), 0)
 xi1, yi1, xi2, yi2 = map(int, (xx1, yy1, xx2, yy2))
 cv2.rectangle(img_bgr, (xi1, yi1), (xi2, yi2), (0, 255, 0), 2)
 label = f"{CLASS_LIST[int(cls)]} {int(conf*100)}%"
 cv2.putText(
 img_bgr, label, (xi1, yi1 - 6),
 cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1
 )
 # 6) Save results
 out_dir = Path("outputs")
 out_dir.mkdir(exist_ok=True)
 out_path = out_dir / f"onnx_full_{img_path.stem}.jpg"
 cv2.imwrite(str(out_path), img_bgr)
 print(f"[{datetime.now():%H:%M:%S}] Saved {out_path}")
 # 7) Display at half size
 h, w = img_bgr.shape[:2]
 disp = cv2.resize(img_bgr, (w // 2, h // 2), interpolation=cv2.INTER_LINEAR)
 cv2.imshow(f"ONNX Inference: {img_path.name}", disp)
 cv2.waitKey(0)
 cv2.destroyAllWindows()

if name == "main":
main()