Spaces:
Running
Running
| import os | |
| import cv2 | |
| import numpy as np | |
| import supervision as sv | |
| from PIL import Image, ImageDraw, ImageFont | |
| from supervision.annotators.utils import resolve_color | |
| # visualization toos based on supervision | |
| BOUNDING_BOX_ANNOTATOR = sv.BoundingBoxAnnotator(thickness=2) | |
| class LabelAnnotator(sv.LabelAnnotator): | |
| def resolve_text_background_xyxy( | |
| center_coordinates, | |
| text_wh, | |
| position, | |
| ): | |
| center_x, center_y = center_coordinates | |
| text_w, text_h = text_wh | |
| return center_x, center_y, center_x + text_w, center_y + text_h | |
| def _draw_labels( | |
| self, | |
| scene: np.ndarray, | |
| labels: list[str], | |
| label_properties: np.ndarray, | |
| detections, | |
| custom_color_lookup, | |
| ) -> None: | |
| assert len(labels) == len(label_properties) == len(detections), ( | |
| f"Number of label properties ({len(label_properties)}), " | |
| f"labels ({len(labels)}) and detections ({len(detections)}) " | |
| "do not match." | |
| ) | |
| color_lookup = ( | |
| custom_color_lookup | |
| if custom_color_lookup is not None | |
| else self.color_lookup | |
| ) | |
| font = ImageFont.truetype("simhei.ttf", int(30 * self.text_scale)) | |
| for idx, label_property in enumerate(label_properties): | |
| background_color = resolve_color( | |
| color=self.color, | |
| detections=detections, | |
| detection_idx=idx, | |
| color_lookup=color_lookup, | |
| ) | |
| text_color = resolve_color( | |
| color=self.text_color, | |
| detections=detections, | |
| detection_idx=idx, | |
| color_lookup=color_lookup, | |
| ) | |
| box_xyxy = label_property[:4] | |
| text_height_padded = label_property[4] | |
| self.draw_rounded_rectangle( | |
| scene=scene, | |
| xyxy=box_xyxy, | |
| color=background_color.as_bgr(), | |
| border_radius=self.border_radius, | |
| ) | |
| text_x = box_xyxy[0] + self.text_padding | |
| text_y = box_xyxy[1] | |
| scene_pil = Image.fromarray(cv2.cvtColor(scene, cv2.COLOR_BGR2RGB)) | |
| draw = ImageDraw.Draw(scene_pil) | |
| draw.text( | |
| (text_x, text_y), | |
| labels[idx], | |
| font=font, | |
| fill=(text_color.r, text_color.g, text_color.b), | |
| ) | |
| scene[:] = cv2.cvtColor(np.array(scene_pil), cv2.COLOR_RGB2BGR) | |
| LABEL_ANNOTATOR = LabelAnnotator(text_padding=4, | |
| text_scale=0.5, | |
| text_thickness=1) | |
| POINT_ANNOTATOR = sv.DotAnnotator(radius=6) | |
| def draw_boxes_points_with_labels( | |
| cv2_image, | |
| boxes=None, | |
| points=None, | |
| classes=None, | |
| output_path=None, | |
| ): | |
| annotated_image = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB) | |
| if boxes is not None and boxes.size: | |
| detections = sv.Detections( | |
| xyxy=boxes, | |
| class_id=np.arange(len(boxes)), | |
| confidence=np.ones(len(boxes)) | |
| ) | |
| annotated_image = BOUNDING_BOX_ANNOTATOR.annotate( | |
| annotated_image, detections) | |
| if points is not None and points.size: | |
| points = np.concatenate([points, points], axis=1) | |
| detections = sv.Detections( | |
| xyxy=points, | |
| class_id=np.arange(len(points)), | |
| confidence=np.ones(len(points)) | |
| ) | |
| annotated_image = POINT_ANNOTATOR.annotate( | |
| annotated_image, detections, | |
| ) | |
| if classes: | |
| annotated_image = LABEL_ANNOTATOR.annotate( | |
| annotated_image, detections, labels=classes | |
| ) | |
| if output_path: | |
| cv2.imwrite( | |
| output_path, | |
| cv2.cvtColor(annotated_image, cv2.COLOR_RGB2BGR) | |
| ) | |
| return annotated_image |