mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 22:36:36 +00:00
3fd726f9c7
add train some catches
32 lines
1.1 KiB
Python
32 lines
1.1 KiB
Python
import cv2
|
|
from matplotlib import pyplot as plt
|
|
|
|
|
|
class ImageLabel:
|
|
def __init__(self, image_path, image, labels_path, labels):
|
|
self.image_path = image_path
|
|
self.image = image
|
|
self.labels_path = labels_path
|
|
self.labels = labels
|
|
|
|
def visualize(self, annotation_classes):
|
|
img = cv2.cvtColor(self.image.copy(), cv2.COLOR_BGR2RGB)
|
|
height, width, channels = img.shape
|
|
for label in self.labels:
|
|
class_num = int(label[-1])
|
|
x_c = float(label[0])
|
|
y_c = float(label[1])
|
|
w = float(label[2])
|
|
h = float(label[3])
|
|
x_min = x_c - w / 2
|
|
y_min = y_c - h / 2
|
|
x_max = x_min + w
|
|
y_max = y_min + h
|
|
color = annotation_classes[class_num].color_tuple
|
|
|
|
cv2.rectangle(img, (int(x_min * width), int(y_min * height)), (int(x_max * width), int(y_max * height)),
|
|
color=color, thickness=3)
|
|
plt.figure(figsize=(12, 12))
|
|
plt.axis('off')
|
|
plt.imshow(img)
|
|
plt.show() |