Spaces:
Running
Running
File size: 6,297 Bytes
6a0e448 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 |
import math
from typing import Any
import cv2
import numpy as np
from pydantic import BaseModel, Field, field_validator
class Point(BaseModel):
x: int
y: int
def __iter__(self):
return iter((self.x, self.y))
def __getitem__(self, index) -> int:
return (self.x, self.y)[index]
def __tuple__(self) -> tuple[int, int]:
return (self.x, self.y)
def __repr__(self) -> str:
return f"Point(x={self.x}, y={self.y})"
class BoundingBox(BaseModel):
label: str = Field(..., description="The label that's given for this bounding box")
left: int = Field(..., description="Left coordinate of the bounding box")
right: int = Field(..., description="Right coordinate of the bounding box")
top: int = Field(..., description="Top coordinate of the bounding box")
bottom: int = Field(..., description="Bottom coordinate of the bounding box")
@field_validator("left", "top", mode="before")
@classmethod
def round_down(cls, v):
return math.floor(float(v))
@field_validator("right", "bottom", mode="before")
@classmethod
def round_up(cls, v):
return math.ceil(float(v))
class POI(BaseModel):
info: dict[str, Any]
element_centroid: Point
bounding_box: BoundingBox
def calculate_dash_points(start, end, dash_length, gap_length):
x1, y1 = start
x2, y2 = end
dx = x2 - x1
dy = y2 - y1
dist = np.sqrt(dx * dx + dy * dy)
if dist == 0:
return []
unit_x = dx / dist
unit_y = dy / dist
dash_points = []
current_dist = 0
while current_dist < dist:
dash_end = min(current_dist + dash_length, dist)
dash_points.extend(
[
(int(x1 + unit_x * current_dist), int(y1 + unit_y * current_dist)),
(int(x1 + unit_x * dash_end), int(y1 + unit_y * dash_end)),
],
)
current_dist += dash_length + gap_length
return dash_points
def draw_dashed_rectangle(
img,
bbox: BoundingBox,
color,
thickness=1,
dash_length=10,
gap_length=5,
):
# Calculate dash points for all sides
top_points = calculate_dash_points(
(bbox.left + 25, bbox.top + 25),
(bbox.right + 25, bbox.top + 25),
dash_length,
gap_length,
)
right_points = calculate_dash_points(
(bbox.right + 25, bbox.top + 25),
(bbox.right + 25, bbox.bottom + 25),
dash_length,
gap_length,
)
bottom_points = calculate_dash_points(
(bbox.right + 25, bbox.bottom + 25),
(bbox.left + 25, bbox.bottom + 25),
dash_length,
gap_length,
)
left_points = calculate_dash_points(
(bbox.left + 25, bbox.bottom + 25),
(bbox.left + 25, bbox.top + 25),
dash_length,
gap_length,
)
# Combine all points
all_points = top_points + right_points + bottom_points + left_points
# Draw all lines at once
if all_points:
all_points = np.array(all_points).reshape((-1, 2, 2))
cv2.polylines(img, all_points, False, color, thickness)
# @time_it(name='Annotate bounding box')
def annotate_bounding_box(image: bytes, bbox: BoundingBox) -> None:
# Draw dashed bounding box
draw_dashed_rectangle(
image,
bbox,
color=(0, 0, 255),
thickness=1,
dash_length=10,
gap_length=5,
)
# Prepare label
font_scale = 0.4 * 4 # Increased by 4x for the larger patch
font = cv2.FONT_HERSHEY_SIMPLEX
thickness = 3 # Increased thickness for the larger patch
# Get text size for the larger patch
(label_width, label_height), _ = cv2.getTextSize(
bbox.label,
font,
font_scale,
thickness,
)
# Create a larger patch (4x)
large_label_patch = np.zeros(
(label_height + 20, label_width + 20, 4),
dtype=np.uint8,
)
large_label_patch[:, :, 0:3] = (0, 0, 255) # BGR color format: Red background
large_label_patch[:, :, 3] = 128 # Alpha channel: 50% opacity (128/255 = 0.5)
# Draw text on the larger patch
cv2.putText(
large_label_patch,
bbox.label,
(8, label_height + 8), # Adjusted position for the larger patch
font,
font_scale,
(255, 255, 255, 128), # White text, 50% opaque (128/255 = 0.5)
thickness,
)
# Scale down the patch to improve anti-aliasing
label_patch = cv2.resize(
large_label_patch,
(label_width // 4 + 5, label_height // 4 + 5),
interpolation=cv2.INTER_AREA,
)
# Calculate position for top-left alignment
offset = 2 # Small offset to prevent touching the bounding box edge
x = min(image.shape[1], max(0, int(bbox.left + 25) - offset))
y = min(image.shape[0], max(0, int(bbox.top + 25) - label_patch.shape[0] - offset))
# Ensure we're not out of bounds
x_end = min(image.shape[1], x + label_patch.shape[1])
y_end = min(image.shape[0], y + label_patch.shape[0])
label_patch = label_patch[: (y_end - y), : (x_end - x)]
# Create a mask for the label patch
alpha_mask = label_patch[:, :, 3] / 255.0
alpha_mask = np.repeat(alpha_mask[:, :, np.newaxis], 3, axis=2)
# Blend the label patch with the image
image_section = image[y:y_end, x:x_end]
blended = (1 - alpha_mask) * image_section + alpha_mask * label_patch[:, :, 0:3]
image[y:y_end, x:x_end] = blended.astype(np.uint8)
def annotate_bounding_boxes(image: bytes, bounding_boxes: list[BoundingBox]) -> bytes:
# Read the image
nparr = np.frombuffer(image, np.uint8)
# Decode the image
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
padded_img = cv2.copyMakeBorder(
img,
top=25, # Value chosen based on label size
bottom=25, # Value chosen based on label size
left=25, # Value chosen based on label size
right=25, # Value chosen based on label size
borderType=cv2.BORDER_CONSTANT,
value=(255, 255, 255),
)
for bounding_box in bounding_boxes:
# Annotate the image in place with the bounding box and the bounding box label
annotate_bounding_box(padded_img, bounding_box)
_, buffer = cv2.imencode(".jpeg", padded_img)
return buffer.tobytes()
|