# 要求保证视频不能移动，且全过程没有任何遮挡物 import os import sys sys.path.append(os.path.dirname(os.path.abspath(__file__))) from utils import draw_grid import cv2 import numpy as np import tkinter as tk from tkinter import simpledialog from dotenv import load_dotenv load_dotenv() # 加载环境变量 class DrawTool: """重新绘制图像的工具""" def __init__(self, patch_w, patch_h, scale): self.img_path = None self.patch_w = patch_w self.patch_h = patch_h self.scale = scale def draw_new_img(self, points): img = cv2.imread(self.img_path) img = cv2.resize(img, (img.shape[1] // self.scale, img.shape[0] // self.scale)) draw_grid(img, self.patch_w // self.scale, self.patch_h // self.scale) for p in points: # 绘制角点 cv2.circle(img, (p[0], p[1]), 5, (255, 0, 0), -1) # 绘制中心点 circle_x_center = p[0] + self.patch_w//(2*self.scale) circle_y_center = p[1] + self.patch_h//(2*self.scale) cv2.circle(img, (circle_x_center, circle_y_center), 5, (0, 0, 255), -1) # 标注类别 cv2.putText(img, str(p[4]), (circle_x_center + 10, circle_y_center + 10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2) return img def set_path(self, img_path): """设置保存路径""" self.img_path = img_path # 所要切分的图块宽高 patch_w = int(os.getenv('PATCH_WIDTH', 256)) patch_h = int(os.getenv('PATCH_HEIGHT', 256)) scale = 2 draw_tool = DrawTool(patch_w=patch_w, patch_h=patch_h, scale=scale) # 存储标记点 clicked_points = [] def get_text_input(prompt): """创建弹窗获取文本输入""" root = tk.Tk() root.withdraw() # 隐藏主窗口 root.attributes('-topmost', True) # 确保弹窗置顶 result = simpledialog.askstring(' ', prompt) root.destroy() if result is None: result = "" if not result.strip().isdigit(): result = -1 return int(result) def det_same_pos_point(x_cor, y_cor): """判断是否重复点击""" global clicked_points for idx, point in enumerate(clicked_points): if point[0] == x_cor and point[1] == y_cor: return idx, True return -1, False def mouse_callback(event, x, y, flags, param): """ 鼠标回调函数 """ global clicked_points global patch_w global patch_h global scale global draw_tool if event == cv2.EVENT_LBUTTONDOWN: # 左键点击 # 在点击位置绘制红色圆点 scale_patch_w = patch_w // scale scale_patch_h = patch_h // scale # 格子角点 circle_x_corner = (x // scale_patch_w)*scale_patch_w circle_y_corner = (y // scale_patch_h)*scale_patch_h # 格子中心点 circle_x_center = circle_x_corner + scale_patch_w//2 circle_y_center = circle_y_corner + scale_patch_h//2 cv2.circle(param, (circle_x_center, circle_y_center), 5, (0, 0, 255), -1) cv2.circle(param, (circle_x_corner, circle_y_corner), 5, (255, 0, 0), -1) # 更新显示 cv2.imshow('img', param) cls = get_text_input('请输入类别:0.背景 1.浑浊 -1.不参与') # 显示标签文本 cv2.putText(param, str(cls), (circle_x_center + 10, circle_y_center + 10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2) # 更新显示 cv2.imshow('img', param) valid_cls = [0, 1, -1] if cls in valid_cls: print(f"点击网格角点: ({circle_x_corner}, {circle_y_corner}) 中心点: ({circle_x_center}, {circle_y_center}) 类别:{cls}") # 记录标注数据，角点 u v w h cls pos, is_exist = det_same_pos_point(circle_x_corner, circle_y_corner) # 判断是否重复点击 if is_exist: print(f"已存在该点: ({clicked_points[pos][0]}, {clicked_points[pos][1]}) 类别: {clicked_points[pos][4]}") clicked_points[pos][4] = cls print(f'重新标注该点: ({clicked_points[pos][0]}, {clicked_points[pos][1]}) 类别: {clicked_points[pos][4]}') else: print(f"添加点: ({circle_x_corner}, {circle_y_corner}) 类别: {cls}") clicked_points.append([circle_x_corner, circle_y_corner, scale_patch_w, scale_patch_h, cls]) else: print("请输入正确的类别！") elif event == cv2.EVENT_RBUTTONDOWN: # 右键点击 removed_point = clicked_points.pop() print(f"撤销标注点: ({removed_point[0]}, {removed_point[1]}) 类别: {removed_point[4]}") # 将撤销点标记为黑色 x = removed_point[0] y = removed_point[1] # 在点击位置绘制黑色圆点 scale_patch_w = patch_w // scale scale_patch_h = patch_h // scale # 格子角点 circle_x_corner = (x // scale_patch_w)*scale_patch_w circle_y_corner = (y // scale_patch_h)*scale_patch_h # 格子中心点 circle_x_center = circle_x_corner + scale_patch_w//2 circle_y_center = circle_y_corner + scale_patch_h//2 cv2.circle(param, (circle_x_center, circle_y_center), 5, (128, 128, 128), -1) cv2.circle(param, (circle_x_corner, circle_y_corner), 5, (128, 128, 128), -1) # 更新显示 cv2.imshow('img', param) def remove_duplicates(arr:list): """列表去重""" unique_list = [] [unique_list.append(item) for item in arr if item not in unique_list] return unique_list def play_video(video_path): global scale dir_name = os.path.dirname(video_path) for i in os.listdir(dir_name): frame = cv2.imread(os.path.join(dir_name, i)) if frame is None: continue # resize frame = cv2.resize(frame, (frame.shape[1] // scale, frame.shape[0] // scale)) cv2.imshow('Video', frame) # 按esc退出 if cv2.waitKey(20) == 27: cv2.destroyAllWindows() break def main(): """ 固定摄像头标注,只需要标注一张图像，后续图像保持一致 1.标注过程，先将图像划分为图框，用cv2划线工具在图像上划网格线 2.用鼠标进行交互，点击图块输入标签，按下空格键完成交互过程，保存标签 3.标签格式：u,v,w,h,label u,v为块左上角坐标，w,h为块的宽和高，label为块的标签 """ global clicked_points global patch_w global patch_h global scale # TODO: 需要更改为准备标注的图像路径，使用当前目录下的000000.jpg，结果保存在当前目录下label.txt img_path = r'D:\code\water_turbidity_det\label_data_tem\4_ch26_20260113145353_1\000300.jpg' play_video(img_path) img = cv2.imread(img_path) draw_tool.set_path(img_path) # resize 图像太大了显示不全 img = cv2.resize(img, (img.shape[1] // scale, img.shape[0] // scale)) # 绘制网格线 draw_grid(img, patch_w // scale, patch_h // scale) # 交互标注 print("操作说明:") print("- 点击鼠标左键在图像上添加红色标记点: 0.其他 1.浑浊 -1.忽略，不参与训练和测试") print("- 点击鼠标右键撤回上一个红色标记点") print("- 按 'c' 键清除所有标记点") print("- 按 ESC 键退出程序") cv2.namedWindow('img') cv2.setMouseCallback('img', mouse_callback, img) # 交互标注 while True: # 更新显示 cv2.imshow('img', draw_tool.draw_new_img(clicked_points)) key = cv2.waitKey(1) & 0xFF # 按 'c' 键清除所有标记点 if key == ord('c'): img = cv2.imread(img_path) img = cv2.resize(img, (img.shape[1] // scale, img.shape[0] // scale)) draw_grid(img, patch_w // scale, patch_h // scale) clicked_points.clear() cv2.setMouseCallback('img', mouse_callback, img) print("已清除所有标记点") # 按 ESC 键退出 elif key == 27: # ESC键 break cv2.destroyAllWindows() # 输出所有点击位置 # 列表去重 clicked_points = remove_duplicates(clicked_points) print(f"总共标记了 {len(clicked_points)} 个点:") for i, point in enumerate(clicked_points): print(f" 点 {i + 1}: ({point[0]}, {point[1]}, {point[2]}, {point[3]}, {point[4]})") # 恢复尺寸 clicked_points = [[p[0]*scale, p[1]*scale, p[2]*scale, p[3]*scale, p[4]] for p in clicked_points] # 写入txt if clicked_points: with open(os.path.join(os.path.dirname(img_path), 'label.txt'), 'w') as fw: for point in clicked_points: fw.write(f"{point[0]},{point[1]},{point[2]},{point[3]},{point[4]}\n") # 保存点 print(f"保存标记点 {len(clicked_points)} 个:") for i, point in enumerate(clicked_points): print(f" 点 {i + 1}: ({point[0]}, {point[1]}, {point[2]}, {point[3]}, {point[4]})") else : print("没有标记点！不保存任何文件！") if __name__ == '__main__': main()