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)