Coverage for orchestrant/yolo/core/postprocess.py: 68%
171 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-19 11:22 +0000
1"""Post-processing utilities for YOLO model outputs."""
3from __future__ import annotations
5from dataclasses import dataclass
6from typing import TYPE_CHECKING
8import numpy as np
9from loguru import logger
11from orchestrant.pipeline.constants import POSTPROCESS_DEFAULT_CONF_THRESHOLD
12from orchestrant.yolo.core.constants import CLASS_NAMES
15if TYPE_CHECKING:
16 from collections.abc import Sequence
19@dataclass
20class DecodeConfig:
21 """Configuration for post-processing model outputs.
23 ``scale`` is one factor for a letterbox, an (x, y) pair for a plain resize.
24 """
26 scale: float | tuple[float, float]
27 pad_x: int
28 pad_y: int
29 input_size: tuple[int, int]
30 conf_threshold: float = POSTPROCESS_DEFAULT_CONF_THRESHOLD
31 debug_boxes: bool = False
34def _softmax(scores: np.ndarray) -> np.ndarray:
35 """Apply softmax normalization to a score array.
37 Args:
38 scores: Raw score array (any shape).
40 Returns:
41 Normalized probability array with same shape.
42 """
43 shifted = scores - np.max(scores)
44 exp_scores = np.exp(shifted)
45 return exp_scores / np.sum(exp_scores)
48def _xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray:
49 """Convert bounding boxes from center format to corner format.
51 Args:
52 boxes: Array of shape (N, 4) with [cx, cy, w, h] format.
54 Returns:
55 Array of shape (N, 4) with [x1, y1, x2, y2] format.
56 """
57 x, y, w, h = boxes.T
58 x1 = x - w / 2
59 y1 = y - h / 2
60 x2 = x + w / 2
61 y2 = y + h / 2
62 return np.stack([x1, y1, x2, y2], axis=1)
65def _sigmoid(values: np.ndarray) -> np.ndarray:
66 """Apply sigmoid activation to values.
68 Args:
69 values: Input array (any shape).
71 Returns:
72 Sigmoid-activated array with same shape.
73 """
74 return 1.0 / (1.0 + np.exp(-values))
77def _squeeze_to_2d(arr: np.ndarray) -> np.ndarray:
78 """Reduce array dimensions to 2D by removing singleton dimensions.
80 Args:
81 arr: Input array with potential singleton dimensions.
83 Returns:
84 2D array with singleton dimensions removed.
85 """
86 data = np.asarray(arr)
87 data = np.squeeze(data)
88 if data.ndim == 3 and data.shape[0] == 1:
89 data = data[0]
90 return data
93def _looks_like_xywh(boxes: np.ndarray) -> bool:
94 """Detect if boxes are in center format (xywh) vs corner format (xyxy).
96 Heuristic: if many boxes have x2 < x1 or y2 < y1, they're likely xywh.
98 Args:
99 boxes: Array of shape (N, 4).
101 Returns:
102 True if boxes appear to be in center format.
103 """
104 if boxes.size == 0:
105 return False
106 x2_lt_x1 = np.mean(boxes[:, 2] < boxes[:, 0])
107 y2_lt_y1 = np.mean(boxes[:, 3] < boxes[:, 1])
108 return bool((x2_lt_x1 > 0.3) or (y2_lt_y1 > 0.3))
111def _decode_classification(output: np.ndarray) -> dict | None:
112 """Decode classification model output to class prediction.
114 Args:
115 output: Model output array.
117 Returns:
118 Dictionary with class_id, score, and label, or None if not classification.
119 """
120 if output.ndim == 1 or (output.ndim == 2 and output.shape[0] == 1):
121 scores = output if output.ndim == 1 else output[0]
122 scores = scores.astype(np.float32)
123 if np.any(scores < 0) or np.any(scores > 1):
124 scores = _softmax(scores)
125 class_id = int(np.argmax(scores))
126 score = float(scores[class_id])
127 label = CLASS_NAMES[class_id] if class_id < len(CLASS_NAMES) else str(class_id)
128 return {
129 "class_id": class_id,
130 "score": score,
131 "label": label,
132 }
133 return None
136def _log_debug_boxes(
137 boxes: np.ndarray,
138 scores: np.ndarray,
139 class_ids: np.ndarray,
140) -> None:
141 """Log the first 3 decoded boxes and their scores for debugging.
143 Args:
144 boxes: Bounding box array.
145 scores: Confidence score array.
146 class_ids: Class ID array.
147 """
148 logger.info(
149 "Decoded boxes (first 3): {}",
150 boxes[:3].round(2).tolist(),
151 )
152 logger.info(
153 "Decoded scores/classes (first 3): {}",
154 list(
155 zip(
156 scores[:3].round(3).tolist(),
157 class_ids[:3].tolist(),
158 strict=False,
159 )
160 ),
161 )
164def _unscale_and_collect(
165 boxes: np.ndarray,
166 scores: np.ndarray,
167 class_ids: np.ndarray,
168 scale: float | tuple[float, float],
169 pad_x: int,
170 pad_y: int,
171 conf_threshold: float,
172) -> list[dict]:
173 """Unscale bounding boxes and collect detections above the threshold.
175 Args:
176 boxes: Bounding boxes in input image coordinates.
177 scores: Confidence scores for each detection.
178 class_ids: Class IDs for each detection.
179 scale: Scale factor applied during preprocessing; (x, y) when per axis.
180 pad_x: Horizontal padding applied during preprocessing.
181 pad_y: Vertical padding applied during preprocessing.
182 conf_threshold: Minimum confidence threshold.
184 Returns:
185 List of detection dictionaries with bbox, score, and class_id.
186 """
187 scale_x, scale_y = scale if isinstance(scale, tuple) else (scale, scale)
188 detections: list[dict] = []
189 for box, score, class_id in zip(boxes, scores, class_ids, strict=False):
190 if score < conf_threshold:
191 continue
192 x1, y1, x2, y2 = box
193 x1 = (x1 - pad_x) / scale_x
194 y1 = (y1 - pad_y) / scale_y
195 x2 = (x2 - pad_x) / scale_x
196 y2 = (y2 - pad_y) / scale_y
197 detections.append(
198 {
199 "bbox": [int(x1), int(y1), int(x2), int(y2)],
200 "score": float(score),
201 "class_id": int(class_id),
202 }
203 )
204 return detections
207def _prepare_boxes(
208 boxes: np.ndarray,
209 input_size: tuple[int, int],
210 *,
211 convert_xywh: bool = True,
212) -> np.ndarray:
213 """Scale boxes to pixel coordinates and optionally convert format.
215 Args:
216 boxes: Raw bounding boxes.
217 input_size: (height, width) of input image.
218 convert_xywh: Whether to convert from xywh to xyxy format.
220 Returns:
221 Processed bounding boxes in pixel coordinates.
222 """
223 height, width = input_size
224 if np.max(boxes) <= 1.5:
225 boxes = boxes * np.array([width, height, width, height], dtype=np.float32)
227 if convert_xywh and _looks_like_xywh(boxes):
228 boxes = _xywh_to_xyxy(boxes)
230 return boxes
233def _decode_triplet_outputs(
234 outputs: Sequence[np.ndarray],
235 input_size: tuple[int, int],
236 scale: float | tuple[float, float],
237 pad_x: int,
238 pad_y: int,
239 conf_threshold: float,
240 *,
241 debug_boxes: bool,
242) -> list[dict] | None:
243 """Decode triplet output format (boxes, scores, class_ids).
245 Args:
246 outputs: Model output tensors.
247 input_size: Input image dimensions.
248 scale: Preprocessing scale factor.
249 pad_x: Horizontal padding.
250 pad_y: Vertical padding.
251 conf_threshold: Confidence threshold.
252 debug_boxes: Whether to log debug info.
254 Returns:
255 List of detections or None if not triplet format.
256 """
257 if len(outputs) < 3:
258 return None
260 boxes = _squeeze_to_2d(outputs[0])
261 scores = _squeeze_to_2d(outputs[1])
262 class_ids = _squeeze_to_2d(outputs[2])
264 if not (boxes.ndim == 2 and boxes.shape[-1] == 4):
265 return None
267 if scores.ndim > 1:
268 scores = scores.reshape(-1)
269 if class_ids.ndim > 1:
270 class_ids = class_ids.reshape(-1)
272 boxes = _prepare_boxes(boxes, input_size, convert_xywh=True)
274 if debug_boxes:
275 _log_debug_boxes(boxes, scores, class_ids)
277 return _unscale_and_collect(
278 boxes, scores, class_ids, scale, pad_x, pad_y, conf_threshold
279 )
282def _decode_pair_outputs(
283 outputs: Sequence[np.ndarray],
284 input_size: tuple[int, int],
285 scale: float | tuple[float, float],
286 pad_x: int,
287 pad_y: int,
288 conf_threshold: float,
289 *,
290 debug_boxes: bool,
291) -> list[dict] | None:
292 """Decode pair output format (scores, boxes).
294 Args:
295 outputs: Model output tensors.
296 input_size: Input image dimensions.
297 scale: Preprocessing scale factor.
298 pad_x: Horizontal padding.
299 pad_y: Vertical padding.
300 conf_threshold: Confidence threshold.
301 debug_boxes: Whether to log debug info.
303 Returns:
304 List of detections or None if not pair format.
305 """
306 if len(outputs) < 2:
307 return None
309 scores = _squeeze_to_2d(outputs[0])
310 boxes = _squeeze_to_2d(outputs[1])
312 if not (scores.ndim == 2 and boxes.ndim == 2 and boxes.shape[1] == 4):
313 return None
315 height, width = input_size
316 probs = _softmax(scores) if scores.shape[1] > 1 else scores
317 class_ids = np.argmax(probs, axis=1)
318 confs = probs[np.arange(len(class_ids)), class_ids]
320 if np.min(boxes) < 0.0 or np.max(boxes) > 1.5:
321 boxes = _sigmoid(boxes)
322 boxes = boxes * np.array([width, height, width, height], dtype=np.float32)
323 boxes = _xywh_to_xyxy(boxes)
325 if debug_boxes:
326 _log_debug_boxes(boxes, confs, class_ids)
328 return _unscale_and_collect(
329 boxes, confs, class_ids, scale, pad_x, pad_y, conf_threshold
330 )
333def _normalize_output_data(output: np.ndarray) -> np.ndarray:
334 """Normalize output data to standard 2D format.
336 Args:
337 output: Raw model output.
339 Returns:
340 Normalized 2D array.
341 """
342 data = output
343 if data.ndim == 3:
344 if data.shape[0] == 1:
345 data = data[0]
346 elif data.shape[0] > 1 and data.shape[-1] == 1:
347 data = np.squeeze(data, axis=-1)
349 if data.ndim == 2 and data.shape[0] < data.shape[1] and data.shape[1] >= 6:
350 data = data.T
351 return data
354def _decode_scores_and_boxes(
355 data: np.ndarray,
356 input_size: tuple[int, int],
357) -> tuple[np.ndarray, np.ndarray, np.ndarray] | None:
358 """Extract boxes, scores, and class IDs from normalized output.
360 Args:
361 data: Normalized 2D output data.
362 input_size: Input image dimensions.
364 Returns:
365 Tuple of (boxes, scores, class_ids) or None if invalid.
366 """
367 if data.ndim != 2 or data.shape[1] < 6:
368 return None
370 channels = data.shape[1]
371 height, width = input_size
373 if channels == 6:
374 boxes = data[:, :4]
375 scores = data[:, 4]
376 class_ids = data[:, 5].astype(int)
377 else:
378 boxes = data[:, :4]
379 if channels >= 85:
380 obj = data[:, 4]
381 class_scores = data[:, 5:]
382 class_ids = np.argmax(class_scores, axis=1)
383 scores = obj * class_scores[np.arange(len(class_ids)), class_ids]
384 else:
385 class_scores = data[:, 4:]
386 class_ids = np.argmax(class_scores, axis=1)
387 scores = class_scores[np.arange(len(class_ids)), class_ids]
389 if np.max(boxes) <= 1.5:
390 boxes = boxes * np.array([width, height, width, height], dtype=np.float32)
392 if channels != 6 or _looks_like_xywh(boxes):
393 boxes = _xywh_to_xyxy(boxes)
395 return boxes, scores, class_ids
398def _decode_generic_output(
399 output: np.ndarray,
400 input_size: tuple[int, int],
401 scale: float | tuple[float, float],
402 pad_x: int,
403 pad_y: int,
404 conf_threshold: float,
405 *,
406 debug_boxes: bool,
407) -> list[dict]:
408 """Decode generic YOLO output format.
410 Args:
411 output: Raw model output.
412 input_size: Input image dimensions.
413 scale: Preprocessing scale factor.
414 pad_x: Horizontal padding.
415 pad_y: Vertical padding.
416 conf_threshold: Confidence threshold.
417 debug_boxes: Whether to log debug info.
419 Returns:
420 List of detections.
421 """
422 data = _normalize_output_data(output)
423 decoded = _decode_scores_and_boxes(data, input_size)
424 if decoded is None:
425 return []
426 boxes, scores, class_ids = decoded
428 if debug_boxes:
429 _log_debug_boxes(boxes, scores, class_ids)
431 return _unscale_and_collect(
432 boxes, scores, class_ids, scale, pad_x, pad_y, conf_threshold
433 )
436def postprocess(
437 outputs: Sequence[np.ndarray],
438 config: DecodeConfig,
439 *,
440 debug_output: bool = False,
441) -> tuple[list, dict | None]:
442 """Parse model outputs for detection or classification models.
444 Args:
445 outputs: Model output tensors.
446 config: Decoding configuration.
447 debug_output: Whether to log output shapes.
449 Returns:
450 Tuple of (detections list, classification dict or None).
451 """
452 detections: list[dict] = []
453 classification: dict | None = None
455 if outputs is None or len(outputs) == 0:
456 return detections, classification
458 if debug_output:
459 logger.info("Model outputs: {}", [np.asarray(out).shape for out in outputs])
461 output = _squeeze_to_2d(outputs[0])
463 if output.ndim == 0:
464 return detections, classification
466 classification = _decode_classification(output)
467 if classification is not None:
468 return detections, classification
470 triplet = _decode_triplet_outputs(
471 outputs,
472 config.input_size,
473 config.scale,
474 config.pad_x,
475 config.pad_y,
476 config.conf_threshold,
477 debug_boxes=config.debug_boxes,
478 )
479 if triplet is not None:
480 return triplet, classification
482 pair = _decode_pair_outputs(
483 outputs,
484 config.input_size,
485 config.scale,
486 config.pad_x,
487 config.pad_y,
488 config.conf_threshold,
489 debug_boxes=config.debug_boxes,
490 )
491 if pair is not None:
492 return pair, classification
494 detections = _decode_generic_output(
495 output,
496 config.input_size,
497 config.scale,
498 config.pad_x,
499 config.pad_y,
500 config.conf_threshold,
501 debug_boxes=config.debug_boxes,
502 )
504 return detections, classification