276 lines
9.1 KiB
Python
276 lines
9.1 KiB
Python
"""Nuclio handler for CVAT automatic annotation using OpenVINO 2025 IR (.xml/.bin).
|
||
|
||
This file combines YOLOv9 inference logic with Nuclio serverless handler structure.
|
||
It loads an OpenVINO Intermediate Representation (IR) model consisting of a
|
||
``.xml`` file (network topology) and a ``.bin`` file (weights).
|
||
|
||
Adjust ``MODEL_XML`` and ``MODEL_BIN`` if your files are located elsewhere.
|
||
"""
|
||
|
||
import base64
|
||
import json
|
||
import os
|
||
from pathlib import Path
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import openvino as ov
|
||
from openvino.preprocess import PrePostProcessor
|
||
from openvino.preprocess import ColorFormat
|
||
from openvino import Layout, Type
|
||
|
||
#MODEL_DIR = "chicken-detection-model-v26n-300e-best-2026-05-02-NEW_openvino_model"
|
||
#MODEL_NAME = "chicken-detection-model-v26n-300e-best-2026-05-02-NEW"
|
||
#MODEL_XML = os.getenv("MODEL_XML", f"/models/{MODEL_DIR}/{MODEL_NAME}.xml")
|
||
#MODEL_BIN = os.getenv("MODEL_BIN", f"/models/{MODEL_DIR}/{MODEL_NAME}.bin")
|
||
MODEL_XML = os.getenv("MODEL_XML","/models/best.xml")
|
||
MODEL_BIN = os.getenv("MODEL_BIN", "/models/best.bin")
|
||
|
||
classes_file = os.getenv("CLASSES_FILE", "/opt/nuclio/classes.json")
|
||
try:
|
||
with open(classes_file, 'r') as f:
|
||
classes_data = json.load(f)
|
||
coconame = [cls['name'] for cls in sorted(classes_data, key=lambda x: x['id'])]
|
||
except (FileNotFoundError, json.JSONDecodeError):
|
||
coconame = ["chicken", "not-chicken", "half-chicken"]
|
||
|
||
|
||
class Yolov9:
|
||
def __init__(
|
||
self, xml_model_path=MODEL_XML, bin_model_path=MODEL_BIN, conf=0.1, nms=0.4
|
||
):
|
||
# Step 1. Initialize OpenVINO Runtime core
|
||
core = ov.Core()
|
||
# Step 2. Read a model
|
||
if bin_model_path:
|
||
model = core.read_model(
|
||
str(Path(xml_model_path)), str(Path(bin_model_path))
|
||
)
|
||
else:
|
||
model = core.read_model(str(Path(xml_model_path)))
|
||
|
||
self.input_shape = model.input(0).shape # NCHW: [1, C, H, W]
|
||
_, _, self.input_height, self.input_width = self.input_shape
|
||
|
||
# Step 3. Initialize Preprocessing for the model
|
||
ppp = PrePostProcessor(model)
|
||
# Specify input image format
|
||
ppp.input().tensor().set_element_type(Type.u8).set_layout(
|
||
Layout("NHWC")
|
||
).set_color_format(ColorFormat.BGR)
|
||
# Specify preprocess pipeline to input image without resizing
|
||
ppp.input().preprocess().convert_element_type(Type.f32).convert_color(
|
||
ColorFormat.RGB
|
||
).scale([255.0, 255.0, 255.0])
|
||
# Specify model's input layout
|
||
ppp.input().model().set_layout(Layout("NCHW"))
|
||
# Specify output results format
|
||
ppp.output().tensor().set_element_type(Type.f32)
|
||
# Embed above steps in the graph
|
||
model = ppp.build()
|
||
|
||
self.compiled_model = core.compile_model(model, "CPU")
|
||
self.conf_thresh = conf
|
||
self.nms_thresh = nms
|
||
self.colors = []
|
||
|
||
# Create random colors
|
||
np.random.seed(42) # Setting seed for reproducibility
|
||
for i in range(len(coconame)):
|
||
color = tuple(np.random.randint(100, 256, size=3))
|
||
self.colors.append(color)
|
||
|
||
def resize_and_pad(self, image):
|
||
old_h, old_w = image.shape[:2]
|
||
ratio = min(self.input_width / old_w, self.input_height / old_h)
|
||
new_w = int(round(old_w * ratio))
|
||
new_h = int(round(old_h * ratio))
|
||
|
||
image = cv2.resize(image, (new_w, new_h))
|
||
|
||
dw = self.input_width - new_w
|
||
dh = self.input_height - new_h
|
||
top = dh // 2
|
||
bottom = dh - top
|
||
left = dw // 2
|
||
right = dw - left
|
||
|
||
color = [114, 114, 114]
|
||
new_im = cv2.copyMakeBorder(
|
||
image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color
|
||
)
|
||
|
||
return new_im, ratio, (left, top)
|
||
|
||
def predict(self, img):
|
||
input_tensor = np.expand_dims(img, 0)
|
||
infer_request = self.compiled_model.create_infer_request()
|
||
infer_request.infer({0: input_tensor})
|
||
output = infer_request.get_output_tensor()
|
||
detections = output.data[0] # [300, 6] end2end: [x1, y1, x2, y2, conf, class_id]
|
||
|
||
boxes = []
|
||
class_ids = []
|
||
confidences = []
|
||
for detection in detections:
|
||
x1, y1, x2, y2, confidence, class_id = detection
|
||
if confidence > self.conf_thresh:
|
||
confidences.append(float(confidence))
|
||
class_ids.append(int(class_id))
|
||
boxes.append(np.array([x1, y1, x2 - x1, y2 - y1]))
|
||
|
||
if not boxes:
|
||
return []
|
||
|
||
indexes = cv2.dnn.NMSBoxes(
|
||
boxes, confidences, self.conf_thresh, self.nms_thresh
|
||
)
|
||
|
||
results = []
|
||
for i in indexes:
|
||
j = i.item()
|
||
box = boxes[j]
|
||
results.append(
|
||
{
|
||
"class_index": class_ids[j],
|
||
"confidence": confidences[j],
|
||
"box": np.array([box[0], box[1], box[0] + box[2], box[1] + box[3]]),
|
||
}
|
||
)
|
||
|
||
return results
|
||
|
||
def draw(self, img, detections, ratio, pad_left, pad_top):
|
||
for detection in detections:
|
||
box = detection["box"]
|
||
classId = detection["class_index"]
|
||
confidence = detection["confidence"]
|
||
|
||
xmin = (box[0] - pad_left) / ratio
|
||
ymin = (box[1] - pad_top) / ratio
|
||
xmax = (box[2] - pad_left) / ratio
|
||
ymax = (box[3] - pad_top) / ratio
|
||
|
||
# Drawing detection box
|
||
cv2.rectangle(
|
||
img,
|
||
(int(xmin), int(ymin)),
|
||
(int(xmax), int(ymax)),
|
||
tuple(map(int, self.colors[classId])),
|
||
3,
|
||
)
|
||
|
||
# Detection box text
|
||
class_string = coconame[classId] + " " + str(confidence)[:4]
|
||
text_size, _ = cv2.getTextSize(class_string, cv2.FONT_HERSHEY_DUPLEX, 1, 2)
|
||
text_rect = (xmin, ymin - 40, text_size[0] + 10, text_size[1] + 20)
|
||
cv2.rectangle(
|
||
img,
|
||
(int(text_rect[0]), int(text_rect[1])),
|
||
(int(text_rect[0] + text_rect[2]), int(text_rect[1] + text_rect[3])),
|
||
tuple(map(int, self.colors[classId])),
|
||
cv2.FILLED,
|
||
)
|
||
cv2.putText(
|
||
img,
|
||
class_string,
|
||
(int(xmin + 5), int(ymin - 10)),
|
||
cv2.FONT_HERSHEY_DUPLEX,
|
||
1,
|
||
(0, 0, 0),
|
||
2,
|
||
cv2.LINE_AA,
|
||
)
|
||
|
||
|
||
def init_context(context):
|
||
"""Nuclio init_context – called once per container.
|
||
|
||
Loads the IR model and compiles it for the CPU.
|
||
"""
|
||
context.logger.info("Init context ----> 0%")
|
||
model = Yolov9(MODEL_XML, MODEL_BIN, conf=0.1, nms=0.4)
|
||
context.user_data.model = model
|
||
context.logger.info("Init context ----> 100%")
|
||
|
||
|
||
def handler(context, event):
|
||
"""Nuclio handler – called for every request.
|
||
|
||
Expects a JSON body with a base‑64 encoded image under the key ``"image"``.
|
||
Returns a CVAT‑compatible JSON with detected objects.
|
||
"""
|
||
context.logger.info("Run OpenVINO YOLOv9 model")
|
||
|
||
# Parse request body
|
||
try:
|
||
data = event.body
|
||
image_b64 = data["image"]
|
||
except Exception as exc:
|
||
context.logger.error(f"Invalid request body: {exc}")
|
||
return context.Response(
|
||
body=json.dumps({"error": "Invalid request body"}),
|
||
status_code=400,
|
||
content_type="application/json",
|
||
)
|
||
|
||
# Decode image
|
||
image_bytes = base64.b64decode(image_b64)
|
||
image = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_COLOR)
|
||
if image is None:
|
||
context.logger.error("Failed to decode image")
|
||
return context.Response(
|
||
body=json.dumps({"error": "Failed to decode image"}),
|
||
status_code=400,
|
||
content_type="application/json",
|
||
)
|
||
|
||
# Get model from context
|
||
model = context.user_data.model
|
||
|
||
print("Prepare Model")
|
||
|
||
# Preprocess: resize and pad
|
||
img_resized, ratio, (pad_left, pad_top) = model.resize_and_pad(image)
|
||
|
||
# Inference
|
||
detections = model.predict(img_resized)
|
||
|
||
# Convert detections to CVAT-compatible format
|
||
shapes = []
|
||
for detection in detections:
|
||
class_id = detection["class_index"]
|
||
confidence = float(detection["confidence"])
|
||
box = detection["box"]
|
||
|
||
x1 = (box[0] - pad_left) / ratio
|
||
y1 = (box[1] - pad_top) / ratio
|
||
x2 = (box[2] - pad_left) / ratio
|
||
y2 = (box[3] - pad_top) / ratio
|
||
|
||
# Convert to pixel coordinates
|
||
x_min_px = int(max(0, x1))
|
||
y_min_px = int(max(0, y1))
|
||
x_max_px = int(min(image.shape[1], x2))
|
||
y_max_px = int(min(image.shape[0], y2))
|
||
|
||
label = coconame[class_id] if class_id < len(coconame) else "unknown"
|
||
|
||
shapes.append(
|
||
{
|
||
"label": label,
|
||
"points": [x_min_px, y_min_px, x_max_px, y_max_px],
|
||
"type": "rectangle",
|
||
"confidence": str(confidence),
|
||
}
|
||
)
|
||
|
||
context.logger.info(f"Detected {len(shapes)} objects")
|
||
|
||
return context.Response(
|
||
body=json.dumps(shapes),
|
||
headers={},
|
||
content_type="application/json",
|
||
status_code=200,
|
||
)
|