SN44 SafetyVision miner trial artifact

1b0739e verified 19 days ago

7.74 kB

import os
from pathlib import Path

# These must be set before NumPy/Ultralytics load their native runtimes. Score's
# sandbox permits only 16 processes/threads, while the host libraries otherwise
# size pools from the underlying machine rather than the two-vCPU quota.
for thread_variable in (
    "OMP_NUM_THREADS",
    "MKL_NUM_THREADS",
    "OPENBLAS_NUM_THREADS",
    "NUMEXPR_NUM_THREADS",
    "VECLIB_MAXIMUM_THREADS",
    "BLIS_NUM_THREADS",
):
    os.environ[thread_variable] = "1"

import cv2
import numpy as np
import onnxruntime as ort
from numpy import ndarray
from pydantic import BaseModel


class BoundingBox(BaseModel):
    x1: int
    y1: int
    x2: int
    y2: int
    cls_id: int
    conf: float


class Polygon(BaseModel):
    cls_id: int
    conf: float
    points: list[tuple[int, int]]


class TVFrameResult(BaseModel):
    frame_id: int
    boxes: list[BoundingBox] | None = None
    polygons: list[Polygon] | None = None
    keypoints: list[tuple[int, int]] | None = None


class Miner:
    """SN44 PPE detector using a 320px SafetyVision v1 ONNX model."""

    # SafetyVision class ID -> active manifest object index:
    # [helmet, vest, gloves, eyewear]
    SOURCE_TO_MANIFEST = {1: 2, 2: 3, 3: 0, 12: 1}
    CONFIDENCE = 0.10
    NMS_IOU = 0.70
    IMAGE_SIZE = 320
    MAX_DETECTIONS = 300

    def __init__(self, path_hf_repo: Path) -> None:
        cv2.setNumThreads(1)
        options = ort.SessionOptions()
        options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
        options.intra_op_num_threads = 1
        options.inter_op_num_threads = 1
        options.add_session_config_entry("session.intra_op.allow_spinning", "0")
        self.session = ort.InferenceSession(
            str(path_hf_repo / "best.onnx"),
            sess_options=options,
            providers=["CPUExecutionProvider"],
        )
        self.input_name = self.session.get_inputs()[0].name

        # This graph was qualified on CPU. Keeping the same provider also
        # avoids provider-dependent output drift during the
        # validator's standardized two-vCPU compliance run.
        self.device = "cpu"

    def __repr__(self) -> str:
        return (
            "SafetyVision-v1-320 PPE ONNX detector "
            f"(confidence={self.CONFIDENCE}, device={self.device})"
        )

    @classmethod
    def _prepare_image(cls, image: ndarray) -> tuple[ndarray, float, tuple[int, int]]:
        height, width = image.shape[:2]
        scale = min(cls.IMAGE_SIZE / height, cls.IMAGE_SIZE / width)
        resized_width = round(width * scale)
        resized_height = round(height * scale)
        if (resized_width, resized_height) != (width, height):
            image = cv2.resize(
                image,
                (resized_width, resized_height),
                interpolation=cv2.INTER_LINEAR,
            )
        horizontal = (cls.IMAGE_SIZE - resized_width) / 2
        vertical = (cls.IMAGE_SIZE - resized_height) / 2
        left, right = round(horizontal - 0.1), round(horizontal + 0.1)
        top, bottom = round(vertical - 0.1), round(vertical + 0.1)
        image = cv2.copyMakeBorder(
            image,
            top,
            bottom,
            left,
            right,
            cv2.BORDER_CONSTANT,
            value=(114, 114, 114),
        )
        tensor = np.ascontiguousarray(image[:, :, ::-1].transpose(2, 0, 1))
        tensor = tensor.astype(np.float32) / 255.0
        return tensor[None], scale, (left, top)

    @staticmethod
    def _nms(boxes: ndarray, scores: ndarray, threshold: float) -> list[int]:
        if len(boxes) == 0:
            return []
        order = np.argsort(-scores)
        areas = np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) * np.maximum(
            0.0, boxes[:, 3] - boxes[:, 1]
        )
        keep: list[int] = []
        while order.size:
            current = int(order[0])
            keep.append(current)
            if order.size == 1:
                break
            remaining = order[1:]
            x1 = np.maximum(boxes[current, 0], boxes[remaining, 0])
            y1 = np.maximum(boxes[current, 1], boxes[remaining, 1])
            x2 = np.minimum(boxes[current, 2], boxes[remaining, 2])
            y2 = np.minimum(boxes[current, 3], boxes[remaining, 3])
            intersection = np.maximum(0.0, x2 - x1) * np.maximum(0.0, y2 - y1)
            union = areas[current] + areas[remaining] - intersection
            overlap = np.divide(
                intersection,
                union,
                out=np.zeros_like(intersection),
                where=union > 0,
            )
            order = remaining[overlap <= threshold]
        return keep

    def _detect(self, image: ndarray) -> list[BoundingBox]:
        height, width = image.shape[:2]
        tensor, scale, (pad_x, pad_y) = self._prepare_image(image)
        prediction = self.session.run(None, {self.input_name: tensor})[0][0].T
        source_boxes = prediction[:, :4]
        class_scores = prediction[:, 4:]
        source_classes = np.argmax(class_scores, axis=1)
        confidences = np.max(class_scores, axis=1)
        target_ids = np.fromiter(self.SOURCE_TO_MANIFEST, dtype=np.int64)
        selected = (confidences > self.CONFIDENCE) & np.isin(source_classes, target_ids)
        source_boxes = source_boxes[selected]
        source_classes = source_classes[selected]
        confidences = confidences[selected]
        if not len(source_boxes):
            return []

        boxes = np.empty_like(source_boxes)
        boxes[:, 0] = source_boxes[:, 0] - source_boxes[:, 2] / 2
        boxes[:, 1] = source_boxes[:, 1] - source_boxes[:, 3] / 2
        boxes[:, 2] = source_boxes[:, 0] + source_boxes[:, 2] / 2
        boxes[:, 3] = source_boxes[:, 1] + source_boxes[:, 3] / 2

        kept: list[int] = []
        for source_class in self.SOURCE_TO_MANIFEST:
            indices = np.flatnonzero(source_classes == source_class)
            local = self._nms(boxes[indices], confidences[indices], self.NMS_IOU)
            kept.extend(int(indices[index]) for index in local)
        kept.sort(key=lambda index: float(confidences[index]), reverse=True)

        output: list[BoundingBox] = []
        for index in kept[: self.MAX_DETECTIONS]:
            x1, y1, x2, y2 = boxes[index]
            x1 = np.clip((x1 - pad_x) / scale, 0, width)
            y1 = np.clip((y1 - pad_y) / scale, 0, height)
            x2 = np.clip((x2 - pad_x) / scale, 0, width)
            y2 = np.clip((y2 - pad_y) / scale, 0, height)
            output.append(
                BoundingBox(
                    x1=round(float(x1)),
                    y1=round(float(y1)),
                    x2=round(float(x2)),
                    y2=round(float(y2)),
                    cls_id=self.SOURCE_TO_MANIFEST[int(source_classes[index])],
                    conf=float(confidences[index]),
                )
            )
        return output

    def predict_batch(
        self,
        batch_images: list[ndarray],
        offset: int,
        n_keypoints: int,
    ) -> list[TVFrameResult]:
        results: list[TVFrameResult] = []
        # The published ONNX graph has a fixed batch dimension of one. The
        # Chute contract may provide multiple frames, so execute them in order
        # rather than passing an invalid larger tensor to ONNX Runtime.
        for index, image in enumerate(batch_images):
            boxes = self._detect(image)
            results.append(
                TVFrameResult(
                    frame_id=offset + index,
                    boxes=boxes,
                    polygons=[],
                    keypoints=[(0, 0) for _ in range(n_keypoints)],
                )
            )
        return results

	import os
	from pathlib import Path

	# These must be set before NumPy/Ultralytics load their native runtimes. Score's
	# sandbox permits only 16 processes/threads, while the host libraries otherwise
	# size pools from the underlying machine rather than the two-vCPU quota.
	for thread_variable in (
	"OMP_NUM_THREADS",
	"MKL_NUM_THREADS",
	"OPENBLAS_NUM_THREADS",
	"NUMEXPR_NUM_THREADS",
	"VECLIB_MAXIMUM_THREADS",
	"BLIS_NUM_THREADS",
	):
	os.environ[thread_variable] = "1"

	import cv2
	import numpy as np
	import onnxruntime as ort
	from numpy import ndarray
	from pydantic import BaseModel


	class BoundingBox(BaseModel):
	x1: int
	y1: int
	x2: int
	y2: int
	cls_id: int
	conf: float


	class Polygon(BaseModel):
	cls_id: int
	conf: float
	points: list[tuple[int, int]]


	class TVFrameResult(BaseModel):
	frame_id: int
	boxes: list[BoundingBox] \| None = None
	polygons: list[Polygon] \| None = None
	keypoints: list[tuple[int, int]] \| None = None


	class Miner:
	"""SN44 PPE detector using a 320px SafetyVision v1 ONNX model."""

	# SafetyVision class ID -> active manifest object index:
	# [helmet, vest, gloves, eyewear]
	SOURCE_TO_MANIFEST = {1: 2, 2: 3, 3: 0, 12: 1}
	CONFIDENCE = 0.10
	NMS_IOU = 0.70
	IMAGE_SIZE = 320
	MAX_DETECTIONS = 300

	def __init__(self, path_hf_repo: Path) -> None:
	cv2.setNumThreads(1)
	options = ort.SessionOptions()
	options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
	options.intra_op_num_threads = 1
	options.inter_op_num_threads = 1
	options.add_session_config_entry("session.intra_op.allow_spinning", "0")
	self.session = ort.InferenceSession(
	str(path_hf_repo / "best.onnx"),
	sess_options=options,
	providers=["CPUExecutionProvider"],
	)
	self.input_name = self.session.get_inputs()[0].name

	# This graph was qualified on CPU. Keeping the same provider also
	# avoids provider-dependent output drift during the
	# validator's standardized two-vCPU compliance run.
	self.device = "cpu"

	def __repr__(self) -> str:
	return (
	"SafetyVision-v1-320 PPE ONNX detector "
	f"(confidence={self.CONFIDENCE}, device={self.device})"
	)

	@classmethod
	def _prepare_image(cls, image: ndarray) -> tuple[ndarray, float, tuple[int, int]]:
	height, width = image.shape[:2]
	scale = min(cls.IMAGE_SIZE / height, cls.IMAGE_SIZE / width)
	resized_width = round(width * scale)
	resized_height = round(height * scale)
	if (resized_width, resized_height) != (width, height):
	image = cv2.resize(
	image,
	(resized_width, resized_height),
	interpolation=cv2.INTER_LINEAR,
	)
	horizontal = (cls.IMAGE_SIZE - resized_width) / 2
	vertical = (cls.IMAGE_SIZE - resized_height) / 2
	left, right = round(horizontal - 0.1), round(horizontal + 0.1)
	top, bottom = round(vertical - 0.1), round(vertical + 0.1)
	image = cv2.copyMakeBorder(
	image,
	top,
	bottom,
	left,
	right,
	cv2.BORDER_CONSTANT,
	value=(114, 114, 114),
	)
	tensor = np.ascontiguousarray(image[:, :, ::-1].transpose(2, 0, 1))
	tensor = tensor.astype(np.float32) / 255.0
	return tensor[None], scale, (left, top)

	@staticmethod
	def _nms(boxes: ndarray, scores: ndarray, threshold: float) -> list[int]:
	if len(boxes) == 0:
	return []
	order = np.argsort(-scores)
	areas = np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) * np.maximum(
	0.0, boxes[:, 3] - boxes[:, 1]
	)
	keep: list[int] = []
	while order.size:
	current = int(order[0])
	keep.append(current)
	if order.size == 1:
	break
	remaining = order[1:]
	x1 = np.maximum(boxes[current, 0], boxes[remaining, 0])
	y1 = np.maximum(boxes[current, 1], boxes[remaining, 1])
	x2 = np.minimum(boxes[current, 2], boxes[remaining, 2])
	y2 = np.minimum(boxes[current, 3], boxes[remaining, 3])
	intersection = np.maximum(0.0, x2 - x1) * np.maximum(0.0, y2 - y1)
	union = areas[current] + areas[remaining] - intersection
	overlap = np.divide(
	intersection,
	union,
	out=np.zeros_like(intersection),
	where=union > 0,
	)
	order = remaining[overlap <= threshold]
	return keep

	def _detect(self, image: ndarray) -> list[BoundingBox]:
	height, width = image.shape[:2]
	tensor, scale, (pad_x, pad_y) = self._prepare_image(image)
	prediction = self.session.run(None, {self.input_name: tensor})[0][0].T
	source_boxes = prediction[:, :4]
	class_scores = prediction[:, 4:]
	source_classes = np.argmax(class_scores, axis=1)
	confidences = np.max(class_scores, axis=1)
	target_ids = np.fromiter(self.SOURCE_TO_MANIFEST, dtype=np.int64)
	selected = (confidences > self.CONFIDENCE) & np.isin(source_classes, target_ids)
	source_boxes = source_boxes[selected]
	source_classes = source_classes[selected]
	confidences = confidences[selected]
	if not len(source_boxes):
	return []

	boxes = np.empty_like(source_boxes)
	boxes[:, 0] = source_boxes[:, 0] - source_boxes[:, 2] / 2
	boxes[:, 1] = source_boxes[:, 1] - source_boxes[:, 3] / 2
	boxes[:, 2] = source_boxes[:, 0] + source_boxes[:, 2] / 2
	boxes[:, 3] = source_boxes[:, 1] + source_boxes[:, 3] / 2

	kept: list[int] = []
	for source_class in self.SOURCE_TO_MANIFEST:
	indices = np.flatnonzero(source_classes == source_class)
	local = self._nms(boxes[indices], confidences[indices], self.NMS_IOU)
	kept.extend(int(indices[index]) for index in local)
	kept.sort(key=lambda index: float(confidences[index]), reverse=True)

	output: list[BoundingBox] = []
	for index in kept[: self.MAX_DETECTIONS]:
	x1, y1, x2, y2 = boxes[index]
	x1 = np.clip((x1 - pad_x) / scale, 0, width)
	y1 = np.clip((y1 - pad_y) / scale, 0, height)
	x2 = np.clip((x2 - pad_x) / scale, 0, width)
	y2 = np.clip((y2 - pad_y) / scale, 0, height)
	output.append(
	BoundingBox(
	x1=round(float(x1)),
	y1=round(float(y1)),
	x2=round(float(x2)),
	y2=round(float(y2)),
	cls_id=self.SOURCE_TO_MANIFEST[int(source_classes[index])],
	conf=float(confidences[index]),
	)
	)
	return output

	def predict_batch(
	self,
	batch_images: list[ndarray],
	offset: int,
	n_keypoints: int,
	) -> list[TVFrameResult]:
	results: list[TVFrameResult] = []
	# The published ONNX graph has a fixed batch dimension of one. The
	# Chute contract may provide multiple frames, so execute them in order
	# rather than passing an invalid larger tensor to ONNX Runtime.
	for index, image in enumerate(batch_images):
	boxes = self._detect(image)
	results.append(
	TVFrameResult(
	frame_id=offset + index,
	boxes=boxes,
	polygons=[],
	keypoints=[(0, 0) for _ in range(n_keypoints)],
	)
	)
	return results