charlie880624
7716a0060f
Some checks failed
deploy / build-n-publish (push) Has been cancelled
lint / lint (push) Has been cancelled
build / build_cpu (3.7, 1.5.1, torch1.5, 0.6.1) (push) Has been cancelled
build / build_cpu (3.7, 1.6.0, torch1.6, 0.7.0) (push) Has been cancelled
build / build_cpu (3.7, 1.7.0, torch1.7, 0.8.1) (push) Has been cancelled
build / build_cpu (3.7, 1.8.0, torch1.8, 0.9.0) (push) Has been cancelled
build / build_cpu (3.7, 1.9.0, torch1.9, 0.10.0) (push) Has been cancelled
build / build_cuda101 (3.7, 1.5.1+cu101, torch1.5, 0.6.1+cu101) (push) Has been cancelled
build / build_cuda101 (3.7, 1.6.0+cu101, torch1.6, 0.7.0+cu101) (push) Has been cancelled
build / build_cuda101 (3.7, 1.7.0+cu101, torch1.7, 0.8.1+cu101) (push) Has been cancelled
build / build_cuda101 (3.7, 1.8.0+cu101, torch1.8, 0.9.0+cu101) (push) Has been cancelled
build / build_cuda102 (3.6, 1.9.0+cu102, torch1.9, 0.10.0+cu102) (push) Has been cancelled
build / build_cuda102 (3.7, 1.9.0+cu102, torch1.9, 0.10.0+cu102) (push) Has been cancelled
build / build_cuda102 (3.8, 1.9.0+cu102, torch1.9, 0.10.0+cu102) (push) Has been cancelled
build / build_cuda102 (3.9, 1.9.0+cu102, torch1.9, 0.10.0+cu102) (push) Has been cancelled
build / test_windows (windows-2022, cpu, 3.8) (push) Has been cancelled
build / test_windows (windows-2022, cu111, 3.8) (push) Has been cancelled
- Add golf1/2/4/7/8 dataset classes for semantic segmentation - Add kneron-specific configs (meconfig series, kn_stdc1_golf4class) - Organize scripts into tools/check/ and tools/kneron/ - Add kneron_preprocessing module - Update README with quick-start guide - Update .gitignore to exclude data dirs, onnx, nef outputs Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
48 lines
1.5 KiB
Python
48 lines
1.5 KiB
Python
import onnxruntime as ort
|
||
import numpy as np
|
||
from PIL import Image
|
||
import cv2
|
||
|
||
# === 1. 載入 ONNX 模型 ===
|
||
onnx_path = "work_dirs/meconfig8/latest.onnx"
|
||
session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
|
||
|
||
# === 2. 前處理輸入圖像（724x362） ===
|
||
def preprocess(img_path):
|
||
image = Image.open(img_path).convert("RGB")
|
||
image = image.resize((724, 362), Image.BILINEAR)
|
||
img = np.array(image) / 255.0
|
||
img = np.transpose(img, (2, 0, 1)) # HWC → CHW
|
||
img = np.expand_dims(img, 0).astype(np.float32) # (1, 3, 362, 724)
|
||
return img
|
||
|
||
img_path = "test.png"
|
||
input_tensor = preprocess(img_path)
|
||
|
||
# === 3. 執行推論 ===
|
||
input_name = session.get_inputs()[0].name
|
||
output = session.run(None, {input_name: input_tensor}) # list of np.array
|
||
|
||
# === 4. 後處理 + 預測 Mask ===
|
||
output_tensor = output[0][0] # shape: (num_classes, H, W)
|
||
pred_mask = np.argmax(output_tensor, axis=0).astype(np.uint8) # (H, W)
|
||
|
||
# === 5. 可視化結果 ===
|
||
colors = [
|
||
[128, 0, 0], # 0: bunker
|
||
[0, 0, 128], # 1: car
|
||
[0, 128, 0], # 2: grass
|
||
[0, 255, 0], # 3: greenery
|
||
[255, 0, 0], # 4: person
|
||
[255, 165, 0], # 5: road
|
||
[0, 255, 255], # 6: tree
|
||
]
|
||
|
||
color_mask = np.zeros((pred_mask.shape[0], pred_mask.shape[1], 3), dtype=np.uint8)
|
||
for cls_id, color in enumerate(colors):
|
||
color_mask[pred_mask == cls_id] = color
|
||
|
||
# 儲存可視化圖片
|
||
cv2.imwrite("onnx_pred_mask.png", color_mask)
|
||
print("✅ 預測結果已儲存為：onnx_pred_mask.png")
|