计算机视觉 目标检测 使用opencv自制一个YOLO常规数据和关键点数据的标注工具

目录

前言普通标注下关键点标注下

标注 YOLO格式的数据集标注关键点一个图中标注多个框和关键点

最终代码

前言

工具由opencv编写，可以直接生成YOLO所需要的标签（pose和常规标签）

代码放到了文章末尾，以及百度云下载链接首先放一段实际操作的视频展示

yolov5数据集标注，yolo-pose数据集标注

普通标注下

按Q切换到下一张图像按T直接退出按Y删除当前图片和对应标签按R隐藏当前内容，继续标注完成之后会和之前的内容合并双击框上的点切换到删除状态，按下E不做操作，按下W删除鼠标左键拉伸和移动当前框

关键点标注下

按Q切换到下一张图像按T退出按Y删除当前图片和对应标签按R隐藏当前内容，继续标注完成之后会和之前的内容合并双击关键点切换关键点的可见状态（0 1 2）可以拖动关键点双击空白区域，输入0和8则->010+18=8,将8号关键点移动到双击处

标注 YOLO格式的数据集

首先需要把22行的代码的初始值从True改为None，为常规标注 更换图像所在的路径和保存路径，保存路径中可以不放任何内容

标注关键点

-把22行的代码的初始值改为True，表示标注关键点，并且设置好相应的key_point_num（要标注的有几个关键点）

和普通标注不同的是这里双击矩形框不会出现任何操作，表示你不可以删除矩形框，只能移动，也不可以删除关键点，也只能移动，但是关键点除了拖动也可以通过双击和键盘输入移动

可以鼠标左键双击关键点

因为关键点有0 1 2三种状态，可见等状态 然后从键盘输入0或者1或者2

设为1仅仅是标志位从2->1，显示颜色发生改变，设置为2因为之前的状态就是2，所以看起来没有变化，设置为0，除了把状态变为0，还要把相应的坐标（x ，y）设置为0 0如果只是想要把状态设置为0，坐标不变，需要注释掉481和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'

如果想把状态为0的点改回来需要把鼠标放到左上角，双击，或者先拖拽到其他显眼区域

接下来选择前三个点展示起始标签和更改完后的标签

初始标签

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

前三个关键点标志位都是2

直接运行，可以发现已经生成好了一个框和key_point_num个关键点，本地同时生成了相应的标签这是因为如果不预先生成一个框和key_point_num个关键点，处理起来会相当复杂，而且当关键点数量多余10个也不好处理，因此预先生成关键点和框，关键点的状态为2，可以自行修改标注的时候把框和点拖拽到相应的位置，把关键点状态设为所需的即可

拖动关键点和框的位置，对于关键点也可以采用另一种标注方式 双击不存在关键点和框的9个点的区域 双击完成后会提示选则本0-10，因为这里关键点一共11个，从键盘输入两次，都要输入0-9的数字才行，如果像选中个位关键点例如0-9则第一次输入0，第二次输入关键点索引

例如要操作8号关键点，第一次输入0第二次输入8，可以看到关键点跑到了鼠标双击的位置

当该图标注完成，点击Q切换，都完成后可以看到本地相应的生成了对应的关键点标签，按T就是总退出 标注的时候也可以拉伸窗口，不影响标注结果，总之注意不要叉掉窗口或者强制退出，尽量还是按Q退出

一个图中标注多个框和关键点

虽然预先生成了一个框和n个关键点，但是比如一个图里面有多个人，比如有两个人，那就需要把两个人都标注在图像中，为了处理这种情况

首先先正常标注 然后这时候不要按Q退出，如果你想标注第二个人，这时候按住R

然后继续标注即可 如果你还想要继续标注就继续按住W，直到你所有的都完成后按住Q切换到下一个图，当所有的都完成之后，可以 看到本地中的标签文件中，已经成功的添加了两个标注物体 此时我们再重新运行程序，可以看到已经成功标注了 但是需要注意的是，这种一个图中含有多个框的，就不要做任何操作了，比如点击，移动之类的，就是只可以看，如果你操作了，只会保留标签中第一行，其余的都会被删掉

最终代码

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)

精彩文章