Coverage for orchestrant/yolo/core/preprocess.py: 77%

1"""Preprocessing for the ONNX detector: letterbox, or a plain resize.""" 

2 

3from __future__ import annotations 

4 

5from typing import TYPE_CHECKING 

6 

7import numpy as np 

8 

9 

10if TYPE_CHECKING: 

11 from collections.abc import Sequence 

12 

13 import onnxruntime as ort 

14 

15 

16def wants_stretch(session: ort.InferenceSession) -> bool: 

17 """Whether the model asks for a plain resize instead of a letterbox. 

18 

19 ModelANT's models say so in their ``kataglyphis.preprocess`` metadata; any 

20 other model keeps the letterbox. 

21 """ 

22 meta = session.get_modelmeta().custom_metadata_map 

23 return meta.get("kataglyphis.preprocess") == "stretch" 

24 

25 

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

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

28 

29 Args: 

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

31 

32 Returns: 

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

34 """ 

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

36 return (640, 640) 

37 

38 height = input_shape[-2] 

39 width = input_shape[-1] 

40 

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

42 return (height, width) 

43 

44 return (640, 640) 

45 

46 

47def _to_blob(image: np.ndarray) -> np.ndarray: 

48 """BGR uint8 (H, W, 3) to an RGB float32 (1, 3, H, W) blob in [0, 1].""" 

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

50 return blob.transpose(2, 0, 1)[np.newaxis, ...] 

51 

52 

53def preprocess( 

54 frame: np.ndarray, 

55 input_size: tuple[int, int] = (640, 640), 

56 *, 

57 stretch: bool = False, 

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

59 """Preprocess a frame for an ONNX detector. 

60 

61 Letterboxes by default: resizes while preserving the aspect ratio and pads 

62 to input_size. With ``stretch`` it resizes each axis on its own instead. 

63 Either way the frame is converted from BGR to RGB and normalized to [0, 1]. 

64 

65 Args: 

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

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

68 stretch: Plain resize to input_size, no padding (see wants_stretch). 

69 

70 Returns: 

71 Tuple of: 

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

73 - scale: Scale factor applied during resize; an (x, y) pair when stretched. 

74 - pad_x: Horizontal padding added. 

75 - pad_y: Vertical padding added. 

76 """ 

77 import cv2 # noqa: PLC0415 - only resizing needs OpenCV, which has no free-threaded wheel 

78 

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

80 

81 if stretch: 

82 resized = cv2.resize(frame, (input_size[1], input_size[0])) 

83 return ( 

84 _to_blob(resized), 

85 (input_size[1] / original_w, input_size[0] / original_h), 

86 0, 

87 0, 

88 ) 

89 

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

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

92 

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

94 

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

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

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

98 

99 return _to_blob(padded), scale, pad_x, pad_y