Table of contents
foreword
The tool is written by opencv, which can directly generate the labels (pose and regular labels) required by YOLO
- The code is placed at the end of the article, as well as the Baidu cloud download link
- First put a video showing the actual operation
yolov5 dataset annotation, yolo-pose dataset annotation
under normal label
- Press Q to switch to the next image
- Press T to exit directly
- Press Y to delete the current picture and the corresponding label
- Press R to hide the current content, and it will be merged with the previous content after continuing to mark
- Double-click the point on the box to switch to the delete state, press E to do nothing, press W to delete
- Left mouse button to stretch and move the current frame
Mark the key points
- Press Q to switch to the next image
- Press T to exit
- Press Y to delete the current picture and the corresponding label
- Press R to hide the current content, and it will be merged with the previous content after continuing to mark
- Double-click the key to toggle the visibility of the key (0 1 2)
- keys can be dragged
- Double-click the blank area, enter 0 and 8 -> 0 10+1 8=8, move the key point 8 to the double-click
Annotate the dataset in YOLO format
- First, you need to change the initial value of the code on line 22 from True to None, which is a regular label
- Replace the path where the image is located and the save path, and nothing can be placed in the save path
Label key points
- -Change the initial value of the code on line 22 to True to indicate key points, and set the corresponding key_point_num (there are several key points to be marked)
The difference from ordinary annotations is that double-clicking the rectangle here will not cause any operation, which means that you can't delete the rectangle, you can only move it, and you can't delete the key point, you can only move it, but the key point can also be double-clicked besides dragging and keyboard input moves
- You can double-click the key point with the left mouse button
Because the key point has three states of 0, 1, and 2, you can see the state
and then enter 0 or 1 or 2 from the keyboard
- Setting it to 1 means that the flag bit changes from 2->1, and the display color changes. Set it to 2 because the previous state is 2, so it seems that there is no change. Setting it to 0, in addition to changing the state to 0, also changes the corresponding Coordinates (x, y) are set to 0 0
- If you just want to set the state to 0 and keep the coordinates unchanged, you need to comment out the code in lines 481 and 482
elif key_insert is not None and key == ord('0'):
with open(label_path, 'r') as f:
label_temp = f.read()
str_temp = label_temp.split(' ')
str_temp[3 * int(key_insert) + 7] = '0'
#str_temp[3 * int(key_insert) + 7 - 1] = '0'
#str_temp[3 * int(key_insert) + 7 - 2] = '0'
If you want to change the point whose status is 0 back, you need to put the mouse on the upper left corner, double-click, or drag it to other conspicuous areas first
Next, select the first three points to display the starting label and the changed label
- initial label
0 0.477661 0.495117 0.895264 0.702474 0.69458 0.65625 2.0 0.895019 0.810872 2.0 0.572509 0.643229 2.0 0.745117 0.577473 2.0 0.642822 0.57194 2.0 0.74707 0.429361 2.0 0.641601 0.431315 2.0 0.748535 0.267252 2.0 0.64624 0.260742 2.0 0.279296 0.255859 2.0 0.050292 0.253906 2.0
The flag bits of the first three key points are all 2
- Run it directly, you can find that a box and key_point_num key points have been generated, and the corresponding labels are generated locally at the same time
- This is because if you do not pre-generate a box and key_point_num key points, it will be quite complicated to deal with, and it is not easy to handle when the number of key points exceeds 10, so pre-generate key points and boxes, the state of the key point is 2, you can Modify by yourself
- When labeling, drag the box and point to the corresponding position, and set the state of the key point to the desired one.
- Drag the position of key points and boxes, and you can also use another labeling method for key points. Double- click
the area of 9 points where there are no key points and boxes.
11, enter twice from the keyboard, you must enter the number of 0-9, if you want to select a key point such as 0-9, enter 0 for the first time, and enter the key point index for the second time
For example, to operate key point 8, enter 0 for the first time and enter 8 for the second time, and you can see that the key point has reached the double-click position of the mouse.
When the drawing is marked, click Q to switch. After the completion, you can see that the corresponding key point labels are generated locally. Press T to always exit the
mark. You can also stretch the window without affecting the mark results. In short, be careful not to cross Close the window or force exit, try to press Q to exit
Annotate multiple boxes and key points in one figure
Although a box and n key points are pre-generated, if there are multiple people in a picture, such as two people, then both people need to be marked in the image. In order to deal with this situation
First mark normally
and then do not press Q to exit at this time, if you want to mark the second person, press and hold R at this time
Then continue labeling.
If you still want to continue labeling, continue to hold down W until you have completed everything and press Q to switch to the next picture. When everything is completed, you can see the label file in the local , two marked objects have been successfully added
At this time, we re-run the program
and we can see that the marked objects have been successfully marked. Click, move, etc., you can only watch. If you operate, only the first line in the label will be kept, and the rest will be deleted.
final code
import math
import os
import cv2
"""
标注关键点只能存在一个框和多个点，并且不能删除点和删除框，读取本地文件的关键点要保证其中的关键点
数和key_point_num的值是一样的，本地标签中如果只存在框的信息就不要使用该脚本标注，不然会出错，
本地文件夹中可以有标签，如果有会优先加载本地标签，没有才会创建一个
"""
draw_line_circle = True # True/None 是否在框上绘制点（8个点）
key_point_is = True # 是否标记关键点 设置为None标注普通yolo标签
# 可以自定义得参数
image_path = R'C:\Users\lengdan\Desktop\YN' # 要标注的图像所在文件夹
label_path = R'C:\Users\lengdan\Desktop\YNlabel' # 标注完成保存到的文件夹
circle_distance = 10 # 半径范围:鼠标进入点的半径范围内会出现光圈
key_point_num = 2 # 关键点个数
box_thickness = 1 # 框的粗细
small_box_thickness = 1 # 框的8个点的粗细
label_thickness = 1 # 框上面的类别字体的粗细
label_fontScale = 0.4 # 框上面的类别字体的倍数
key_thick = -1 # 关键点的粗细
key_text_thick = 2 # 关键点上文字粗细
key_text_scale = 0.6 # 关键点上文字的放大倍数
key_radius = 4 # 关键点绘制半径
dot = 6 # 选择保留几位小数
key_color = {
0: (0, 0, 200),
1: (255, 0, 0),
2: (0, 222, 0)
} # 关键点的颜色
key_text_color = {
0: (0, 100, 200),
1: (255, 0, 0),
2: (0, 255, 125)
} # 关键点上文本的颜色
box_color = {
0: (125, 125, 125),
1: (0, 255, 0),
2: (0, 255, 0),
3: (255, 0, 0),
4: (0, 255, 255),
5: (255, 255, 0),
6: (255, 0, 255),
7: (0, 125, 125),
8: (125, 125, 125),
9: (125, 125, 125),
10: (125, 0, 125),
11: (125, 0, 125),
12: (125, 0, 125),
13: (125, 0, 125),
14: (125, 0, 125),
15: (125, 0, 125)
} # 每个不同类别框的颜色
my_cls = {
0: 'person',
1: '1',
2: '10',
3: 'Stiffeners',
4: 'test',
} # 添加自己的框的标签，如果没有就用i:'i'替代
final_class = {
i: my_cls[i] if i in my_cls else str(i) for i in range(16)
} # 框的默认名字
# 不要修改的参数
position = None # 这里判断鼠标放到了哪个点上，方便后面移动的时候做计算
label = None # 操作图像对应的标签
img = None # 操作的图像
Mouse_move = None # 选择移动框的标志位
label_index = None # 鼠标选中的框在标签中的位置
label_index_pos = None # 记录选中了框的8个点位的哪一个
Mouse_insert = None # 用来记录是否进入删除状态
draw_rectangle = None # 用来记录开始添加新框
end_draw_rectangle = None # 用来记录结束绘制新框
append_str_temp = None # 用来保存新增加的框的信息
empty_label = None # 本地是否存在标签文件标志
# 关键点相关的参数
key_points = None
key_points_move = None
key_x = None # 移动关键点的时候记录其每个关键点的x
key_y = None # 移动关键点的时候记录其每个关键点的y
key_v = None # 移动关键点的时候记录其每个关键点的状态
key_box = None
box_move = None # 移动的是框的时候的标志位
key_insert = None # 对某个关键点双击，切换其状态
move_key_point = None # 把其他位置的关键点移动到这个地方
la_path = None
key_point_one = None # 使用双击移动关键点的时候，记录第一个按下的键
key_point_two = None # 使用双击移动关键点的时候，记录第二个按下的键
append_new_key_point = None # 增加第二个关键点
append_new_key_point_index = 0 # 增加第二个关键点
window_w = None # 获取创建窗口的宽度
window_h = None # 获取创建窗口的高度
def flag_init():
# 初始化下参数
global position, label, img, Mouse_insert, Mouse_move, label_index, draw_rectangle, end_draw_rectangle, append_str_temp, empty_label, \
label_index_pos, key_v, key_x, key_y, key_x, key_box, key_points_move, key_points, box_move, move_key_point, window_w, window_h
position = None
Mouse_move = None
label_index = None
label_index_pos = None
Mouse_insert = None
draw_rectangle = None
end_draw_rectangle = None
append_str_temp = None
empty_label = None
key_points = None
key_points_move = None
key_x = None
key_y = None
key_v = None
key_box = None
box_move = None
move_key_point = None
window_w = None
window_h = None
# 用来绘制小的填充矩形框
def draw_rect_box(img, center, length_1, color=(0, 0, 255)):
x1, y1 = center[0] - length_1, center[1] - length_1
x2, y2 = center[0] + length_1, center[1] + length_1
cv2.rectangle(img, (x1, y1), (x2, y2), color, thickness=-1)
# 用来读取本地图像
def img_read(img_path, scale_):
global window_w, window_h
# scale_填写屏幕的最小尺寸
image = cv2.imread(img_path)
scale_x, scale_y, _ = image.shape
if max(scale_x, scale_y) > scale_ and window_w is None:
scale = max(scale_x, scale_y) / scale_
image = cv2.resize(image, (int(image.shape[1] / scale), int(image.shape[0] / scale)))
if window_w is not None:
image = cv2.resize(image, (window_w, window_h))
return image
# 判断两点的间距，用来判断鼠标所在位置是否进入了8个点所在的区域
def distance(p1, p2):
global circle_distance
if math.sqrt((p2[0] - p1[0]) ** 2 + (p2[1] - p1[1]) ** 2) < circle_distance:
return True
else:
return False
# 绘制虚线矩形框，当切换到删除时，由实线框转为虚线框
def draw_dotted_rectangle(img, pt1, pt2, length_1=5, gap=6, thick=2, color=(100, 254, 100)):
(x1, y1), (x2, y2) = pt1, pt2
temp1, temp2 = x1, y1
while x1 + length_1 < x2:
cv2.line(img, (x1, y1), (x1 + length_1, y1), color, thickness=thick)
cv2.line(img, (x1, y2), (x1 + length_1, y2), color, thickness=thick)
x1 += length_1 + gap
while y1 + length_1 < y2:
cv2.line(img, (temp1, y1), (temp1, y1 + length_1), color, thickness=thick)
cv2.line(img, (x1, y1), (x1, y1 + length_1), color, thickness=thick)
y1 += length_1 + gap
# 把本地标签展示到图像中
def label_show(img1, label_path, index):
global small_box_thickness, box_thickness, label_fontScale, label_thickness, key_point_is, key_points, \
key_radius, key_color, key_thick, key_text_scale, key_text_thick, key_text_color, label, draw_line_circle
with open(la_path) as f:
label = f.readlines()
if len(label) == 0:
return
for i, points in enumerate(label):
if key_point_is:
# 获取关键点参数
key_points = points.split(' ')[5:]
points = points.split(' ')[0:5]
classify = int(float(points[0]))
points.pop(0)
point = [float(s.strip('\n')) for s in points]
# point = list(map(float, points))
scale_y, scale_x, _ = img1.shape
x, y, w, h = int((point[0] - point[2] / 2) * scale_x), int(
(point[1] - point[3] / 2) * scale_y), int(
point[2] * scale_x), int(point[3] * scale_y)
if i == index:
draw_dotted_rectangle(img1, (x, y), (x + w, y + h), box_thickness)
else:
cv2.rectangle(img1, (x, y), (x + w, y + h), box_color[classify], thickness=box_thickness)
if draw_line_circle:
# 绘制边上中心点，与四个顶点,矩形框中心点
draw_rect_box(img1, (x, int(0.5 * (y + y + h))), length_1=small_box_thickness)
draw_rect_box(img1, (x + w - 1, int(0.5 * (y + y + h))), length_1=small_box_thickness)
draw_rect_box(img1, (int(0.5 * (x + x + w)), y), length_1=small_box_thickness)
draw_rect_box(img1, (int(0.5 * (x + x + w)), y + h), length_1=small_box_thickness)
draw_rect_box(img1, (x, y), length_1=small_box_thickness)
draw_rect_box(img1, (x + w, y), length_1=small_box_thickness)
draw_rect_box(img1, (x + w, y + h), length_1=small_box_thickness)
draw_rect_box(img1, (x, y + h), length_1=small_box_thickness)
draw_rect_box(img1, (int(x + 0.5 * w), int(y + 0.5 * h)), length_1=small_box_thickness)
cv2.putText(img1, str(final_class[classify]), (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, label_fontScale,
(255, 0, 255), label_thickness)
if key_point_is:
# 依次获取每个关键点
key_x = [float(i) for i in key_points[::3]]
key_y = [float(i) for i in key_points[1::3]]
key_v = [int(float(i)) for i in key_points[2::3]]
index = 0
key_point = zip(key_x, key_y)
for p in key_point:
cv2.circle(img, (int(p[0] * scale_x), int(p[1] * scale_y)), key_radius, key_color[key_v[index]],
thickness=key_thick,
lineType=cv2.LINE_AA)
cv2.putText(img, str(index), (int(p[0] * scale_x - 5), int(p[1] * scale_y - 10)),
cv2.FONT_HERSHEY_SIMPLEX,
key_text_scale, key_text_color[0], key_text_thick)
index += 1
key_points = None
# 回调函数，用于记录鼠标操作
def mouse_event(event, x, y, flag, param):
global label, img, position, Mouse_move, label_index, label_index_pos, dot, Mouse_insert, draw_rectangle, \
end_draw_rectangle, key_points, key_v, key_x, key_y, key_x, key_box, key_points_move, box_move, \
key_insert, label_path, move_key_point
scale_y, scale_x, _ = img.shape
# 鼠标如果位于8个点左右，即通过position记录当前位置，通过主函数在鼠标附近绘制空心圈
# 通过label_index记录鼠标选择了第几个框，通过label_index_pos记录该框第几个点被选中了
with open(la_path) as f:
label = f.readlines()
if move_key_point is None and key_insert is None and Mouse_insert is None and empty_label is None and event == cv2.EVENT_MOUSEMOVE and img is not None and label is not None and \
Mouse_move is None:
for i, la in enumerate(label):
la = la.strip('\n').split(' ')
if key_point_is:
key_points = list(map(float, la))[5:]
la = list(map(float, la))[0:5]
x1, y1 = int((la[1] - la[3] / 2) * scale_x), int((la[2] - la[4] / 2) * scale_y)
x2, y2 = x1 + int(la[3] * scale_x), y1 + int(la[4] * scale_y)
# 这里判断鼠标放到了哪个点上，方便后面移动的时候做计算
if distance((x, y), (x1, y1)):
label_index_pos = 0
position = (x, y)
label_index = i
box_move = True
key_points_move = None
break
elif distance((x, y), (x2, y2)):
label_index_pos = 1
position = (x, y)
label_index = i
box_move = True
key_points_move = None
break
elif distance((x, y), (x1, int(0.5 * y1 + 0.5 * y2))):
label_index_pos = 2
position = (x, y)
label_index = i
box_move = True
key_points_move = None
break
elif distance((x, y), (int((x1 + x2) / 2), y2)):
label_index_pos = 3
position = (x, y)
label_index = i
box_move = True
key_points_move = None
break
elif distance((x, y), (int((x1 + x2) / 2), y1)):
label_index_pos = 4
position = (x, y)
label_index = i
box_move = True
key_points_move = None
break
elif distance((x, y), (x2, int(0.5 * y1 + 0.5 * y2))):
label_index_pos = 5
position = (x, y)
label_index = i
box_move = True
key_points_move = None
break
elif distance((x, y), (x1, y2)):
label_index_pos = 6
position = (x, y)
label_index = i
box_move = True
key_points_move = None
break
elif distance((x, y), (x2, y1)):
label_index_pos = 7
position = (x, y)
label_index = i
box_move = True
key_points_move = None
break
elif distance((x, y), ((x1 + x2) / 2, (y1 + y2) / 2)):
# 框中心
label_index_pos = 8
position = (x, y)
label_index = i
box_move = True
key_points_move = None
break
else:
label_index_pos = None
position = None
label_index = None
if key_point_is:
# 判断鼠标是不是放到了关键点上
key_x = [float(i) for i in key_points[::3]]
key_y = [float(i) for i in key_points[1::3]]
key_v = [float(i) for i in key_points[2::3]] # 能见度
if len(key_x) == len(key_v) and len(key_x) == len(key_y):
for index, key_ in enumerate(key_x):
if distance((x, y), (int(key_ * scale_x), int(key_y[index] * scale_y))):
position = (x, y)
label_index, label_index_pos = i, index
key_box = la
key_points_move = True
box_move = None
break
# 这里到下一个注释都是为了移动已有的框做准备
if position is not None and event == cv2.EVENT_LBUTTONDOWN:
Mouse_move = True
position = None
# 首先判断鼠标选择了该框的第几个点，然后移动鼠标的时候只负责移动该点
if Mouse_move and box_move:
# 先把要移动的框的标签记录下来，然后删除，添加到末尾，不断修改末尾标签来达到移动框的目的
# temp_label用来记录标签
temp_label = label[label_index]
label.pop(label_index)
temp_label = temp_label.strip('\n').split(' ')
temp_label = [float(i) for i in temp_label]
x_1, y_1 = (temp_label[1] - 0.5 * temp_label[3]), (temp_label[2] - 0.5 * temp_label[4])
x_2, y_2 = x_1 + temp_label[3], y_1 + temp_label[4]
# 判断移动的是8个点中的哪个
if label_index_pos == 0:
x_1, y_1 = x / scale_x, y / scale_y
elif label_index_pos == 1:
x_2, y_2 = x / scale_x, y / scale_y
elif label_index_pos == 2:
x_1 = x / scale_x
elif label_index_pos == 3:
y_2 = y / scale_y
elif label_index_pos == 4:
y_1 = y / scale_y
elif label_index_pos == 5:
x_2 = x / scale_x
elif label_index_pos == 6:
x_1, y_2 = x / scale_x, y / scale_y
elif label_index_pos == 7:
y_1, x_2 = y / scale_y, x / scale_x
elif label_index_pos == 8:
x_1, y_1 = x / scale_x - (abs(temp_label[3]) / 2), y / scale_y - (abs(temp_label[4]) / 2)
x_2, y_2 = x / scale_x + (abs(temp_label[3]) / 2), y / scale_y + (abs(temp_label[4]) / 2)
# 把移动后的点信息保存下来添加到标签中，以此形成动态绘制一个框的效果
temp_label[0], temp_label[1], temp_label[2], temp_label[3], temp_label[4] = str(
round((int(temp_label[0])), dot)), \
str(round(((x_1 + x_2) * 0.5), dot)), str(round(((y_1 + y_2) * 0.5), dot)), str(
round((abs(x_1 - x_2)), dot)), str(round((abs(y_1 - y_2)), dot))
temp_label = [str(i) for i in temp_label]
str_temp = ' '.join(temp_label) + '\n'
label.append(str_temp)
label_index = len(label) - 1
elif Mouse_move and key_points_move:
label.pop(label_index)
key_x[label_index_pos] = round(x / scale_x, dot)
key_y[label_index_pos] = round(y / scale_y, dot)
key_box[0] = int(key_box[0])
str_temp = ' '.join([str(j) for j in key_box])
for index, kx in enumerate(key_x):
str_temp += ' ' + str(kx) + ' ' + str(key_y[index]) + ' ' + str(int(key_v[index]))
label.append(str_temp)
label_index = len(label) - 1
if Mouse_move and event == cv2.EVENT_LBUTTONUP:
flag_init()
# 这里是为了删除框
if key_point_is is None and Mouse_insert is None and position is not None and event == cv2.EVENT_LBUTTONDBLCLK and Mouse_move is None:
Mouse_insert = label_index
if key_point_is and event == cv2.EVENT_LBUTTONDBLCLK and Mouse_move is None and key_points_move and box_move is None:
key_insert = label_index_pos
if key_point_is and event == cv2.EVENT_LBUTTONDBLCLK and Mouse_insert is None and key_insert is None and position is None:
move_key_point = (x, y)
# 这里是为了增加新的框
if key_point_is is None and Mouse_insert is None and position is None and Mouse_move is None and event == cv2.EVENT_LBUTTONDOWN and end_draw_rectangle is None:
draw_rectangle = [(x, y), (x, y)]
# 如果鼠标左键一直没有松开，则不断更新第二个点的位置
elif Mouse_insert is None and draw_rectangle is not None and event == cv2.EVENT_MOUSEMOVE and end_draw_rectangle is None:
draw_rectangle[1] = (x, y)
# 鼠标松开了，最后记录松开时鼠标的位置，现在则记录了开始和松开鼠标的两个位置
# 如果两个位置太近，则不添加
elif Mouse_insert is None and draw_rectangle is not None and event == cv2.EVENT_LBUTTONUP:
if end_draw_rectangle is None:
draw_rectangle[1] = (x, y)
if not distance(draw_rectangle[0], draw_rectangle[1]):
end_draw_rectangle = True
else:
draw_rectangle = None
def create_file_key(img_path, label_path):
empty_la = None
if not os.path.exists(label_path):
with open(label_path, 'w') as f:
pass
empty_la = True
with open(label_path) as f:
label_ = f.readlines()
if len(label_) == 0 or label_[0] == '\n':
empty_la = True
img_s = img_read(img_path, 950) # 950调整图像的大小
if key_point_is and empty_la:
box_create = '0 0.5 0.5 0.3 0.3 '
len_t = img_s.shape[1] // key_point_num
key_num_x = [str(round((i * len_t + 20) / img_s.shape[1], dot)) + ' ' + str(0.5) + ' ' + '2' for i in
range(key_point_num)]
with open(label_path, 'w') as f:
f.write(box_create + ' '.join(key_num_x))
def main(img_path, label_path):
global img, position, label, Mouse_insert, draw_rectangle, end_draw_rectangle, append_str_temp, empty_label, \
Mouse_move, dot, box_move, key_insert, key_point_one, key_point_two, key_x, key_y, key_v, \
move_key_point, append_new_key_point, append_new_key_point_index, window_w, window_h
# 判断本地是否存在文件，或者文件中是否为空或者存在一个换行符，就先把标签删除，添加'0 0 0 0 0\n'
# 如果不预先添加一个处理起来有点麻烦，这里就先加一个，然后后面删掉就行了
if not os.path.exists(label_path):
empty_label = True
with open(label_path, 'w') as f:
pass
with open(label_path) as f:
label = f.readlines()
if len(label) == 0 or label[0] == '\n':
empty_label = True
# 这里的2是将原图缩小为2分之一
print(img_path)
img_s = img_read(img_path, 900)
if key_point_is and empty_label:
box_create = '0 0.5 0.5 0.3 0.3 '
len_t = img_s.shape[1] // key_point_num
key_num_x = [str(round((i * len_t + 20) / img_s.shape[1], dot)) + ' ' + str(0.5) + ' ' + '2' for i in
range(key_point_num)]
with open(label_path, 'w') as f:
f.write(box_create + ' '.join(key_num_x))
label = box_create + ' '.join(key_num_x)
# 创建回调函数，绑定窗口
cv2.namedWindow('image', cv2.WINDOW_NORMAL)
_, _, window_w, window_h = cv2.getWindowImageRect('image')
cv2.resizeWindow('image', img_s.shape[1], img_s.shape[0])
cv2.setMouseCallback('image', mouse_event)
# 刷新图像的地方
while True:
# 首先读取下标签，用来初始化显示
with open(label_path, 'w') as f:
for i in label:
f.write(i)
# 如果鼠标选中了框的8个点之一，就在鼠标周围绘制空心圈
if Mouse_insert is None and draw_rectangle is None and position is not None and key_insert is None:
img = img_s.copy()
label_show(img, label_path, Mouse_insert)
cv2.circle(img, position, 10, (0, 255, 100), 2)
# 如果选择开始增加新的框，则不断绘制鼠标起始点和移动过程之间形成的框
elif draw_rectangle is not None and end_draw_rectangle is None:
img = img_s.copy()
label_show(img, label_path, Mouse_insert)
cv2.rectangle(img, draw_rectangle[0], draw_rectangle[1], color=box_color[1], thickness=2)
# 当松开鼠标后，记录两点位置，并提示选择类别
elif draw_rectangle is not None and end_draw_rectangle:
scale_y, scale_x, _ = img.shape
x1, y1 = draw_rectangle[0]
x2, y2 = draw_rectangle[1]
w1, h1 = abs(x2 - x1), abs(y2 - y1)
append_str_temp = str(round((x1 + x2) / 2 / scale_x, dot)) + ' ' + str(
round((y1 + y2) / 2 / scale_y, dot)) + ' ' + \
str(round((w1 / scale_x), dot)) + ' ' + str(round((h1 / scale_y), dot)) + '\n'
cv2.putText(img, 'choose your classify', (0, img.shape[0] // 2 - 30),
cv2.FONT_HERSHEY_SIMPLEX,
0.7, (255, 0, 255), 2)
cv2.putText(img, ' '.join([str(i) + ':' + my_cls[i] for i in my_cls]), (0, img.shape[0] // 2 + 30),
cv2.FONT_HERSHEY_SIMPLEX,
0.7, (100, 255, 255), 2)
elif key_insert is not None:
position, Mouse_move, box_move = None, None, None # 禁用其他操作
cv2.putText(img, 'Switching visibility: 0 1 2', (0, img.shape[0] // 2 - 30),
cv2.FONT_HERSHEY_SIMPLEX,
1, (100, 255, 255), 2,
lineType=cv2.LINE_AA)
elif move_key_point is not None:
position, Mouse_move, box_move = None, None, None # 禁用其他操作
cv2.putText(img, 'choose point: 0 - {}'.format(key_point_num - 1), (0, img.shape[0] // 2 - 30),
cv2.FONT_HERSHEY_SIMPLEX,
1, (100, 255, 255), 2,
lineType=cv2.LINE_AA)
# 如果什么标志都没有，就正常显示一个图
else:
img = img_s.copy()
if Mouse_insert is not None:
position, Mouse_move = None, None
cv2.putText(img, 'delete: W, exit: E', (0, img.shape[0] // 2 - 30),
cv2.FONT_HERSHEY_SIMPLEX,
1, (100, 255, 255), 2,
lineType=cv2.LINE_AA)
label_show(img, label_path, Mouse_insert)
cv2.imshow('image', img)
# key用来获取键盘输入
key = cv2.waitKey(10)
# 输入为Q则退出
if key == ord('Q'):
append_new_key_point = None
# 退出按键
break
if move_key_point is not None and key_point_one is None and 48 <= key <= 57:
key_point_one = int(chr(key))
key = 0
if move_key_point is not None and key_point_two is None and 48 <= key <= 57:
key_point_two = int(chr(key))
key = 0
if (move_key_point is not None) and (key_point_one is not None) and (key_point_two is not None):
with open(la_path) as f:
label = f.readlines()
for i, la in enumerate(label):
la = la.strip('\n').split(' ')
key_points_ = list(map(float, la))[5:]
key_box_ = list(map(float, la))[0:5]
key_x_ = [float(i) for i in key_points_[::3]]
key_y_ = [float(i) for i in key_points_[1::3]]
key_v_ = [float(i) for i in key_points_[2::3]] # 能见度
key_box_[0] = int(key_box_[0])
index_ = key_point_one * 10 + key_point_two
if index_ >= key_point_num:
break
key_x_[index_] = round(move_key_point[0] / img.shape[1], dot)
key_y_[index_] = round(move_key_point[1] / img.shape[0], dot)
str_temp = ' '.join([str(j) for j in key_box_])
for index, kx in enumerate(key_x_):
str_temp += ' ' + str(kx) + ' ' + str(key_y_[index]) + ' ' + str(int(key_v_[index]))
label = str_temp
with open(la_path, 'w') as f:
f.write(str_temp)
move_key_point, key_point_one, key_point_two = None, None, None
break
move_key_point, key_point_one, key_point_two = None, None, None
# 如果按键输入为W则删除选中的框
if Mouse_insert is not None and key == ord('W'):
label.pop(Mouse_insert)
Mouse_insert = None
elif key_insert is not None and key == ord('0'):
with open(label_path, 'r') as f:
label_temp = f.read()
str_temp = label_temp.split(' ')
str_temp[3 * int(key_insert) + 7] = '0'
str_temp[3 * int(key_insert) + 7 - 1] = '0'
str_temp[3 * int(key_insert) + 7 - 2] = '0'
with open(label_path, 'w') as f:
f.write(' '.join(str_temp))
label = ' '.join(str_temp)
key_insert = None
elif key_insert is not None and key == ord('1'):
with open(label_path, 'r') as f:
label_temp = f.read()
str_temp = label_temp.split(' ')
str_temp[3 * int(key_insert) + 7] = '1'
with open(label_path, 'w') as f:
f.write(' '.join(str_temp))
label = ' '.join(str_temp)
key_insert = None
elif key_insert is not None and key == ord('2'):
with open(label_path, 'r') as f:
label_temp = f.read()
str_temp = label_temp.split(' ')
str_temp[3 * int(key_insert) + 7] = '2'
with open(label_path, 'w') as f:
f.write(' '.join(str_temp))
label = ' '.join(str_temp)
key_insert = None
# 如果输入为E则从选中框的状态退出
elif key == ord('E'):
Mouse_insert = None
# ——————————————————————————————————————————————————————————————————————————————————————
# 通过键盘获取输入的类别
# 原本是设置了从键盘输入0-9为框加入标签，这里又添加了 Z-》10， X-》11，C-》12......
# 如果要增加新的映射按键应该仿照这个格式 条件中加入 or key == ord('P')
# 下面追加else if: key == ord('P') str_temp = str(num) + ' ' + append_str_temp
# ——————————————————————————————————————————————————————————————————————————————————————
elif Mouse_move is None and Mouse_insert is None and draw_rectangle is not None and end_draw_rectangle is not None \
and (48 <= key <= 57 or key == ord('Z') or key == ord('X') or key == ord('C') or key == ord('V')
or key == ord('B') or key == ord('N')):
if 48 <= key <= 57:
str_temp = str(chr(key)) + ' ' + append_str_temp
elif key == ord('Z'):
str_temp = str(10) + ' ' + append_str_temp
elif key == ord('X'):
str_temp = str(11) + ' ' + append_str_temp
elif key == ord('C'):
str_temp = str(12) + ' ' + append_str_temp
elif key == ord('V'):
str_temp = str(13) + ' ' + append_str_temp
elif key == ord('B'):
str_temp = str(14) + ' ' + append_str_temp
elif key == ord('N'):
str_temp = str(15) + ' ' + append_str_temp
label.append(str_temp)
append_str_temp, draw_rectangle, end_draw_rectangle, empty_label = None, None, None, None
elif key == ord('R'):
flag_init()
append_new_key_point = True
break
# 按下T退出
elif key == ord('T'):
exit(0)
# 按下Y删除当前图片和对应的标签
elif key == ord('Y'):
os.remove(img_path)
os.remove(label_path)
break
def delete_line_feed(label_path):
# 去掉最后一行的换行符'\n'，保存的时候需要
if os.path.exists(label_path):
with open(label_path) as f:
label_ = f.read()
label_ = label_.rstrip('\n')
with open(label_path, 'w') as f:
f.write(label_)
def append__line_feed(label_path):
# 加上最后一行的换行符'\n'，标注的时候增加新的框的时候需要
with open(label_path) as f:
label_ = f.read()
if len(label_) < 4:
with open(label_path, 'w') as f:
pass
return
label_ = label_.rstrip('\n') + '\n'
with open(label_path, 'w') as f:
f.write(label_)
def key_check(label_path):
# 检查开启关键点之后本地标签是否满足要求, 如果本地标签中和预设关键点数不等以及关键点数量不是3的倍数都会将原有标签重置
if os.path.exists(label_path):
with open(label_path) as f:
label_ = f.readlines()
for label_ in label_:
label_ = label_.strip('\n').split(' ')
if ((len(label_) - 5) % 3) or ((len(label_) - 5) // 3 - key_point_num):
with open(label_path, 'w') as f:
pass
def label_check(label_path):
# 检查普通标签,判断每行是包含5个数值
if os.path.exists(label_path):
with open(label_path) as f:
label_ = f.readlines()
for i in label_:
i = i.strip('\n').split(' ')
if len(i) - 5 != 0:
with open(label_path, 'w'):
pass
def merge_file_key(la_path, index):
with open(la_path) as f:
text = f.read().strip('\n')
for i in range(index):
with open(la_path.split('.')[0] + str(i) + '.txt') as f:
text += '\n' + f.read().strip('\n')
os.remove(la_path.split('.')[0] + str(i) + '.txt')
with open(la_path, 'w') as f:
f.write(text)
if __name__ == '__main__':
image_ = os.listdir(image_path)
for im in image_:
flag_init()
im_path = os.path.join(image_path, im)
la_path = os.path.join(label_path, im.split('.')[0] + '.txt')
if key_point_is:
key_check(la_path) # 检查本地标签的关键点数量是否和预设的关键点数量相等，以及去除框的5点后点数是否满足为3的倍数
create_file_key(im_path, la_path)
else:
delete_line_feed(la_path)
label_check(la_path)
if os.path.exists(la_path):
# 先增加一个换行符为了后面的增加框的操作
append__line_feed(la_path)
while True:
main(im_path, la_path)
if append_new_key_point is None:
break
else:
la_path = os.path.join(label_path, im.split('.')[0] + str(append_new_key_point_index) + '.txt')
with open(la_path, 'w') as f:
pass
if key_point_is:
key_check(la_path) # 检查本地标签的关键点数量是否和预设的关键点数量相等，以及去除框的5点后点数是否满足为3的倍数
create_file_key(im_path, la_path)
else:
delete_line_feed(la_path)
label_check(la_path)
append_new_key_point_index += 1
if append_new_key_point_index != 0:
merge_file_key(os.path.join(label_path, im.split('.')[0] + '.txt'), append_new_key_point_index)
append_new_key_point_index = 0
if os.path.exists(la_path):
# 去掉最后一行的换行符
delete_line_feed(la_path)