# model.py — DETECT/SEG/POSE 공통 스키마 (Triton Python Backend용) # # [모듈 역할] # - Triton Python backend 진입점 클래스(TritonPythonModel)를 정의한다. # - YOLO TensorRT 엔진(model.engine)을 로드해서 # RAW_IMAGE / ORIGINAL_IMAGE_SHAPE 입력을 받아 추론하고, # SSIMS 파이프라인에서 공통으로 사용하는 JSON 스키마("FINAL_RESULT")로 반환한다. # # [상위/하위 모듈 연계] # - 상위 호출자: # - Triton 서버 → python backend → TritonPythonModel.execute() # - ai_engine.run_appsrc / ai_engine.result_processor 등은 # Triton HTTP 클라이언트(triton_http_client.py)를 통해 # 이 backend의 출력("FINAL_RESULT")을 받아서 후처리한다. # - 출력 JSON 형식: # - "task": "detect" | "pose" | "seg" | "classify" # - "img_size": [H, W] # - "coord_type": "pixel" # - "bboxes": [[x1, y1, x2, y2], ...] # - "relative_bboxes":[[x1/W, y1/H, x2/W, y2/H], ...] # - "labels": [str, ...] # - "scores": [float, ...] # - (옵션) "keypoints": 포즈일 때 [N, K, 3] (x, y, conf) # - (옵션) "masks": 세그멘테이션일 때, 폴리곤 리스트 # # [성능/안전 설계 원칙] # - initialize() / finalize() 에는 무한 루프나 별도 스레드를 두지 않는다. # - execute() 는 "요청이 있을 때만" 호출되며, 함수 내부에 while True 루프를 두지 않는다. # → CPU 폭주 문제는 이 backend 자체보다는, # 상위에서 얼마나 많은 요청을 보내는지(추론 FPS/재시도 루프)와 더 연관이 크다. import json from pathlib import Path import numpy as np import torch import triton_python_backend_utils as pb_utils from ultralytics import YOLO class TritonPythonModel: """ Triton Python backend에서 요구하는 엔트리 클래스. - initialize(self, args): 모델/리소스 로드 (프로세스 시작 시 1회 호출) - execute(self, requests): 추론 처리 (요청이 들어올 때마다 호출) - finalize(self): 정리 작업 (모델 언로드 시 1회 호출) """ def initialize(self, args): """ Triton이 모델 인스턴스를 초기화할 때 호출된다. args 예시: { 'model_repository': '/workspace/ssims.ai/model_repository/yolo11_960_ppe_og_9class', 'model_version': '1', ... } """ # 디바이스 결정 (YOLO 엔진은 내부적으로 TensorRT를 사용하지만, # ultralytics API와의 일관성을 위해 device 문자열 유지) self.device = "cuda" if torch.cuda.is_available() else "cpu" # 모델 repository / version 경로 repo_dir = Path(args["model_repository"]) model_dir = repo_dir / args["model_version"] self.repo_name = repo_dir.name self.model_version = str(args["model_version"]) # --- task 타입 추론 --- # repo 이름에 pose/seg/cls 키워드가 있으면 task_type을 해당 값으로 설정하고, # 아니면 기본값 "detect"로 동작한다. name_lower = self.repo_name.lower() if "pose" in name_lower: self.task_type = "pose" elif "seg" in name_lower or "segment" in name_lower: self.task_type = "seg" elif "cls" in name_lower or "classify" in name_lower: self.task_type = "classify" else: self.task_type = "detect" # --- YOLO TensorRT 엔진 로드 --- # model.engine 파일은 Triton model_repository 안에 version 디렉토리 기준으로 위치한다. engine_path = model_dir / "model.engine" if not engine_path.exists(): raise RuntimeError( f"[INIT][{self.repo_name}] model.engine not found at: {engine_path}" ) # YOLO 엔진 로드 (task_type을 명시해서 헤드를 강제) # ※ 여기에는 while 루프나 warmup 반복을 절대 넣지 않는다. self.model = YOLO(engine_path, task=self.task_type) # --- 클래스 이름 맵 로드 (선택사항) --- # labels.json 이 있으면 {id(str): name(str)} 형식으로 로드해서 int → str 매핑으로 변환. labels_path = model_dir / "labels.json" if labels_path.exists(): try: with open(labels_path, "r", encoding="utf-8") as f: raw_labels = json.load(f) # 키를 int로 강제 변환하여 cls 인덱스에 안전하게 대응 self.class_names = {int(k): v for k, v in raw_labels.items()} except Exception as e: # labels.json 문제로 전체 모델이 죽지 않도록, # 일단 빈 dict 로 두고 로그만 남긴다. print( f"[INIT][{self.repo_name}] failed to load labels.json ({labels_path}): {e}" ) self.class_names = {} else: self.class_names = {} print( f"[INIT] repo='{self.repo_name}' v={self.model_version} " f"task='{self.task_type}' device='{self.device}'" ) # ------------------------------------------------------------------ # 내부 유틸 함수들 # ------------------------------------------------------------------ def _to_numpy(self, x): """ torch.Tensor 또는 numpy-like 객체를 numpy.ndarray로 변환. - keypoints, boxes.xyxy, boxes.conf 등의 값을 일관되게 처리하기 위한 헬퍼. """ if isinstance(x, torch.Tensor): return x.detach().cpu().numpy() if x is None: return None return np.array(x) def _kp_to_x_y_conf(self, kobj, H, W, max_k=17): """ [포즈 전용] keypoints를 [N, K, 3] 형식의 (x, y, conf) 리스트로 변환. - 입력: ultralytics 결과 객체의 keypoints (res.keypoints) - 출력: 파이썬 리스트 (JSON 직렬화 용이) 우선순위: 1) keypoints.data: [N, K, 3] (x, y, conf) 2) keypoints.xy + keypoints.conf 3) keypoints.xyn (정규화 좌표) → 픽셀 좌표로 환산 """ if kobj is None: return [] # 1) data([N,K,3]) try: arr = self._to_numpy(getattr(kobj, "data", None)) if arr is not None and arr.ndim == 3 and arr.shape[2] >= 2: N = arr.shape[0] K = min(max_k, arr.shape[1]) out = [] for i in range(N): one = [] for j in range(K): x = float(arr[i, j, 0]) y = float(arr[i, j, 1]) c = float(arr[i, j, 2]) if arr.shape[2] >= 3 else 1.0 one.append([x, y, c]) out.append(one) return out except Exception: # 포맷이 예상과 다를 경우, 다음 후보로 넘어간다. pass # 2) xy + conf try: xy = self._to_numpy(getattr(kobj, "xy", None)) cf = self._to_numpy(getattr(kobj, "conf", None)) if xy is not None and xy.ndim == 3 and xy.shape[2] == 2: N = xy.shape[0] K = min(max_k, xy.shape[1]) out = [] for i in range(N): one = [] for j in range(K): x = float(xy[i, j, 0]) y = float(xy[i, j, 1]) if cf is not None and cf.ndim == 2: c = float(cf[i, j]) else: c = 1.0 one.append([x, y, c]) out.append(one) return out except Exception: pass # 3) xyn 정규화 → 픽셀 좌표로 환산 (conf는 1.0으로 고정) try: xyn = self._to_numpy(getattr(kobj, "xyn", None)) if xyn is not None and xyn.ndim == 3 and xyn.shape[2] == 2: N = xyn.shape[0] K = min(max_k, xyn.shape[1]) out = [] for i in range(N): one = [] for j in range(K): x = float(xyn[i, j, 0] * W) y = float(xyn[i, j, 1] * H) one.append([x, y, 1.0]) out.append(one) return out except Exception: pass # 어떤 포맷도 맞지 않을 경우 빈 리스트 반환 return [] # ------------------------------------------------------------------ # Triton 필수 메서드: execute # ------------------------------------------------------------------ def execute(self, requests): """ Triton이 실제 추론을 수행할 때 호출되는 메서드. 입력 텐서 이름: - "RAW_IMAGE" : [N, ...] (ultralytics YOLO가 받을 수 있는 이미지 배열/배치) - "ORIGINAL_IMAGE_SHAPE" : [N, 2] = [H, W] 출력 텐서: - "FINAL_RESULT" : dtype=object, shape=[1], JSON 문자열(batch_out)을 담고 있음. """ responses = [] # 이 함수 내부에는 while True 등의 무한 루프를 두지 않는다. # → 요청 수 만큼 한 번씩만 처리하고 반환. for request in requests: try: # ---- 입력 텐서 추출 ---- raw_tensor = pb_utils.get_input_tensor_by_name(request, "RAW_IMAGE") shape_tensor = pb_utils.get_input_tensor_by_name( request, "ORIGINAL_IMAGE_SHAPE" ) if raw_tensor is None or shape_tensor is None: raise ValueError( "Missing required input tensors: RAW_IMAGE and ORIGINAL_IMAGE_SHAPE" ) raw = raw_tensor.as_numpy() shapes = shape_tensor.as_numpy() # [N,2] = [H,W] # ultralytics YOLO는 list[np.ndarray] 또는 list[이미지] 형태를 기대하므로 # numpy 배치를 파이썬 리스트로 변환. imgs = list(raw) # ---- YOLO 추론 ---- # - device는 initialize에서 결정된 self.device 사용 # - verbose=False 로 로그 최소화 results_batch = self.model.predict( imgs, device=self.device, verbose=False ) batch_out = [] # 요청 단위로 결과 파싱 for i, res in enumerate(results_batch): # shapes[i] : [H, W] H = int(shapes[i][0]) W = int(shapes[i][1]) # boxes, labels, scores 파싱 boxes = getattr(res, "boxes", None) has_boxes = boxes is not None and len(boxes) > 0 b_abs, b_rel, labels, scores = [], [], [], [] if has_boxes: # 절대 좌표 [x1, y1, x2, y2] xyxy = self._to_numpy(getattr(boxes, "xyxy", None)) if xyxy is not None: xyxy = xyxy.astype(float) b_abs = xyxy.tolist() # 상대 좌표 [x1/W, y1/H, x2/W, y2/H] scale = np.array([W, H, W, H], dtype=float) b_rel = (xyxy / scale).tolist() # 클래스 ID → 라벨 문자열 매핑 cls_ids = self._to_numpy(getattr(boxes, "cls", None)) if cls_ids is not None: cls_ids = cls_ids.astype(int) labels = [ self.class_names.get(int(c), str(int(c))) for c in cls_ids ] # confidence confs = self._to_numpy(getattr(boxes, "conf", None)) if confs is not None: scores = confs.astype(float).tolist() # 공통 출력 필드 구성 item = { "task": self.task_type, "img_size": [H, W], "coord_type": "pixel", "bboxes": b_abs, "relative_bboxes": b_rel, "labels": labels, "scores": scores, } # ---- pose 결과 추가 ---- if self.task_type == "pose": kpts = getattr(res, "keypoints", None) item["keypoints"] = self._kp_to_x_y_conf( kpts, H, W, max_k=17 ) # ---- seg 결과 추가 ---- if self.task_type == "seg": masks = getattr(res, "masks", None) polys = [] if masks is not None and hasattr(masks, "xy"): try: polys = [poly.tolist() for poly in masks.xy] except Exception: polys = [] item["masks"] = polys batch_out.append(item) # ---- JSON 직렬화 및 Triton 출력 텐서 구성 ---- out_json = json.dumps(batch_out, ensure_ascii=False) out_tensor = pb_utils.Tensor( "FINAL_RESULT", np.array([out_json], dtype=object) ) responses.append( pb_utils.InferenceResponse(output_tensors=[out_tensor]) ) except Exception as e: # 에러가 발생해도 Triton 프로세스가 죽지 않도록, # TritonError 로 감싸서 응답에 포함시킨다. err_msg = f"Model execution failed in repo='{getattr(self, 'repo_name', '?')}': {e}" print(f"[ERROR][TritonPythonModel.execute] {err_msg}") responses.append( pb_utils.InferenceResponse( error=pb_utils.TritonError(err_msg) ) ) return responses def finalize(self): """ Triton이 모델 인스턴스를 언로드할 때 호출된다. 현재는 별도의 정리 작업 없이 로그만 남긴다. """ print(f"[FINALIZE] TritonPythonModel unloaded for repo='{self.repo_name}'")