File size: 3,927 Bytes
c76173d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | """
model_runner.py
ทำงานใน child process (ProcessPoolExecutor)
- init_worker() → โหลด ONNX session 1 ครั้งต่อ process
- predict_emotion() → รับ image bytes, คืน dict ผลลัพธ์
แยกไฟล์นี้ออกมาเพื่อหลีกเลี่ยง pickling ปัญหาของ ONNX session
"""
import io
import time
from typing import Dict, Any
import numpy as np
import onnxruntime as ort
from PIL import Image
# ─── Global state ต่อ process ─────────────────────────────
_session: ort.InferenceSession = None
_input_name: str = None
_id2label: Dict[int, str] = None
# Label map ของ dima806/facial_emotions_image_detection
EMOTION_LABELS = {
0: "angry",
1: "disgust",
2: "fear",
3: "happy",
4: "neutral",
5: "sad",
6: "surprise",
}
# Preprocessing constants (ViT-style, ตรงกับ Hugging Face preprocessor_config.json)
IMAGE_SIZE = 224
MEAN = np.array([0.5, 0.5, 0.5], dtype=np.float32)
STD = np.array([0.5, 0.5, 0.5], dtype=np.float32)
def init_worker(model_path: str) -> None:
"""
เรียกครั้งเดียวตอน process pool สร้าง worker
โหลด ONNX session เข้า global variable ของ process นั้น
"""
global _session, _input_name, _id2label
sess_options = ort.SessionOptions()
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
# inter/intra thread ต่อ worker — ป้องกัน over-subscription
sess_options.inter_op_num_threads = 1
sess_options.intra_op_num_threads = 2
_session = ort.InferenceSession(
model_path,
sess_options=sess_options,
providers=["CPUExecutionProvider"],
)
_input_name = _session.get_inputs()[0].name
_id2label = EMOTION_LABELS
def preprocess(image_bytes: bytes) -> np.ndarray:
"""
แปลง raw image bytes → numpy tensor [1, 3, 224, 224] float32
ใช้ manual preprocessing แทน AutoFeatureExtractor เพื่อลด overhead
"""
img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
img = img.resize((IMAGE_SIZE, IMAGE_SIZE), Image.BICUBIC)
arr = np.array(img, dtype=np.float32) / 255.0 # [0,1]
arr = (arr - MEAN) / STD # normalize
arr = arr.transpose(2, 0, 1) # HWC → CHW
arr = np.expand_dims(arr, axis=0) # [1, 3, H, W]
return arr.astype(np.float32)
def softmax(x: np.ndarray) -> np.ndarray:
e = np.exp(x - np.max(x))
return e / e.sum()
def predict_emotion(image_bytes: bytes) -> Dict[str, Any]:
"""
ฟังก์ชันหลักที่ถูกเรียกจาก main process ผ่าน executor
รับ image bytes, คืน dict ผลลัพธ์
"""
if _session is None:
raise RuntimeError("Model session not initialized. init_worker() was not called.")
# Preprocess
input_tensor = preprocess(image_bytes)
# Inference
t0 = time.perf_counter()
outputs = _session.run(None, {_input_name: input_tensor})
t1 = time.perf_counter()
inference_ms = (t1 - t0) * 1000
# Postprocess
logits = outputs[0][0] # shape: (num_classes,)
probs = softmax(logits)
predicted_idx = int(np.argmax(probs))
predicted_label = _id2label[predicted_idx]
confidence = float(probs[predicted_idx])
all_scores = {
_id2label[i]: round(float(probs[i]), 4)
for i in range(len(probs))
}
return {
"predicted_emotion": predicted_label,
"confidence": round(confidence, 4),
"all_scores": all_scores,
"inference_time_ms": round(inference_ms, 2),
}
|