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()