water_turbidity_det/labelme/fixed_label.py

# 要求保证视频不能移动，且全过程没有任何遮挡物
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'/frame_data/test/video4_20251129120320_20251129123514\000000.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()