735 lines
32 KiB
Python
735 lines
32 KiB
Python
import math
|
||
import os
|
||
import cv2
|
||
|
||
"""
|
||
标注关键点只能存在一个框和多个点，并且不能删除点和删除框，读取本地文件的关键点要保证其中的关键点
|
||
数和key_point_num的值是一样的，本地标签中如果只存在框的信息就不要使用该脚本标注，不然会出错，
|
||
本地文件夹中可以有标签，如果有会优先加载本地标签，没有才会创建一个
|
||
关键点标注
|
||
这个是默认的标注就是普通标注
|
||
按Q切换下一张,R清除干扰再标注（把之前的框屏蔽）
|
||
按Y从本地把图像和标签一切删掉
|
||
按T 退出
|
||
鼠标双击可以删除框
|
||
单机就是拖拽
|
||
|
||
"""
|
||
draw_line_circle = True # True/None 是否在框上绘制点（8个点）
|
||
key_point_is = None # 是否标记关键点 设置为None标注普通yolo标签
|
||
# 可以自定义得参数
|
||
image_path = R'C:\Users\lengdan\Desktop\data1\images\images' # 标注完成保存到的文件夹
|
||
label_path = R'C:\Users\lengdan\Desktop\data1\images\labels' # 要标注的图像所在文件夹
|
||
circle_distance = 10 # 半径范围:鼠标进入点的半径范围内会出现光圈
|
||
key_point_num = 5 # 关键点个数
|
||
box_thickness = 2 # 框的粗细
|
||
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: '0',
|
||
1: '1',
|
||
2: '10',
|
||
3: '11',
|
||
4: '12',
|
||
5: '13',
|
||
6: '14',
|
||
7: '15',
|
||
8: '2',
|
||
9: '3',
|
||
10: '4',
|
||
11: '5',
|
||
12: '6',
|
||
13: '7',
|
||
14: '8',
|
||
15: '9'
|
||
} # 添加自己的框的标签，如果没有就用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
|
||
# 通过键盘获取输入的类别
|
||
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
|
||
elif key == ord('T'):
|
||
exit(0)
|
||
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)
|