Draw clusters on an image
(
image: Image.Image, clusters: list[Cluster], scale_x: float, scale_y: float
)
| 6 | |
| 7 | |
| 8 | def draw_clusters( |
| 9 | image: Image.Image, clusters: list[Cluster], scale_x: float, scale_y: float |
| 10 | ) -> None: |
| 11 | """ |
| 12 | Draw clusters on an image |
| 13 | """ |
| 14 | draw = ImageDraw.Draw(image, "RGBA") |
| 15 | # Create a smaller font for the labels |
| 16 | font: ImageFont.ImageFont | FreeTypeFont |
| 17 | try: |
| 18 | font = ImageFont.truetype("arial.ttf", 12) |
| 19 | except OSError: |
| 20 | # Fallback to default font if arial is not available |
| 21 | font = ImageFont.load_default() |
| 22 | for c_tl in clusters: |
| 23 | all_clusters = [c_tl, *c_tl.children] |
| 24 | for c in all_clusters: |
| 25 | # Draw cells first (underneath) |
| 26 | cell_color = (0, 0, 0, 40) # Transparent black for cells |
| 27 | for tc in c.cells: |
| 28 | cx0, cy0, cx1, cy1 = tc.bbox.as_tuple() |
| 29 | cx0 *= scale_x |
| 30 | cx1 *= scale_x |
| 31 | cy0 *= scale_x |
| 32 | cy1 *= scale_y |
| 33 | |
| 34 | draw.rectangle( |
| 35 | [(cx0, cy0), (cx1, cy1)], |
| 36 | outline=None, |
| 37 | fill=cell_color, |
| 38 | ) |
| 39 | # Draw cluster rectangle |
| 40 | x0, y0, x1, y1 = c.bbox.as_tuple() |
| 41 | x0 *= scale_x |
| 42 | x1 *= scale_x |
| 43 | y0 *= scale_x |
| 44 | y1 *= scale_y |
| 45 | |
| 46 | cluster_fill_color = (*list(DocItemLabel.get_color(c.label)), 70) |
| 47 | cluster_outline_color = ( |
| 48 | *list(DocItemLabel.get_color(c.label)), |
| 49 | 255, |
| 50 | ) |
| 51 | draw.rectangle( |
| 52 | [(x0, y0), (x1, y1)], |
| 53 | outline=cluster_outline_color, |
| 54 | fill=cluster_fill_color, |
| 55 | ) |
| 56 | # Add label name and confidence |
| 57 | label_text = f"{c.label.name} ({c.confidence:.2f})" |
| 58 | # Create semi-transparent background for text |
| 59 | text_bbox = draw.textbbox((x0, y0), label_text, font=font) |
| 60 | text_bg_padding = 2 |
| 61 | draw.rectangle( |
| 62 | [ |
| 63 | ( |
| 64 | text_bbox[0] - text_bg_padding, |
| 65 | text_bbox[1] - text_bg_padding, |
no outgoing calls
no test coverage detected