diff --git a/supervision/__init__.py b/supervision/__init__.py index 31d01314..6ddc5c57 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -17,6 +17,7 @@ from supervision.annotators.core import ( LabelAnnotator, MaskAnnotator, TraceAnnotator, + ClassificationAnnotator ) from supervision.annotators.utils import ColorLookup from supervision.classification.core import Classifications diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 8005f49b..1415d15a 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -7,6 +7,7 @@ import numpy as np from supervision.annotators.base import BaseAnnotator from supervision.annotators.utils import ColorLookup, Trace, resolve_color from supervision.detection.core import Detections +from supervision.classification.core import Classifications from supervision.draw.color import Color, ColorPalette from supervision.geometry.core import Position @@ -919,3 +920,127 @@ class TraceAnnotator: thickness=self.thickness, ) return scene + + +class ClassificationAnnotator: + """ + Annotate classification results on an image. + """ + + def __init__( + self, + color: Union[Color, ColorPalette] = ColorPalette.default(), + text_color: Color = Color.black(), + text_scale: float = 0.5, + text_thickness: int = 1, + text_padding: int = 10, + text_position: Position = Position.TOP_LEFT, + color_lookup: ColorLookup = ColorLookup.CLASS, + ): + """ + Args: + color (Union[Color, ColorPalette]): The color or color palette to use for + annotating the text background. + text_color (Color): The color to use for the text. + text_scale (float): Font scale for the text. + text_thickness (int): Thickness of the text characters. + text_position (Position): Position of the text relative to the image. + Possible values are defined in the `Position` enum. + color_lookup (str): Strategy for mapping colors to annotations. + Options are `INDEX`, `CLASS`, `TRACE`. + """ + self.color: Union[Color, ColorPalette] = color + self.text_color: Color = text_color + self.text_scale: float = text_scale + self.text_thickness: int = text_thickness + self.text_position: Position = text_position + self.color_lookup: ColorLookup = color_lookup + self.text_padding: int = text_padding + + def annotate( + self, + scene: np.ndarray, + classifications: Classifications, + labels: List[str] = None, + custom_color_lookup: Optional[np.ndarray] = None, + ) -> np.ndarray: + """ + Annotates the given scene with labels based on the provided detections. + + Args: + scene (np.ndarray): The image where labels will be drawn. + detections (Detections): Object detections to annotate. + labels (List[str]): Optional. Custom labels for each detection. + custom_color_lookup (Optional[np.ndarray]): Custom color lookup array. + Allows to override the default color mapping strategy. + + Returns: + np.ndarray: The annotated image. + + Example: + ```python + >>> import supervision as sv + + >>> image = ... + >>> detections = sv.Detections(...) + + >>> classification_annotator = sv.ClassificationAnnotator(text_position=sv.Position.CENTER) + >>> annotated_frame = classification_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` + """ + font = cv2.FONT_HERSHEY_SIMPLEX + + classification = classifications.get_top_k(k=1) + classification_idx = classification[0][0] + + color = resolve_color( + color=self.color, + detections=classifications, + detection_idx=classification_idx, + color_lookup=self.color_lookup + if custom_color_lookup is None + else custom_color_lookup, + ) + + self.text_color = color + + text = f"{labels[classification_idx]} ({classifications.confidence[classification_idx] * 100:.2f}%)" + + text_wh = cv2.getTextSize( + text=text, + fontFace=font, + fontScale=self.text_scale, + thickness=self.text_thickness, + )[0] + + if self.text_position == Position.BOTTOM_LEFT: + text_x = self.text_padding + text_y = scene.shape[0] - self.text_padding + elif self.text_position == Position.BOTTOM_RIGHT: + text_x = scene.shape[1] - text_wh[0] - self.text_padding + text_y = scene.shape[0] - self.text_padding + elif self.text_position == Position.TOP_LEFT: + text_x = self.text_padding + text_y = text_wh[1] + self.text_padding + elif self.text_position == Position.TOP_RIGHT: + text_x = scene.shape[1] - text_wh[0] - self.text_padding + text_y = text_wh[1] + self.text_padding + else: + raise ValueError( + f"Invalid position {self.text_position} for classification annotator." + ) + + cv2.putText( + img=scene, + text=text, + org=(text_x, text_y), + fontFace=font, + fontScale=self.text_scale, + color=self.text_color.as_rgb(), + thickness=self.text_thickness, + lineType=cv2.LINE_AA, + ) + return scene diff --git a/supervision/classification/core.py b/supervision/classification/core.py index 157275dd..c62d758a 100644 --- a/supervision/classification/core.py +++ b/supervision/classification/core.py @@ -39,6 +39,43 @@ class Classifications: _validate_class_ids(self.class_id, n) _validate_confidence(self.confidence, n) + def __len__(self) -> int: + """ + Returns the number of classifications. + """ + return len(self.class_id) + + @classmethod + def from_clip(cls, clip_results) -> Classifications: + """ + Creates a Classifications instance from a + (https://github.com/openai/clip) inference result. + + Args: + clip_results (np.ndarray): The output result from clip model + + Returns: + Classifications: A new Classifications object. + + Example: + ```python + >>> import cv2 + >>> import clip + >>> import supervision as sv + + >>> image = cv2.imread(SOURCE_IMAGE_PATH) + >>> model, preprocess = clip.load('ViT-B/32') + >>> text = clip.tokenize(["a diagram", "a dog", "a cat"]) + >>> classifications = sv.Classifications.from_clip(model(image, text)) + ``` + """ + + probs = clip_results.softmax(dim=-1).cpu().numpy() + class_ids = np.arange(probs.shape[1]) + confidence = probs[0] + + return cls(class_id=class_ids, confidence=confidence) + @classmethod def from_ultralytics(cls, ultralytics_results) -> Classifications: """