testObbOpenCv.py

import argparse
import math
from pathlib import Path

import cv2
import numpy as np

REPO_DIR = Path(__file__).parent.parent
DEFAULT_MODEL = REPO_DIR / "model" / "best.onnx"

CLASSES = {
    0: "Body",
    1: "Lens",
    2: "System",
}

COLORS = {
    0: (76, 175, 80),
    1: (33, 150, 243),
    2: (255, 152, 0),
}


def letterbox(
    image: np.ndarray,
    imageSize: int,
    fillColor: tuple[int, int, int] = (114, 114, 114),
) -> tuple[np.ndarray, float, int, int]:
    height, width = image.shape[:2]
    scale = min(imageSize / width, imageSize / height)

    resizedWidth = int(round(width * scale))
    resizedHeight = int(round(height * scale))

    resized = cv2.resize(
        image,
        (resizedWidth, resizedHeight),
        interpolation=cv2.INTER_LINEAR,
    )

    padX = (imageSize - resizedWidth) // 2
    padY = (imageSize - resizedHeight) // 2

    canvas = np.full(
        (imageSize, imageSize, 3),
        fillColor,
        dtype=np.uint8,
    )
    canvas[
        padY : padY + resizedHeight,
        padX : padX + resizedWidth,
    ] = resized

    return canvas, scale, padX, padY


def parseOutput(
    output: np.ndarray,
    confidenceThreshold: float,
    scale: float,
    padX: int,
    padY: int,
    imageShape: tuple[int, int, int],
) -> tuple[
    list[tuple[tuple[float, float], tuple[float, float], float]], list[float], list[int]
]:
    if output.ndim == 3:
        output = output[0]

    if output.shape[0] < output.shape[1]:
        predictions = output.T
    else:
        predictions = output

    classCount = len(CLASSES)
    expectedValues = 5 + classCount
    if predictions.shape[1] != expectedValues:
        raise ValueError(
            f"Expected {expectedValues} ONNX outputs per prediction, "
            f"got {predictions.shape[1]}"
        )

    imageHeight, imageWidth = imageShape[:2]
    boxes = []
    scores = []
    classIds = []

    for prediction in predictions:
        classScores = prediction[4 : 4 + classCount]
        classId = int(np.argmax(classScores))
        score = float(classScores[classId])
        if score < confidenceThreshold:
            continue

        centerX, centerY, width, height = prediction[:4]
        angleRadians = float(prediction[4 + classCount])

        centerX = (float(centerX) - padX) / scale
        centerY = (float(centerY) - padY) / scale
        width = float(width) / scale
        height = float(height) / scale

        centerX = min(max(centerX, 0.0), float(imageWidth - 1))
        centerY = min(max(centerY, 0.0), float(imageHeight - 1))
        width = max(width, 1.0)
        height = max(height, 1.0)

        boxes.append(
            (
                (centerX, centerY),
                (width, height),
                math.degrees(angleRadians),
            )
        )
        scores.append(score)
        classIds.append(classId)

    return boxes, scores, classIds


def runRotatedNms(
    boxes: list[tuple[tuple[float, float], tuple[float, float], float]],
    scores: list[float],
    classIds: list[int],
    confidenceThreshold: float,
    nmsThreshold: float,
) -> list[int]:
    keptIndexes = []

    for classId in sorted(set(classIds)):
        localIndexes = [
            index for index, boxClassId in enumerate(classIds) if boxClassId == classId
        ]
        localBoxes = [boxes[index] for index in localIndexes]
        localScores = [scores[index] for index in localIndexes]

        selected = cv2.dnn.NMSBoxesRotated(
            localBoxes,
            localScores,
            confidenceThreshold,
            nmsThreshold,
        )

        if len(selected) == 0:
            continue

        for selectedIndex in np.array(selected).flatten():
            keptIndexes.append(localIndexes[int(selectedIndex)])

    return keptIndexes


def drawDetections(
    image: np.ndarray,
    boxes: list[tuple[tuple[float, float], tuple[float, float], float]],
    scores: list[float],
    classIds: list[int],
    indexes: list[int],
) -> np.ndarray:
    output = image.copy()

    for index in indexes:
        classId = classIds[index]
        color = COLORS.get(classId, (255, 255, 255))
        label = f"{CLASSES.get(classId, classId)} {scores[index]:.2f}"

        points = cv2.boxPoints(boxes[index])
        points = np.intp(points)

        cv2.polylines(output, [points], True, color, 2, cv2.LINE_AA)

        labelX = int(points[:, 0].min())
        labelY = int(points[:, 1].min()) - 8
        labelY = max(labelY, 20)

        textSize, baseline = cv2.getTextSize(
            label,
            cv2.FONT_HERSHEY_SIMPLEX,
            0.6,
            2,
        )
        textWidth, textHeight = textSize

        cv2.rectangle(
            output,
            (labelX, labelY - textHeight - baseline),
            (labelX + textWidth + 6, labelY + baseline),
            color,
            -1,
        )
        cv2.putText(
            output,
            label,
            (labelX + 3, labelY),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.6,
            (255, 255, 255),
            2,
            cv2.LINE_AA,
        )

    return output


def getOutputPath(imagePath: Path, outputPath: Path | None) -> Path:
    if outputPath:
        return outputPath

    return imagePath.with_name(f"{imagePath.stem}_obb{imagePath.suffix}")


def parseArgs() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("image", type=Path, help="Image to run inference on.")
    parser.add_argument(
        "--model",
        type=Path,
        default=DEFAULT_MODEL,
        help="Path to the exported YOLO OBB ONNX model.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=None,
        help="Path to save the image with bounding box overlays.",
    )
    parser.add_argument(
        "--imgsz",
        type=int,
        default=960,
        help="Input image size used when exporting the ONNX model.",
    )
    parser.add_argument(
        "--conf",
        type=float,
        default=0.25,
        help="Confidence threshold for detections.",
    )
    parser.add_argument(
        "--iou",
        type=float,
        default=0.45,
        help="Rotated NMS IoU threshold.",
    )
    return parser.parse_args()


def main() -> None:
    args = parseArgs()
    imagePath = args.image
    modelPath = args.model
    outputPath = getOutputPath(imagePath, args.output)

    if not imagePath.exists():
        raise FileNotFoundError(f"Image not found: {imagePath}")

    if not modelPath.exists():
        raise FileNotFoundError(f"ONNX model not found: {modelPath}")

    image = cv2.imread(str(imagePath))
    if image is None:
        raise ValueError(f"Could not read image: {imagePath}")

    inputImage, scale, padX, padY = letterbox(image, args.imgsz)
    blob = cv2.dnn.blobFromImage(
        inputImage,
        scalefactor=1 / 255.0,
        size=(args.imgsz, args.imgsz),
        mean=(0, 0, 0),
        swapRB=True,
        crop=False,
    )

    net = cv2.dnn.readNetFromONNX(str(modelPath))
    net.setInput(blob)
    output = net.forward()

    boxes, scores, classIds = parseOutput(
        output,
        args.conf,
        scale,
        padX,
        padY,
        image.shape,
    )
    keptIndexes = runRotatedNms(
        boxes,
        scores,
        classIds,
        args.conf,
        args.iou,
    )

    result = drawDetections(
        image,
        boxes,
        scores,
        classIds,
        keptIndexes,
    )

    outputPath.parent.mkdir(parents=True, exist_ok=True)
    if not cv2.imwrite(str(outputPath), result):
        raise ValueError(f"Could not write output image: {outputPath}")

    print(f"Detections: {len(keptIndexes)}")
    print(f"Saved: {outputPath}")


if __name__ == "__main__":
    main()