Python script, ASCII text executable
1
import os
2
import time
3
from os import path
4
import cv2 as cv
6
import numpy as np
7
from PIL import Image, ImageFont, ImageDraw
8
import matplotlib.pyplot as plt
10
import matplotlib.patches as patches
11
import math
13
font = ImageFont.truetype("/home/vlad/.local/share/fonts/Roboto-Medium.ttf", 12)
15
image_path = "./yolo_data/IMG-20240514-WA0087.jpg"
17
assert path.exists(image_path), f"Image {image_path} not found"
18
from ruamel.yaml import YAML
19
yaml = YAML()
21
with open("yolo.yaml", "r") as file:
23
data = yaml.load(file)
24
class_mapping = data["names"]
25
img = cv.imread(image_path)
27
annotations_file = path.splitext(image_path)[0] + ".txt"
28
with open(annotations_file, "r") as f:
29
annotations = []
30
for line in f:
31
klass, cx, cy, w, h = line.split()
32
annotations.append((int(klass), float(cx), float(cy), float(w), float(h)))
33
for box in annotations:
35
# Get coordinates and draw the bounding box
36
klass, cx, cy, w, h = box
37
x1 = int((cx - w / 2) * img.shape[1])
38
y1 = int((cy - h / 2) * img.shape[0])
39
x2 = int((cx + w / 2) * img.shape[1])
40
y2 = int((cy + h / 2) * img.shape[0])
41
cv.rectangle(img, (x1, y1), (x2, y2), (255, 255, 255), 2)
43
cv.rectangle(img, (x1 - 1, y1 - 1), (x2 + 1, y2 + 1), (0, 152, 255), 2)
44
cls = int(klass)
46
img_pil = Image.fromarray(img)
48
draw = ImageDraw.Draw(img_pil)
49
draw.text((x1, y1),
51
f"[VAL] {class_mapping[cls]}",
52
font=font,
53
fill=(255, 255, 255),
54
stroke_width=1,
55
stroke_fill=(0, 0, 0))
56
img = np.array(img_pil)
57
result = cv.imwrite("./annotated.jpg", img)
59
assert result, "Failed to save image"
60
assert path.exists("./annotated.jpg"), "Failed to save image"
61