mirror of
https://github.com/yuruotong1/autoMate.git
synced 2025-12-26 05:16:21 +08:00
263 lines
11 KiB
Python
263 lines
11 KiB
Python
# from typing import List, Optional, Union, Tuple
|
|
|
|
# import cv2
|
|
# import numpy as np
|
|
|
|
# from supervision.detection.core import Detections
|
|
# from supervision.draw.color import Color, ColorPalette
|
|
|
|
|
|
# class BoxAnnotator:
|
|
# """
|
|
# A class for drawing bounding boxes on an image using detections provided.
|
|
|
|
# Attributes:
|
|
# color (Union[Color, ColorPalette]): The color to draw the bounding box,
|
|
# can be a single color or a color palette
|
|
# thickness (int): The thickness of the bounding box lines, default is 2
|
|
# text_color (Color): The color of the text on the bounding box, default is white
|
|
# text_scale (float): The scale of the text on the bounding box, default is 0.5
|
|
# text_thickness (int): The thickness of the text on the bounding box,
|
|
# default is 1
|
|
# text_padding (int): The padding around the text on the bounding box,
|
|
# default is 5
|
|
|
|
# """
|
|
|
|
# def __init__(
|
|
# self,
|
|
# color: Union[Color, ColorPalette] = ColorPalette.DEFAULT,
|
|
# thickness: int = 3, # 1 for seeclick 2 for mind2web and 3 for demo
|
|
# text_color: Color = Color.BLACK,
|
|
# text_scale: float = 0.5, # 0.8 for mobile/web, 0.3 for desktop # 0.4 for mind2web
|
|
# text_thickness: int = 2, #1, # 2 for demo
|
|
# text_padding: int = 10,
|
|
# avoid_overlap: bool = True,
|
|
# ):
|
|
# self.color: Union[Color, ColorPalette] = color
|
|
# self.thickness: int = thickness
|
|
# self.text_color: Color = text_color
|
|
# self.text_scale: float = text_scale
|
|
# self.text_thickness: int = text_thickness
|
|
# self.text_padding: int = text_padding
|
|
# self.avoid_overlap: bool = avoid_overlap
|
|
|
|
# def annotate(
|
|
# self,
|
|
# scene: np.ndarray,
|
|
# detections: Detections,
|
|
# labels: Optional[List[str]] = None,
|
|
# skip_label: bool = False,
|
|
# image_size: Optional[Tuple[int, int]] = None,
|
|
# ) -> np.ndarray:
|
|
# """
|
|
# Draws bounding boxes on the frame using the detections provided.
|
|
|
|
# Args:
|
|
# scene (np.ndarray): The image on which the bounding boxes will be drawn
|
|
# detections (Detections): The detections for which the
|
|
# bounding boxes will be drawn
|
|
# labels (Optional[List[str]]): An optional list of labels
|
|
# corresponding to each detection. If `labels` are not provided,
|
|
# corresponding `class_id` will be used as label.
|
|
# skip_label (bool): Is set to `True`, skips bounding box label annotation.
|
|
# Returns:
|
|
# np.ndarray: The image with the bounding boxes drawn on it
|
|
|
|
# Example:
|
|
# ```python
|
|
# import supervision as sv
|
|
|
|
# classes = ['person', ...]
|
|
# image = ...
|
|
# detections = sv.Detections(...)
|
|
|
|
# box_annotator = sv.BoxAnnotator()
|
|
# labels = [
|
|
# f"{classes[class_id]} {confidence:0.2f}"
|
|
# for _, _, confidence, class_id, _ in detections
|
|
# ]
|
|
# annotated_frame = box_annotator.annotate(
|
|
# scene=image.copy(),
|
|
# detections=detections,
|
|
# labels=labels
|
|
# )
|
|
# ```
|
|
# """
|
|
# font = cv2.FONT_HERSHEY_SIMPLEX
|
|
# for i in range(len(detections)):
|
|
# x1, y1, x2, y2 = detections.xyxy[i].astype(int)
|
|
# class_id = (
|
|
# detections.class_id[i] if detections.class_id is not None else None
|
|
# )
|
|
# idx = class_id if class_id is not None else i
|
|
# color = (
|
|
# self.color.by_idx(idx)
|
|
# if isinstance(self.color, ColorPalette)
|
|
# else self.color
|
|
# )
|
|
# cv2.rectangle(
|
|
# img=scene,
|
|
# pt1=(x1, y1),
|
|
# pt2=(x2, y2),
|
|
# color=color.as_bgr(),
|
|
# thickness=self.thickness,
|
|
# )
|
|
# if skip_label:
|
|
# continue
|
|
|
|
# text = (
|
|
# f"{class_id}"
|
|
# if (labels is None or len(detections) != len(labels))
|
|
# else labels[i]
|
|
# )
|
|
|
|
# text_width, text_height = cv2.getTextSize(
|
|
# text=text,
|
|
# fontFace=font,
|
|
# fontScale=self.text_scale,
|
|
# thickness=self.text_thickness,
|
|
# )[0]
|
|
|
|
# if not self.avoid_overlap:
|
|
# text_x = x1 + self.text_padding
|
|
# text_y = y1 - self.text_padding
|
|
|
|
# text_background_x1 = x1
|
|
# text_background_y1 = y1 - 2 * self.text_padding - text_height
|
|
|
|
# text_background_x2 = x1 + 2 * self.text_padding + text_width
|
|
# text_background_y2 = y1
|
|
# # text_x = x1 - self.text_padding - text_width
|
|
# # text_y = y1 + self.text_padding + text_height
|
|
# # text_background_x1 = x1 - 2 * self.text_padding - text_width
|
|
# # text_background_y1 = y1
|
|
# # text_background_x2 = x1
|
|
# # text_background_y2 = y1 + 2 * self.text_padding + text_height
|
|
# else:
|
|
# text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2 = get_optimal_label_pos(self.text_padding, text_width, text_height, x1, y1, x2, y2, detections, image_size)
|
|
|
|
# cv2.rectangle(
|
|
# img=scene,
|
|
# pt1=(text_background_x1, text_background_y1),
|
|
# pt2=(text_background_x2, text_background_y2),
|
|
# color=color.as_bgr(),
|
|
# thickness=cv2.FILLED,
|
|
# )
|
|
# # import pdb; pdb.set_trace()
|
|
# box_color = color.as_rgb()
|
|
# luminance = 0.299 * box_color[0] + 0.587 * box_color[1] + 0.114 * box_color[2]
|
|
# text_color = (0,0,0) if luminance > 160 else (255,255,255)
|
|
# cv2.putText(
|
|
# img=scene,
|
|
# text=text,
|
|
# org=(text_x, text_y),
|
|
# fontFace=font,
|
|
# fontScale=self.text_scale,
|
|
# # color=self.text_color.as_rgb(),
|
|
# color=text_color,
|
|
# thickness=self.text_thickness,
|
|
# lineType=cv2.LINE_AA,
|
|
# )
|
|
# return scene
|
|
|
|
|
|
# def box_area(box):
|
|
# return (box[2] - box[0]) * (box[3] - box[1])
|
|
|
|
# def intersection_area(box1, box2):
|
|
# x1 = max(box1[0], box2[0])
|
|
# y1 = max(box1[1], box2[1])
|
|
# x2 = min(box1[2], box2[2])
|
|
# y2 = min(box1[3], box2[3])
|
|
# return max(0, x2 - x1) * max(0, y2 - y1)
|
|
|
|
# def IoU(box1, box2, return_max=True):
|
|
# intersection = intersection_area(box1, box2)
|
|
# union = box_area(box1) + box_area(box2) - intersection
|
|
# if box_area(box1) > 0 and box_area(box2) > 0:
|
|
# ratio1 = intersection / box_area(box1)
|
|
# ratio2 = intersection / box_area(box2)
|
|
# else:
|
|
# ratio1, ratio2 = 0, 0
|
|
# if return_max:
|
|
# return max(intersection / union, ratio1, ratio2)
|
|
# else:
|
|
# return intersection / union
|
|
|
|
|
|
# def get_optimal_label_pos(text_padding, text_width, text_height, x1, y1, x2, y2, detections, image_size):
|
|
# """ check overlap of text and background detection box, and get_optimal_label_pos,
|
|
# pos: str, position of the text, must be one of 'top left', 'top right', 'outer left', 'outer right' TODO: if all are overlapping, return the last one, i.e. outer right
|
|
# Threshold: default to 0.3
|
|
# """
|
|
|
|
# def get_is_overlap(detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size):
|
|
# is_overlap = False
|
|
# for i in range(len(detections)):
|
|
# detection = detections.xyxy[i].astype(int)
|
|
# if IoU([text_background_x1, text_background_y1, text_background_x2, text_background_y2], detection) > 0.3:
|
|
# is_overlap = True
|
|
# break
|
|
# # check if the text is out of the image
|
|
# if text_background_x1 < 0 or text_background_x2 > image_size[0] or text_background_y1 < 0 or text_background_y2 > image_size[1]:
|
|
# is_overlap = True
|
|
# return is_overlap
|
|
|
|
# # if pos == 'top left':
|
|
# text_x = x1 + text_padding
|
|
# text_y = y1 - text_padding
|
|
|
|
# text_background_x1 = x1
|
|
# text_background_y1 = y1 - 2 * text_padding - text_height
|
|
|
|
# text_background_x2 = x1 + 2 * text_padding + text_width
|
|
# text_background_y2 = y1
|
|
# is_overlap = get_is_overlap(detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size)
|
|
# if not is_overlap:
|
|
# return text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2
|
|
|
|
# # elif pos == 'outer left':
|
|
# text_x = x1 - text_padding - text_width
|
|
# text_y = y1 + text_padding + text_height
|
|
|
|
# text_background_x1 = x1 - 2 * text_padding - text_width
|
|
# text_background_y1 = y1
|
|
|
|
# text_background_x2 = x1
|
|
# text_background_y2 = y1 + 2 * text_padding + text_height
|
|
# is_overlap = get_is_overlap(detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size)
|
|
# if not is_overlap:
|
|
# return text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2
|
|
|
|
|
|
# # elif pos == 'outer right':
|
|
# text_x = x2 + text_padding
|
|
# text_y = y1 + text_padding + text_height
|
|
|
|
# text_background_x1 = x2
|
|
# text_background_y1 = y1
|
|
|
|
# text_background_x2 = x2 + 2 * text_padding + text_width
|
|
# text_background_y2 = y1 + 2 * text_padding + text_height
|
|
|
|
# is_overlap = get_is_overlap(detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size)
|
|
# if not is_overlap:
|
|
# return text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2
|
|
|
|
# # elif pos == 'top right':
|
|
# text_x = x2 - text_padding - text_width
|
|
# text_y = y1 - text_padding
|
|
|
|
# text_background_x1 = x2 - 2 * text_padding - text_width
|
|
# text_background_y1 = y1 - 2 * text_padding - text_height
|
|
|
|
# text_background_x2 = x2
|
|
# text_background_y2 = y1
|
|
|
|
# is_overlap = get_is_overlap(detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size)
|
|
# if not is_overlap:
|
|
# return text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2
|
|
|
|
# return text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2
|