297 lines
14 KiB
Python
297 lines
14 KiB
Python
import kp
|
||
import time
|
||
import numpy as np
|
||
from typing import List, Dict, Any, Callable, Optional
|
||
import queue
|
||
import threading
|
||
import multiprocessing
|
||
import cv2
|
||
import os
|
||
|
||
# 定義一個 Dongle 的設定結構
|
||
class DongleConfig:
|
||
def __init__(self, port_id: list, scpu_fw_path: str, ncpu_fw_path: str, model_path: str, device_type: str = "KL520"):
|
||
self.port_id = port_id
|
||
self.scpu_fw_path = scpu_fw_path
|
||
self.ncpu_fw_path = ncpu_fw_path
|
||
self.model_path = model_path
|
||
self.device_type = device_type
|
||
|
||
# 定義一個 Pipeline 的層級結構
|
||
class PipelineLayer:
|
||
def __init__(self, name: str, dongle_config: DongleConfig, preprocess_func: Optional[Callable] = None, postprocess_func: Optional[Callable] = None):
|
||
self.name = name
|
||
self.dongle_config = dongle_config
|
||
self.preprocess_func = preprocess_func
|
||
self.postprocess_func = postprocess_func
|
||
|
||
class KneronPipeline:
|
||
def __init__(self, pipeline_layers: List[PipelineLayer]):
|
||
if not pipeline_layers:
|
||
raise ValueError("Pipeline must have at least one layer.")
|
||
self.pipeline_layers = pipeline_layers
|
||
self._dongles: Dict[str, Any] = {} # 儲存 kp.core.DeviceGroup 實例
|
||
self._model_descriptors: Dict[str, Any] = {} # 儲存模型描述符
|
||
self._layer_connections: List[tuple] = [] # 儲存層之間的連接關係
|
||
self._initialized = False
|
||
self._lock = threading.Lock() # 用於初始化保護
|
||
|
||
def add_layer_connection(self, from_layer_name: str, to_layer_name: str):
|
||
"""
|
||
定義不同層之間的資料流向。
|
||
例如: pipeline.add_layer_connection("layer1", "layer2")
|
||
表示 layer1 的輸出作為 layer2 的輸入。
|
||
更複雜的連接方式可能需要更詳細的定義，例如指定輸出節點到輸入節點的對應。
|
||
"""
|
||
from_layer = next((layer for layer in self.pipeline_layers if layer.name == from_layer_name), None)
|
||
to_layer = next((layer for layer in self.pipeline_layers if layer.name == to_layer_name), None)
|
||
if not from_layer or not to_layer:
|
||
raise ValueError(f"Invalid layer names: {from_layer_name} or {to_layer_name} not found.")
|
||
self._layer_connections.append((from_layer_name, to_layer_name))
|
||
|
||
def initialize(self):
|
||
"""
|
||
初始化所有 dongles, 載入韌體和模型。
|
||
"""
|
||
with self._lock:
|
||
if self._initialized:
|
||
print("Pipeline already initialized.")
|
||
return
|
||
|
||
print("[初始化 Pipeline...]")
|
||
for layer in self.pipeline_layers:
|
||
config = layer.dongle_config
|
||
print(f"[連接設備] Layer: {layer.name}, Port: {config.port_id}")
|
||
try:
|
||
# 使用單獨的 DeviceGroup 來管理每個 dongle
|
||
device_group = kp.core.connect_devices(usb_port_ids=config.port_id)
|
||
self._dongles[layer.name] = device_group
|
||
print(f" - {layer.name}: 連接成功")
|
||
|
||
print(f"[設置超時] Layer: {layer.name}")
|
||
kp.core.set_timeout(device_group=device_group, milliseconds=5000)
|
||
print(f" - {layer.name}: 超時設置成功")
|
||
|
||
print(f"[上傳韌體] Layer: {layer.name}")
|
||
kp.core.load_firmware_from_file(device_group=device_group,
|
||
scpu_fw_path=config.scpu_fw_path,
|
||
ncpu_fw_path=config.ncpu_fw_path)
|
||
print(f" - {layer.name}: 韌體上傳成功")
|
||
|
||
print(f"[上傳模型] Layer: {layer.name}")
|
||
model_descriptor = kp.core.load_model_from_file(device_group=device_group,
|
||
file_path=config.model_path)
|
||
self._model_descriptors[layer.name] = model_descriptor
|
||
print(f" - {layer.name}: 模型上傳成功")
|
||
|
||
except Exception as e:
|
||
print(f"錯誤: 初始化 Layer {layer.name} 失敗: {str(e)}")
|
||
# 清理已連接的設備
|
||
self.release()
|
||
raise e
|
||
|
||
self._initialized = True
|
||
print("[Pipeline 初始化完成]")
|
||
|
||
def run(self, input_data: Any) -> Dict[str, Any]:
|
||
"""
|
||
執行整個 pipeline。
|
||
這部分需要處理平行和串行的執行邏輯。
|
||
輸入可以是原始數據 (例如圖片路徑)，第一個 layer 的 preprocess 會處理它。
|
||
"""
|
||
if not self._initialized:
|
||
raise RuntimeError("Pipeline not initialized. Call .initialize() first.")
|
||
|
||
# 這裡需要實現平行和多層邏輯。
|
||
# 一種方式是使用 ThreadPoolExecutor 或 ProcessPoolExecutor。
|
||
# 另一種是手動管理 Thread/Process。
|
||
# 考慮到 dongle 通訊的 I/O 綁定特性，Thread 可能更適合平行處理。
|
||
# 但如果 preprocess/postprocess 是 CPU 綁定，則 multiprocessing 更優。
|
||
# 我們先假設 dongle 通訊是主要瓶頸，使用 threading。
|
||
# 如果 preprocess/postprocess 也是瓶頸，可以考慮在 pipeline 內部針對這些步驟使用 Process。
|
||
|
||
results: Dict[str, Any] = {}
|
||
# 這個範例只處理簡單的順序 pipeline，平行和複雜串接需要更多邏輯
|
||
# TODO: 實現平行和複雜串接邏輯
|
||
|
||
current_input = input_data
|
||
for i, layer in enumerate(self.pipeline_layers):
|
||
print(f"[執行 Layer] {layer.name}")
|
||
dongle = self._dongles[layer.name]
|
||
model_descriptor = self._model_descriptors[layer.name]
|
||
|
||
# 預處理
|
||
processed_input = current_input
|
||
if layer.preprocess_func:
|
||
print(f" - 執行 {layer.name} 的預處理")
|
||
processed_input = layer.preprocess_func(current_input)
|
||
|
||
# 推論
|
||
print(f" - 執行 {layer.name} 的推論")
|
||
try:
|
||
# 假設 processed_input 是 kp.GenericInputNodeImage 列表或可轉換為它
|
||
# 這裡需要根據實際的 preprocess 輸出和模型輸入來調整
|
||
if isinstance(processed_input, list) and all(isinstance(item, kp.GenericInputNodeImage) for item in processed_input):
|
||
inference_input_descriptor = kp.GenericImageInferenceDescriptor(
|
||
model_id=model_descriptor.models[0].id,
|
||
inference_number=0,
|
||
input_node_image_list=processed_input
|
||
)
|
||
elif isinstance(processed_input, np.ndarray):
|
||
# 假設 preprocess 輸出了 numpy array, 需要轉換為 GenericInputNodeImage
|
||
# 這需要更詳細的 info, 例如圖像格式, resize, padding, normalize
|
||
# 這裡先給一個簡易範例，假設是BGR565, 128x128
|
||
inference_input_descriptor = kp.GenericImageInferenceDescriptor(
|
||
model_id=model_descriptor.models[0].id,
|
||
inference_number=0,
|
||
input_node_image_list=[
|
||
kp.GenericInputNodeImage(
|
||
image=processed_input,
|
||
image_format=kp.ImageFormat.KP_IMAGE_FORMAT_RGB565, # 這裡需要根據你的 preprocess 輸出調整
|
||
resize_mode=kp.ResizeMode.KP_RESIZE_ENABLE,
|
||
padding_mode=kp.PaddingMode.KP_PADDING_CORNER,
|
||
normalize_mode=kp.NormalizeMode.KP_NORMALIZE_KNERON
|
||
)
|
||
]
|
||
)
|
||
else:
|
||
raise TypeError(f"Unsupported processed input type for layer {layer.name}: {type(processed_input)}")
|
||
|
||
|
||
kp.inference.generic_image_inference_send(device_group=dongle,
|
||
generic_inference_input_descriptor=inference_input_descriptor)
|
||
generic_raw_result = kp.inference.generic_image_inference_receive(device_group=dongle)
|
||
|
||
# 處理原始結果
|
||
inf_node_output_list = []
|
||
for node_idx in range(generic_raw_result.header.num_output_node):
|
||
# 這裡假設輸出是 float 類型，需要根據你的模型輸出類型調整
|
||
inference_float_node_output = kp.inference.generic_inference_retrieve_float_node(
|
||
node_idx=node_idx,
|
||
generic_raw_result=generic_raw_result,
|
||
channels_ordering=kp.ChannelOrdering.KP_CHANNEL_ORDERING_CHW # 需要根據模型輸出調整
|
||
)
|
||
inf_node_output_list.append(inference_float_node_output.ndarray.copy())
|
||
|
||
raw_output = inf_node_output_list # 可以是 list of numpy arrays
|
||
|
||
except Exception as e:
|
||
print(f"錯誤: Layer {layer.name} 推論失敗: {str(e)}")
|
||
raise e
|
||
|
||
# 後處理
|
||
final_output = raw_output
|
||
if layer.postprocess_func:
|
||
print(f" - 執行 {layer.name} 的後處理")
|
||
final_output = layer.postprocess_func(raw_output)
|
||
|
||
results[layer.name] = final_output
|
||
|
||
# 設定下一個 layer 的輸入 (簡易串接，更複雜需要 _layer_connections 邏輯)
|
||
current_input = final_output
|
||
|
||
return results
|
||
|
||
def release(self):
|
||
"""
|
||
釋放所有 dongles 連接。
|
||
"""
|
||
with self._lock:
|
||
if not self._initialized:
|
||
print("Pipeline not initialized.")
|
||
return
|
||
|
||
print("[釋放 Pipeline...]")
|
||
for layer_name, dongle in self._dongles.items():
|
||
try:
|
||
kp.core.disconnect_devices(device_group=dongle)
|
||
print(f" - {layer_name}: 已斷開連接")
|
||
except Exception as e:
|
||
print(f"錯誤: 斷開 Layer {layer_name} 連接失敗: {str(e)}")
|
||
self._dongles = {}
|
||
self._model_descriptors = {}
|
||
self._initialized = False
|
||
print("[Pipeline 釋放完成]")
|
||
|
||
def __enter__(self):
|
||
self.initialize()
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||
self.release()
|
||
|
||
# 範例使用
|
||
if __name__ == '__main__':
|
||
# 定義你的 preprocess 和 postprocess 函數
|
||
def my_preprocess(image_path: str):
|
||
# 參照你提供的 2_2nef_test.py
|
||
img = cv2.imread(image_path)
|
||
if img is None:
|
||
raise Exception(f"無法讀取圖片: {image_path}")
|
||
img_resized = cv2.resize(img, (128, 128))
|
||
img_bgr565 = cv2.cvtColor(img_resized, cv2.COLOR_BGR2BGR565)
|
||
# 返回 numpy array，KneronPipeline.run 中會轉換為 GenericInputNodeImage
|
||
return img_bgr565
|
||
|
||
def my_postprocess(raw_output: List[np.ndarray]):
|
||
# 參照你提供的 2_2nef_test.py
|
||
probability = raw_output[0].flatten()[0] # 假設是單一輸出節點，取第一個值
|
||
result = "Fire" if probability > 0.5 else "No Fire"
|
||
return {"result": result, "confidence": probability}
|
||
|
||
def another_preprocess(data: Any):
|
||
# 另一個 layer 的預處理
|
||
print("執行第二層的預處理...")
|
||
return data # 這裡只是範例，實際需要根據前一层的輸出和當前層模型輸入來處理
|
||
|
||
def another_postprocess(raw_output: List[np.ndarray]):
|
||
# 另一個 layer 的後處理
|
||
print("執行第二層的後處理...")
|
||
# 假設這層輸出是另一個分類結果
|
||
class_id = np.argmax(raw_output[0].flatten())
|
||
return {"class_id": class_id}
|
||
|
||
|
||
# 定義 Dongle 配置
|
||
dongle_config1 = DongleConfig(port_id=0, scpu_fw_path='fw_scpu.bin', ncpu_fw_path='fw_ncpu.bin', model_path='models_520.nef')
|
||
# 如果有另一個 dongle 和模型
|
||
dongle_config2 = DongleConfig(port_id=1, scpu_fw_path='fw_scpu.bin', ncpu_fw_path='fw_ncpu.bin', model_path='another_model.nef')
|
||
|
||
|
||
# 定義 Pipeline 層
|
||
# 單層 pipeline (平行處理多個輸入可以使用這個 structure, 但 run 方法需要修改)
|
||
# layers_single = [
|
||
# PipelineLayer(name="detector_dongle_0", dongle_config=dongle_config1, preprocess_func=my_preprocess, postprocess_func=my_postprocess),
|
||
# # 如果想平行處理，可以在這裡加更多使用不同 dongle 的 layer，但 run 方法需要平行化
|
||
# # PipelineLayer(name="detector_dongle_1", dongle_config=dongle_config2, preprocess_func=my_preprocess, postprocess_func=my_postprocess),
|
||
# ]
|
||
|
||
# 多層 pipeline (串接不同 dongles)
|
||
layers_multi = [
|
||
PipelineLayer(name="detector_layer", dongle_config=dongle_config1, preprocess_func=my_preprocess, postprocess_func=my_postprocess),
|
||
PipelineLayer(name="classifier_layer", dongle_config=dongle_config2, preprocess_func=another_preprocess, postprocess_func=another_postprocess),
|
||
]
|
||
|
||
|
||
# 建立 Pipeline 實例
|
||
# pipeline = KneronPipeline(pipeline_layers=layers_single) # 單層範例
|
||
pipeline = KneronPipeline(pipeline_layers=layers_multi) # 多層範例
|
||
|
||
# 定義層之間的連接 (僅多層時需要，目前 run 方法只支持簡單順序串接)
|
||
# pipeline.add_layer_connection("detector_layer", "classifier_layer")
|
||
|
||
# 使用 with 語句確保釋放資源
|
||
try:
|
||
with pipeline:
|
||
# 執行推論
|
||
image_path = r'C:\Users\USER\Desktop\Yu-An\Firedetection\test_images\fire4.jpeg'
|
||
results = pipeline.run(input_data=image_path)
|
||
|
||
print("\nPipeline 執行結果:")
|
||
for layer_name, output in results.items():
|
||
print(f" Layer '{layer_name}' 輸出: {output}")
|
||
|
||
# 如果是平行處理，可以在這裡輸入多個 image paths，然後在 run 方法裡分派給不同的 dongle
|
||
|
||
except Exception as e:
|
||
print(f"Pipeline 執行過程中發生錯誤: {str(e)}") |