Coverage for drawyolo/draw.py: 100.00%
113 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-05-21 06:37 +0000
1import cv2
2import random
3import numpy as np
4from pathlib import Path
6FONT = cv2.FONT_HERSHEY_SIMPLEX
8def get_font_scale(height) -> float:
9 target_height = min(max(0.03 * height, 10), 100) # Limit to 10-100 pixels
10 text = "Hg"
11 low, high = 0.01, 10 # Start with a wide range
12 best_fontScale = low
14 while high - low > 0.01: # Precision threshold
15 mid = (low + high) / 2
16 (_, text_height), _ = cv2.getTextSize(text, FONT, mid, thickness=1)
18 if text_height < target_height:
19 best_fontScale = mid # Update best found value
20 low = mid # Increase font scale
21 else:
22 high = mid # Decrease font scale
24 return best_fontScale
27def rescale_image(
28 image: np.ndarray, width: int | None = None, height: int | None = None
29) -> np.ndarray:
30 """Rescale an image based on given width and height.
32 - If both `width` and `height` are None or 0, return the original image.
33 - If only one is provided, rescale while maintaining the aspect ratio.
34 - If both are provided, rescale to the exact size (ignoring aspect ratio).
35 """
36 if not width and not height:
37 return image
39 original_height, original_width = image.shape[:2]
40 if not width:
41 scale = height / original_height
42 width = int(original_width * scale)
43 elif not height:
44 scale = width / original_width
45 height = int(original_height * scale)
47 return cv2.resize(image, (width, height), interpolation=cv2.INTER_AREA)
50def save_image(
51 image: np.ndarray,
52 output: str | Path | None,
53):
54 """Save image to file"""
55 if output:
56 output = Path(output)
57 output.parent.mkdir(parents=True, exist_ok=True)
58 cv2.imwrite(str(output), image)
61def default_line_thickness(width: int, height: int) -> int:
62 return max(round(0.002 * min(height,width)),2)
65def plot_one_box(
66 x, image, color=None, label: str = None,
67 line_thickness: int | None = None,
68 font_scale:float|None = None,
69):
70 """Plots one bounding box on image img"""
71 line_thickness = line_thickness or default_line_thickness(
72 image.shape[1], image.shape[0]
73 )
74 color = color or [random.randint(0, 255) for _ in range(3)]
75 c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
76 cv2.rectangle(image, c1, c2, color, thickness=line_thickness, lineType=cv2.LINE_AA)
77 font_scale = font_scale or get_font_scale(image.shape[0])
78 if label:
79 label = str(label).upper()
80 t_size = cv2.getTextSize(label, FONT, fontScale=font_scale, thickness=1)[0]
81 font_thickness = max(round(t_size[1] * 0.08), 1)
82 c2 = c1[0] + t_size[0], c1[1] - round(t_size[1] * 1.1) - font_thickness * 2
83 cv2.rectangle(image, c1, c2, color, -1, cv2.LINE_AA) # filled
84 cv2.putText(
85 image,
86 label,
87 (c1[0], c1[1] - round(0.05 * t_size[1]) - font_thickness),
88 FONT,
89 font_scale,
90 [225, 255, 255],
91 thickness=font_thickness,
92 lineType=cv2.LINE_AA,
93 )
96def make_colors(count: int):
97 """Generate random colors"""
98 random.seed(42)
99 colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(count)]
100 return colors
103def draw_box_on_image_with_labels(
104 image: Path,
105 labels: Path,
106 output: Path,
107 classes: list[str],
108 colors: list[str] = None,
109 width: int | None = None,
110 height: int | None = None,
111 line_thickness: int | None = None,
112) -> np.ndarray:
113 """
114 Adds rectangle boxes on the images.
115 """
116 colors = colors or make_colors(len(classes))
118 # Read image
119 try:
120 image = cv2.imread(str(image))
121 image = rescale_image(image, width, height)
122 height, width = image.shape[:2]
123 except Exception as e:
124 print(f"Cannot read image: {e}")
125 return
127 line_thickness = line_thickness or default_line_thickness(width, height)
128 font_scale = get_font_scale(height)
130 # Get Labels
131 labels = Path(labels)
132 labels = labels.read_text().strip().split("\n") if labels.exists() else []
133 for line in labels:
134 staff = line.split()
135 class_idx = int(staff[0])
137 x_center, y_center, w, h = (
138 float(staff[1]) * width,
139 float(staff[2]) * height,
140 float(staff[3]) * width,
141 float(staff[4]) * height,
142 )
143 x1 = round(x_center - w / 2)
144 y1 = round(y_center - h / 2)
145 x2 = round(x_center + w / 2)
146 y2 = round(y_center + h / 2)
148 plot_one_box(
149 [x1, y1, x2, y2],
150 image,
151 color=colors[class_idx],
152 label=classes[class_idx],
153 line_thickness=line_thickness,
154 font_scale=font_scale,
155 )
157 save_image(image, output)
158 return image
161def draw_box_on_image_with_yolo_result(
162 image: Path,
163 results,
164 output: Path,
165 classes: list[str] = None,
166 colors: list[str] = None,
167 highest: bool = False,
168 width: int | None = None,
169 height: int | None = None,
170 line_thickness: int | None = None,
171) -> np.ndarray:
172 colors = colors or make_colors(len(classes))
174 image = cv2.imread(str(image))
176 boxes_list = results.boxes
177 if highest:
178 best_confidence = 0
179 best_boxes = None
180 for boxes in boxes_list:
181 confidence = boxes.conf.item()
182 if confidence > best_confidence:
183 best_boxes = boxes
184 best_confidence = confidence
185 boxes_list = [best_boxes]
187 original_height, original_width = image.shape[:2]
188 image = rescale_image(image, width, height)
189 height, width = image.shape[:2]
191 font_scale = get_font_scale(height)
192 line_thickness = line_thickness or default_line_thickness(width, height)
194 for boxes in boxes_list:
195 class_idx = boxes.cls[0].int().item()
196 xyxy = boxes.xyxy.cpu()[0].clone().detach().numpy()
197 xyxy[0] = xyxy[0] * width / original_width
198 xyxy[1] = xyxy[1] * height / original_height
199 xyxy[2] = xyxy[2] * width / original_width
200 xyxy[3] = xyxy[3] * height / original_height
201 plot_one_box(
202 xyxy,
203 image,
204 color=colors[class_idx],
205 label=classes[class_idx],
206 line_thickness=line_thickness,
207 font_scale=font_scale,
208 )
210 save_image(image, output)
211 return image
214def draw_box_on_image_with_model(
215 image: Path,
216 weights: Path | str,
217 output: Path,
218 res: int = 1280,
219 classes: list[str] = None,
220 colors: list[str] = None,
221 highest: bool = False,
222 width: int | None = None,
223 height: int | None = None,
224 line_thickness: int | None = None,
225) -> np.ndarray:
226 """Draw boxes on image using YOLO model"""
227 from ultralytics import YOLO
229 model = YOLO(str(weights))
231 res = res or 1280
232 classes = classes or model.names
234 results = model.predict(source=[image], show=False, save=False, imgsz=res)
235 assert len(results) == 1
237 return draw_box_on_image_with_yolo_result(
238 image,
239 results[0],
240 output=output,
241 classes=classes,
242 colors=colors,
243 highest=highest,
244 width=width,
245 height=height,
246 line_thickness=line_thickness,
247 )