| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960 |
- # Copyright 2025 Yakhyokhuja Valikhujaev
- # Author: Yakhyokhuja Valikhujaev
- # GitHub: https://github.com/yakhyo
- from typing import List, Union
- import cv2
- import numpy as np
- def draw_detections(
- image: np.ndarray,
- bboxes: Union[List[np.ndarray], List[List[float]]],
- scores: Union[np.ndarray, List[float]],
- landmarks: Union[List[np.ndarray], List[List[List[float]]]],
- vis_threshold: float = 0.6,
- ):
- """
- Draws bounding boxes, scores, and landmarks from separate lists onto an image.
- Args:
- image (np.ndarray): The image to draw on.
- bboxes (List[np.ndarray] or List[List[float]]): List of bounding boxes. Each bbox can be
- np.ndarray with shape (4,) or list [x1, y1, x2, y2].
- scores (List[float] or np.ndarray): List or array of confidence scores.
- landmarks (List[np.ndarray] or List[List[List[float]]]): List of landmark sets. Each landmark
- set can be np.ndarray with shape (5, 2) or nested list [[[x,y],...],...].
- vis_threshold (float): Confidence threshold for filtering which detections to draw.
- """
- _colors = [(0, 0, 255), (0, 255, 255), (255, 0, 255), (0, 255, 0), (255, 0, 0)]
- # Filter detections by score
- keep_indices = [i for i, score in enumerate(scores) if score >= vis_threshold]
- # Draw the filtered detections
- for i in keep_indices:
- bbox = np.array(bboxes[i], dtype=np.int32)
- score = scores[i]
- landmark_set = np.array(landmarks[i], dtype=np.int32)
- # Calculate adaptive thickness
- thickness = max(1, int(min(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 100))
- # Draw bounding box
- cv2.rectangle(image, tuple(bbox[:2]), tuple(bbox[2:]), (0, 0, 255), thickness)
- # Draw score
- cv2.putText(
- image,
- f'{score:.2f}',
- (bbox[0], bbox[1] - 10),
- cv2.FONT_HERSHEY_SIMPLEX,
- 0.5,
- (255, 255, 255),
- thickness,
- )
- # Draw landmarks
- for j, point in enumerate(landmark_set):
- cv2.circle(image, tuple(point), thickness + 1, _colors[j], -1)
|
