Coverage for orchestr_ant_ion / yolo / core / preprocess.py: 0%

1"""Preprocessing utilities for YOLO inference.""" 

2 

3from __future__ import annotations 

4 

5from typing import TYPE_CHECKING 

6 

7import cv2 

8import numpy as np 

9 

10 

11if TYPE_CHECKING: 

12 from collections.abc import Sequence 

13 

14 

15def infer_input_size(input_shape: Sequence[object] | None) -> tuple[int, int]: 

16 """Infer (height, width) from an ONNX input shape. 

17 

18 Args: 

19 input_shape: Shape array from ONNX model input, typically [batch, channels, height, width]. 

20 

21 Returns: 

22 Tuple of (height, width). Defaults to (640, 640) if shape cannot be inferred. 

23 """ 

24 if not input_shape or len(input_shape) < 4: 

25 return (640, 640) 

26 

27 height = input_shape[-2] 

28 width = input_shape[-1] 

29 

30 if isinstance(height, int) and isinstance(width, int): 

31 return (height, width) 

32 

33 return (640, 640) 

34 

35 

36def preprocess( 

37 frame: np.ndarray, input_size: tuple[int, int] = (640, 640) 

38) -> tuple[np.ndarray, float, int, int]: 

39 """Preprocess frame for YOLOv10 inference. 

40 

41 Resizes frame while preserving aspect ratio, pads to input_size, 

42 converts BGR to RGB, and normalizes to [0, 1]. 

43 

44 Args: 

45 frame: Input image in BGR format, shape (H, W, 3). 

46 input_size: Target size as (height, width). Defaults to (640, 640). 

47 

48 Returns: 

49 Tuple of: 

50 - blob: Preprocessed image blob, shape (1, 3, H, W), float32. 

51 - scale: Scale factor applied during resize. 

52 - pad_x: Horizontal padding added. 

53 - pad_y: Vertical padding added. 

54 """ 

55 original_h, original_w = frame.shape[:2] 

56 

57 scale = min(input_size[0] / original_h, input_size[1] / original_w) 

58 new_w, new_h = int(original_w * scale), int(original_h * scale) 

59 

60 resized = cv2.resize(frame, (new_w, new_h)) 

61 

62 padded = np.full((input_size[0], input_size[1], 3), 114, dtype=np.uint8) 

63 pad_x, pad_y = (input_size[1] - new_w) // 2, (input_size[0] - new_h) // 2 

64 padded[pad_y : pad_y + new_h, pad_x : pad_x + new_w] = resized 

65 

66 blob = padded[:, :, ::-1].astype(np.float32) / 255.0 

67 blob = blob.transpose(2, 0, 1)[np.newaxis, ...] 

68 

69 return blob, scale, pad_x, pad_y