forked from dsutanto/nuclio-handler
336 lines
11 KiB
Python
336 lines
11 KiB
Python
"""Nuclio handler for Ultralytics YOLO-seg models exported to OpenVINO (CVAT mask output)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import io
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from openvino import Core
|
|
from PIL import Image
|
|
|
|
MODELS_DIR = Path(os.environ.get("MODELS_DIR", "/models"))
|
|
CLASSES_PATH = Path(os.environ.get("CLASSES_PATH", "/opt/nuclio/classes.json"))
|
|
DEFAULT_IMGSZ = int(os.environ.get("MODEL_IMGSZ", "320"))
|
|
DEFAULT_CONF_THRESHOLD = float(os.environ.get("CONF_THRESHOLD", "0.25"))
|
|
DEFAULT_MASK_THRESHOLD = float(os.environ.get("MASK_THRESHOLD", "0.0"))
|
|
BOX_CHANNELS = 6
|
|
|
|
|
|
def load_classes(path: Path) -> dict[int, str]:
|
|
with path.open(encoding="utf-8") as handle:
|
|
payload = json.load(handle)
|
|
|
|
if isinstance(payload, list):
|
|
return {int(item["id"]): str(item["name"]) for item in payload}
|
|
|
|
class_map = payload.get("class", payload)
|
|
return {int(class_id): str(name) for class_id, name in class_map.items()}
|
|
|
|
|
|
def find_model_xml(models_dir: Path) -> Path:
|
|
candidates = [
|
|
models_dir / "best.xml",
|
|
models_dir / "best_openvino_model" / "best.xml",
|
|
*sorted(models_dir.glob("*.xml")),
|
|
*sorted(models_dir.glob("**/*.xml")),
|
|
]
|
|
for candidate in candidates:
|
|
if candidate.is_file():
|
|
return candidate
|
|
raise FileNotFoundError(
|
|
f"No OpenVINO XML model found under {models_dir}. "
|
|
"Mount the exported openvino_model directory at /models."
|
|
)
|
|
|
|
|
|
def letterbox(
|
|
image: np.ndarray,
|
|
new_shape: tuple[int, int] = (320, 320),
|
|
color: tuple[int, int, int] = (114, 114, 114),
|
|
) -> tuple[np.ndarray, float, tuple[float, float]]:
|
|
height, width = image.shape[:2]
|
|
target_height, target_width = new_shape
|
|
scale = min(target_height / height, target_width / width)
|
|
new_unpad_width = int(round(width * scale))
|
|
new_unpad_height = int(round(height * scale))
|
|
resized = cv2.resize(image, (new_unpad_width, new_unpad_height), interpolation=cv2.INTER_LINEAR)
|
|
|
|
pad_width = target_width - new_unpad_width
|
|
pad_height = target_height - new_unpad_height
|
|
pad_left = pad_width / 2
|
|
pad_top = pad_height / 2
|
|
padded = cv2.copyMakeBorder(
|
|
resized,
|
|
int(round(pad_top - 0.1)),
|
|
int(round(pad_height - pad_top)),
|
|
int(round(pad_left - 0.1)),
|
|
int(round(pad_width - pad_left)),
|
|
cv2.BORDER_CONSTANT,
|
|
value=color,
|
|
)
|
|
return padded, scale, (pad_left, pad_top)
|
|
|
|
|
|
def preprocess_image(image_bgr: np.ndarray, imgsz: int) -> tuple[np.ndarray, float, tuple[float, float]]:
|
|
letterboxed, scale, pad = letterbox(image_bgr, new_shape=(imgsz, imgsz))
|
|
rgb = letterboxed[:, :, ::-1].transpose(2, 0, 1)
|
|
tensor = np.expand_dims(rgb, axis=0).astype(np.float32) / 255.0
|
|
return tensor, scale, pad
|
|
|
|
|
|
def scale_boxes(
|
|
boxes: np.ndarray,
|
|
image_shape: tuple[int, int],
|
|
scale: float,
|
|
pad: tuple[float, float],
|
|
) -> np.ndarray:
|
|
boxes = boxes.copy()
|
|
pad_x, pad_y = pad
|
|
boxes[:, [0, 2]] -= pad_x
|
|
boxes[:, [1, 3]] -= pad_y
|
|
boxes[:, :4] /= scale
|
|
boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, image_shape[1])
|
|
boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, image_shape[0])
|
|
return boxes
|
|
|
|
|
|
def to_cvat_mask(box: list[int], mask: np.ndarray) -> list[int]:
|
|
xtl, ytl, xbr, ybr = box
|
|
flattened = mask[ytl : ybr + 1, xtl : xbr + 1].astype(np.uint8).ravel().tolist()
|
|
flattened.extend([xtl, ytl, xbr, ybr])
|
|
return flattened
|
|
|
|
|
|
def split_seg_outputs(raw_outputs: dict[Any, np.ndarray]) -> tuple[np.ndarray, np.ndarray]:
|
|
detections = None
|
|
proto = None
|
|
|
|
for value in raw_outputs.values():
|
|
array = np.array(value)
|
|
if array.ndim == 3 and array.shape[-1] >= BOX_CHANNELS:
|
|
detections = array
|
|
elif array.ndim == 4:
|
|
proto = array
|
|
|
|
if detections is None or proto is None:
|
|
shapes = {str(key): np.array(value).shape for key, value in raw_outputs.items()}
|
|
raise ValueError(
|
|
"Expected YOLO-seg outputs: detections (1, N, 6+nm) and proto (1, nm, H, W). "
|
|
f"Got shapes: {shapes}. Re-export the segmentation model to OpenVINO."
|
|
)
|
|
return detections, proto
|
|
|
|
|
|
def decode_end2end_detections(
|
|
detections: np.ndarray,
|
|
class_names: dict[int, str],
|
|
conf_threshold: float,
|
|
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
|
predictions = detections[0] if detections.ndim == 3 else detections
|
|
if predictions.ndim != 2 or predictions.shape[1] < BOX_CHANNELS:
|
|
raise ValueError(
|
|
f"Unexpected detections shape {predictions.shape}. "
|
|
"End-to-end YOLO-seg expects (N, 6+nm) with [x1,y1,x2,y2,conf,cls,...]."
|
|
)
|
|
|
|
boxes = predictions[:, :4]
|
|
confidences = predictions[:, 4]
|
|
class_ids = predictions[:, 5].astype(np.int32)
|
|
mask_coeffs = predictions[:, BOX_CHANNELS:]
|
|
|
|
keep = confidences >= conf_threshold
|
|
if not np.any(keep):
|
|
empty = np.empty((0, 4), dtype=np.float32)
|
|
return empty, empty, np.empty((0,), dtype=np.int32), np.empty((0,), dtype=np.float32)
|
|
|
|
boxes = boxes[keep]
|
|
confidences = confidences[keep]
|
|
class_ids = class_ids[keep]
|
|
mask_coeffs = mask_coeffs[keep]
|
|
|
|
known = np.array([class_id in class_names for class_id in class_ids], dtype=bool)
|
|
if not np.any(known):
|
|
empty = np.empty((0, 4), dtype=np.float32)
|
|
return empty, empty, np.empty((0,), dtype=np.int32), np.empty((0,), dtype=np.float32)
|
|
|
|
return boxes[known], mask_coeffs[known], class_ids[known], confidences[known]
|
|
|
|
|
|
def process_masks(
|
|
proto: np.ndarray,
|
|
mask_coeffs: np.ndarray,
|
|
boxes_letterbox: np.ndarray,
|
|
imgsz: int,
|
|
image_shape: tuple[int, int],
|
|
scale: float,
|
|
pad: tuple[float, float],
|
|
mask_threshold: float,
|
|
) -> list[np.ndarray]:
|
|
proto_maps = proto[0] if proto.ndim == 4 else proto
|
|
mask_dim, proto_h, proto_w = proto_maps.shape
|
|
if mask_coeffs.shape[1] != mask_dim:
|
|
raise ValueError(
|
|
f"Mask coefficient dim {mask_coeffs.shape[1]} does not match proto channels {mask_dim}."
|
|
)
|
|
|
|
flat_proto = proto_maps.reshape(mask_dim, -1)
|
|
masks = mask_coeffs @ flat_proto
|
|
masks = masks.reshape(-1, proto_h, proto_w)
|
|
|
|
pad_x, pad_y = pad
|
|
image_h, image_w = image_shape
|
|
binary_masks: list[np.ndarray] = []
|
|
|
|
for index, mask in enumerate(masks):
|
|
mask_letterbox = cv2.resize(mask, (imgsz, imgsz), interpolation=cv2.INTER_LINEAR)
|
|
|
|
x1, y1, x2, y2 = boxes_letterbox[index]
|
|
x1_i = max(0, int(np.floor(x1)))
|
|
y1_i = max(0, int(np.floor(y1)))
|
|
x2_i = min(imgsz, int(np.ceil(x2)))
|
|
y2_i = min(imgsz, int(np.ceil(y2)))
|
|
cropped = np.zeros_like(mask_letterbox)
|
|
cropped[y1_i:y2_i, x1_i:x2_i] = mask_letterbox[y1_i:y2_i, x1_i:x2_i]
|
|
|
|
top = int(round(pad_y - 0.1))
|
|
left = int(round(pad_x - 0.1))
|
|
bottom = imgsz - int(round(pad_y + 0.1))
|
|
right = imgsz - int(round(pad_x + 0.1))
|
|
unpadded = cropped[top:bottom, left:right]
|
|
if unpadded.size == 0:
|
|
binary_masks.append(np.zeros((image_h, image_w), dtype=np.uint8))
|
|
continue
|
|
|
|
resized = cv2.resize(unpadded, (image_w, image_h), interpolation=cv2.INTER_LINEAR)
|
|
binary_masks.append((resized > mask_threshold).astype(np.uint8))
|
|
|
|
return binary_masks
|
|
|
|
|
|
def build_mask_results(
|
|
boxes: np.ndarray,
|
|
masks: list[np.ndarray],
|
|
class_ids: np.ndarray,
|
|
confidences: np.ndarray,
|
|
class_names: dict[int, str],
|
|
) -> list[dict[str, Any]]:
|
|
results: list[dict[str, Any]] = []
|
|
image_h, image_w = masks[0].shape if masks else (0, 0)
|
|
|
|
for box, mask, class_id, confidence in zip(boxes, masks, class_ids, confidences):
|
|
if int(mask.sum()) == 0:
|
|
continue
|
|
|
|
xtl = max(0, int(np.floor(box[0])))
|
|
ytl = max(0, int(np.floor(box[1])))
|
|
xbr = min(image_w - 1, int(np.ceil(box[2])))
|
|
ybr = min(image_h - 1, int(np.ceil(box[3])))
|
|
if xbr <= xtl or ybr <= ytl:
|
|
ys, xs = np.where(mask > 0)
|
|
if len(xs) == 0:
|
|
continue
|
|
xtl, xbr = int(xs.min()), int(xs.max())
|
|
ytl, ybr = int(ys.min()), int(ys.max())
|
|
|
|
results.append(
|
|
{
|
|
"confidence": str(float(confidence)),
|
|
"label": class_names[int(class_id)],
|
|
"type": "mask",
|
|
"mask": to_cvat_mask([xtl, ytl, xbr, ybr], mask),
|
|
}
|
|
)
|
|
return results
|
|
|
|
|
|
def init_context(context) -> None:
|
|
context.logger.info("Initializing OpenVINO YOLO segmentation handler")
|
|
if not CLASSES_PATH.is_file():
|
|
raise FileNotFoundError(
|
|
f"Classes file not found: {CLASSES_PATH}. "
|
|
"Mount classes.json at /opt/nuclio/classes.json."
|
|
)
|
|
|
|
class_names = load_classes(CLASSES_PATH)
|
|
model_xml = find_model_xml(MODELS_DIR)
|
|
context.logger.info(f"Loading OpenVINO model from {model_xml}")
|
|
|
|
core = Core()
|
|
compiled_model = core.compile_model(model_xml, "CPU")
|
|
input_layer = compiled_model.input(0)
|
|
input_shape = tuple(input_layer.shape)
|
|
imgsz = int(input_shape[2]) if len(input_shape) == 4 else DEFAULT_IMGSZ
|
|
|
|
context.user_data.compiled_model = compiled_model
|
|
context.user_data.class_names = class_names
|
|
context.user_data.imgsz = imgsz
|
|
context.logger.info(f"Ready: classes={class_names}, imgsz={imgsz}")
|
|
|
|
|
|
def handler(context, event):
|
|
try:
|
|
payload = event.body
|
|
if isinstance(payload, (bytes, bytearray)):
|
|
payload = json.loads(payload.decode("utf-8"))
|
|
if isinstance(payload, str):
|
|
payload = json.loads(payload)
|
|
|
|
image_b64 = payload["image"]
|
|
threshold = float(payload.get("threshold", DEFAULT_CONF_THRESHOLD))
|
|
image_bytes = base64.b64decode(image_b64)
|
|
image_rgb = np.array(Image.open(io.BytesIO(image_bytes)).convert("RGB"))
|
|
image_bgr = image_rgb[:, :, ::-1]
|
|
image_hw = (image_bgr.shape[0], image_bgr.shape[1])
|
|
|
|
tensor, scale, pad = preprocess_image(image_bgr, context.user_data.imgsz)
|
|
outputs = context.user_data.compiled_model([tensor])
|
|
detections, proto = split_seg_outputs(outputs)
|
|
|
|
boxes_letterbox, mask_coeffs, class_ids, confidences = decode_end2end_detections(
|
|
detections,
|
|
class_names=context.user_data.class_names,
|
|
conf_threshold=threshold,
|
|
)
|
|
if len(confidences) == 0:
|
|
return context.Response(
|
|
body=json.dumps([]),
|
|
headers={},
|
|
content_type="application/json",
|
|
status_code=200,
|
|
)
|
|
|
|
masks = process_masks(
|
|
proto=proto,
|
|
mask_coeffs=mask_coeffs,
|
|
boxes_letterbox=boxes_letterbox,
|
|
imgsz=context.user_data.imgsz,
|
|
image_shape=image_hw,
|
|
scale=scale,
|
|
pad=pad,
|
|
mask_threshold=DEFAULT_MASK_THRESHOLD,
|
|
)
|
|
boxes = scale_boxes(boxes_letterbox, image_hw, scale, pad)
|
|
results = build_mask_results(
|
|
boxes=boxes,
|
|
masks=masks,
|
|
class_ids=class_ids,
|
|
confidences=confidences,
|
|
class_names=context.user_data.class_names,
|
|
)
|
|
context.logger.info(f"Returning {len(results)} mask detections")
|
|
return context.Response(
|
|
body=json.dumps(results),
|
|
headers={},
|
|
content_type="application/json",
|
|
status_code=200,
|
|
)
|
|
except Exception as exc:
|
|
context.logger.error(f"Segmentation handler failed: {exc}", exc_info=True)
|
|
raise
|