water_turbidity_det/labelme/check_label.py

"""检查标注是否正确，读取标签，信息"""
import cv2
import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from utils import draw_grid
from dotenv import load_dotenv
load_dotenv() # 加载环境变量

patch_w = int(os.getenv('PATCH_WIDTH', 256))
patch_h = int(os.getenv('PATCH_HEIGHT', 256))
scale = 2
clicked_points = []  # u v w h cls

def main():
    """检查标注结果是否正确"""
    # 标注文件路径
    global patch_w
    global patch_h
    global scale
    global clicked_points
    # TODO:修改为要检查的图片路径
    imgs_path = r'D:\code\water_turbidity_det\data\video4_20251129120320_20251129123514'
    output_root = r'D:\code\water_turbidity_det\check'
    label_path = os.path.join(imgs_path, 'label.txt')
    if os.path.exists(label_path):
        with open(label_path, 'r') as fr:
            lines = fr.readlines()
        lines = [line.strip() for line in lines]
        for line in lines:
            point = line.split(',')
            clicked_points.append([int(point[0])//scale, int(point[1])//scale, int(point[2])//scale, int(point[3])//scale, int(point[4])])
        del lines
    # 检查结果输出路径
    output_path = os.path.join(output_root, os.path.basename(imgs_path)+'_check')
    if not os.path.exists(output_path):
        os.makedirs(output_path)
    # 获取所有照片
    all_imgs = os.listdir(imgs_path)
    all_imgs = [img for img in all_imgs if img.split('.')[-1] == 'jpg' or img.split('.')[-1] == 'png']

    for img in all_imgs:
        img_path = os.path.join(imgs_path, img)
        img = cv2.imread(img_path)
        img = cv2.resize(img, (img.shape[1] // scale, img.shape[0] // scale))
        # 绘制网格线
        img = draw_grid(img, patch_w // scale, patch_h // scale)
        # 绘制标记点
        for point in clicked_points:
            # 计算中心点
            center_x = point[0]+point[2]//scale
            center_y = point[1]+point[3]//scale
            cv2.circle(img, (center_x, center_y), 5, (0, 0, 255), -1)
            # 显示标签文本
            cv2.putText(img, str(point[4]), (center_x + 10, center_y + 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
        cv2.imwrite(os.path.join(output_path, os.path.basename(img_path)),
                    img)
        print(f"检查结果保存在: {os.path.join(output_path, os.path.basename(img_path))}")

if __name__ == '__main__':
    main()