# 根据标注文件，生成patch，每个类别放在一个文件夹下
import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

import numpy as np
import cv2
import gc
from dotenv import load_dotenv
load_dotenv()

# 从图像中截取patch_w*patch_h大小的图像，并打上标签
patch_w = int(os.getenv('PATCH_WIDTH', 256))
patch_h = int(os.getenv('PATCH_HEIGHT', 256))

def main():
    # TODO：需要修改为标注好的图片路径
    input_path = r'D:\code\water_turbidity_det\data\4_video_202511211127'
    # TODO: 需要修改为保存patch的根目录
    output_path_root = r'D:\code\water_turbidity_det\label_data\test'

    # 读取标注文件
    label_path = os.path.join(input_path, 'label.txt')
    if not os.path.exists(label_path):  # 强制要求必须标注点什么
        raise FileNotFoundError(f"{label_path} 不存在")
    with open(label_path, 'r') as fr:
        lines = fr.readlines()
        lines = [line.strip() for line in lines]
    # 恢复标注的格网
    grids_info = [] # clicked_points
    for line in lines:
        point = line.split(',')
        grids_info.append([int(point[0]), int(point[1]), int(point[2]), int(point[3]), int(point[4])])
    # 我们先创建一些类别文件夹
    # 0类
    if not os.path.exists(os.path.join(output_path_root, str(0))):
        os.makedirs(os.path.join(output_path_root, str(0)))
    # 其余类
    for grid in grids_info:
        if grid[4] <= 0:
            continue
        if not os.path.exists(os.path.join(output_path_root, str(grid[4]))):
            os.makedirs(os.path.join(output_path_root, str(grid[4])))
    # 获取图像
    all_imgs = [os.path.join(input_path, i) for i in os.listdir(input_path) if i.split('.')[-1] == 'jpg' or i.split('.')[-1] == 'png']
    for img_path in all_imgs:
        img_base_name = os.path.basename(img_path).split('.')[0]
        img = cv2.imread(img_path)
        # 获取图像高宽
        img_h, img_w, _ = img.shape
        # 先将不参与训练的patch重置为0
        for g in grids_info:
            if g[4] < 0:  # 标签小于零的不参与训练
                img[g[1]:min(g[1]+g[3], img_h), g[0]:min(g[0]+g[2], img_w), :] = 0
        # 再将大于0的patch保存到对应的类别文件夹下
        for g in grids_info:
            if g[4] > 0:  # 标签大于零的放到相应的文件夹下
                patch_name = f'{img_base_name}_{g[0]}_{g[1]}_{g[4]}.jpg'  # 图块保存名称：图片名_左上角x_左上角y_类别.jpg
                patch = img[g[1]:min(g[1]+g[3], img_h), g[0]:min(g[0]+g[2], img_w), :]
                # 保存图块
                cv2.imwrite(os.path.join(output_path_root,str(g[4]), patch_name), patch)
                # 置零已经保存的patch区域
                img[g[1]:min(g[1]+g[3], img_h), g[0]:min(g[0]+g[2], img_w), :] = 0
        # 最后将剩余的patch保存到0类文件夹下
        for i in range(img_h // patch_h + 1):
            for j in range(img_w // patch_w + 1):
                patch = img[i*patch_h:min(i*patch_h+patch_h, img_h), j*patch_w:min(j*patch_w+patch_w, img_w), :]
                patch_name = f'{img_base_name}_{j*patch_w}_{i*patch_h}_0.jpg'
                # 长宽比过滤
                if patch.shape[0] / (patch.shape[1]+1e-6) > 1.314 or patch.shape[0] / (patch.shape[1]+1e-6) < 0.75:
                    #print(f"长宽比过滤: {patch_name}")
                    continue
                # 纯黑图像过滤
                if np.mean(patch) < 10.10:
                    #print(f"纯黑图像过滤: {patch_name}")
                    continue
                cv2.imwrite(os.path.join(output_path_root, '0', patch_name), patch)
                #print(f"保存图块: {patch_name}到{os.path.join(output_path_root, '0', patch_name)}")
        print(f"处理图片: {img_path}完成")
        # del patch, img
        # gc.collect()

if __name__ == '__main__':
    main()