import os
import shutil
from ultralytics import YOLO

# ================= 配置区 =================
DIR_NORMAL = 'datasets/raw_normal'
DIR_PROHIBITED = 'datasets/raw_prohibited'
IMG_TARGET = 'datasets/knife_data/images/train'
LBL_TARGET = 'datasets/knife_data/labels/train'
VERIFY_DIR = 'runs/verify_labels'

# 【关键优化 1】：换用中型模型 (yolov8m.pt)，识别能力比 n 强很多
# 如果你的 Mac 内存足够（8G以上），可以直接写 yolov8x.pt 效果最好
model = YOLO('yolov8m.pt') 
# ==========================================

def safe_makedirs(path):
    if not os.path.exists(path):
        os.makedirs(path, exist_ok=True)

def process_folder(src_dir, target_class_id, prefix):
    if not os.path.exists(src_dir) or not any(f.lower().endswith(('.jpg', '.png', '.jpeg')) for f in os.listdir(src_dir)):
        return

    print(f"\n🚀 正在对【{prefix}】进行深度扫描，ID -> {target_class_id}")
    
    # 【关键优化 2】：调低 conf (置信度阈值)
    # 0.1 代表只要有 10% 的把握是刀，就把它标出来。
    results = model.predict(
        source=src_dir,
        conf=0.1, 
        save_txt=True,
        save=True,
        classes=[43], 
        project='runs/detect',
        name=f'tmp_{prefix}',
        exist_ok=True
    )

    actual_save_dir = results[0].save_dir
    tmp_label_dir = os.path.join(actual_save_dir, 'labels')

    # 1. 搬运并修正标签
    label_count = 0
    if os.path.exists(tmp_label_dir):
        for f in os.listdir(tmp_label_dir):
            if f.endswith('.txt'):
                src_path = os.path.join(tmp_label_dir, f)
                dst_path = os.path.join(LBL_TARGET, f)
                
                with open(src_path, 'r') as file:
                    lines = file.readlines()
                
                corrected_lines = [f"{target_class_id} " + " ".join(l.split()[1:]) + "\n" for l in lines]
                
                with open(dst_path, 'w') as file:
                    file.writelines(corrected_lines)
                label_count += 1
    
    # 2. 同步图片
    img_count = 0
    for f in os.listdir(src_dir):
        if f.lower().endswith(('.jpg', '.jpeg', '.png')):
            shutil.copy(os.path.join(src_dir, f), os.path.join(IMG_TARGET, f))
            img_count += 1
            
    # 3. 整理验证图
    for f in os.listdir(actual_save_dir):
        if f.lower().endswith(('.jpg', '.jpeg', '.png')):
            shutil.move(os.path.join(actual_save_dir, f), os.path.join(VERIFY_DIR, f"{prefix}_{f}"))
    
    print(f"   ✅ 完成：图片 {img_count} 张，AI 捕获标注 {label_count} 份")

def run_workflow():
    # 初始化环境
    for d in [IMG_TARGET, LBL_TARGET, VERIFY_DIR]:
        if os.path.exists(d): shutil.rmtree(d)
        safe_makedirs(d)
    
    process_folder(DIR_NORMAL, 0, "normal")
    process_folder(DIR_PROHIBITED, 1, "prohibited")

    print("\n" + "★"*40)
    print(f"✅ 自动标注执行完毕")
    print(f"📊 识别统计：共生成 {len(os.listdir(LBL_TARGET))} 个标签文件")
    print(f"⚠️  警告：如果标签文件少于图片总数，说明 AI 还是没认全。")
    print(f"🛠️  下一步建议：使用 LabelImg 打开 {IMG_TARGET} 手动补齐剩下的框。")
    print("★"*40)

if __name__ == "__main__":
    run_workflow()