migrate_annotations.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
标注文件迁移脚本

将现有的标注文件从平铺结构迁移到按场景分组的结构：
- 原结构: annotations/episode_id.json
- 新结构: annotations/scene_X/episode_id.json

作者: Lozumi
日期: 2025-07-22
"""

import os
import json
import shutil
from datetime import datetime
from typing import Dict, List

def load_annotation_file(filepath: str) -> Dict:
    """加载标注文件"""
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            return json.load(f)
    except Exception as e:
        print(f"加载文件失败 {filepath}: {e}")
        return None

def migrate_annotations():
    """迁移标注文件到按场景分组的结构"""
    annotations_dir = "annotations"
    
    if not os.path.exists(annotations_dir):
        print("annotations 目录不存在")
        return
    
    # 创建备份目录
    backup_time = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup_dir = os.path.join(annotations_dir, f"migration_backup_{backup_time}")
    
    migrated_files = []
    skipped_files = []
    error_files = []
    
    print("开始迁移标注文件...")
    print(f"备份目录: {backup_dir}")
    
    # 遍历annotations目录下的所有.json文件
    for filename in os.listdir(annotations_dir):
        if not filename.endswith('.json'):
            continue
            
        # 跳过特殊文件
        if filename in ['progress.json']:
            continue
            
        filepath = os.path.join(annotations_dir, filename)
        
        # 跳过已经是目录的项目
        if os.path.isdir(filepath):
            continue
            
        print(f"处理文件: {filename}")
        
        # 加载标注数据
        annotation_data = load_annotation_file(filepath)
        if annotation_data is None:
            error_files.append(filename)
            continue
            
        # 获取场景ID
        scene_id = annotation_data.get('scene_id')
        if scene_id is None:
            print(f"  警告: 文件 {filename} 缺少 scene_id，跳过")
            skipped_files.append(filename)
            continue
            
        # 创建场景目录
        scene_dir = os.path.join(annotations_dir, f"scene_{scene_id}")
        os.makedirs(scene_dir, exist_ok=True)
        
        # 目标文件路径
        target_filepath = os.path.join(scene_dir, filename)
        
        # 如果目标文件已存在，跳过
        if os.path.exists(target_filepath):
            print(f"  跳过: 目标文件已存在 {target_filepath}")
            skipped_files.append(filename)
            continue
            
        try:
            # 创建备份（首次创建备份目录）
            if not os.path.exists(backup_dir):
                os.makedirs(backup_dir, exist_ok=True)
            backup_filepath = os.path.join(backup_dir, filename)
            shutil.copy2(filepath, backup_filepath)
            
            # 移动文件到场景目录
            shutil.move(filepath, target_filepath)
            migrated_files.append((filename, f"scene_{scene_id}"))
            print(f"  成功: {filename} -> scene_{scene_id}/")
            
        except Exception as e:
            print(f"  错误: 移动文件失败 {filename}: {e}")
            error_files.append(filename)
    
    # 输出迁移结果
    print("\n" + "="*50)
    print("迁移完成！")
    print(f"成功迁移: {len(migrated_files)} 个文件")
    print(f"跳过文件: {len(skipped_files)} 个文件")
    print(f"错误文件: {len(error_files)} 个文件")
    
    if migrated_files:
        print("\n成功迁移的文件:")
        for filename, scene_dir in migrated_files:
            print(f"  {filename} -> {scene_dir}/")
    
    if skipped_files:
        print(f"\n跳过的文件: {skipped_files}")
    
    if error_files:
        print(f"\n错误的文件: {error_files}")
    
    # 显示新的目录结构
    print(f"\n新的目录结构:")
    scene_dirs = [d for d in os.listdir(annotations_dir) 
                  if os.path.isdir(os.path.join(annotations_dir, d)) and d.startswith('scene_')]
    scene_dirs.sort()
    
    for scene_dir in scene_dirs:
        scene_path = os.path.join(annotations_dir, scene_dir)
        files = [f for f in os.listdir(scene_path) if f.endswith('.json')]
        print(f"  {scene_dir}/: {len(files)} 个文件")

def verify_migration():
    """验证迁移结果"""
    annotations_dir = "annotations"
    
    print("验证迁移结果...")
    
    # 检查根目录是否还有.json文件（除了progress.json）
    root_json_files = []
    for filename in os.listdir(annotations_dir):
        if filename.endswith('.json') and filename != 'progress.json':
            filepath = os.path.join(annotations_dir, filename)
            if not os.path.isdir(filepath):
                root_json_files.append(filename)
    
    if root_json_files:
        print(f"警告: 根目录仍有 {len(root_json_files)} 个 .json 文件未迁移:")
        for filename in root_json_files:
            print(f"  {filename}")
    else:
        print("✓ 根目录清理完成，所有 .json 文件已迁移")
    
    # 统计场景目录
    scene_dirs = [d for d in os.listdir(annotations_dir) 
                  if os.path.isdir(os.path.join(annotations_dir, d)) and d.startswith('scene_')]
    
    total_files = 0
    for scene_dir in scene_dirs:
        scene_path = os.path.join(annotations_dir, scene_dir)
        files = [f for f in os.listdir(scene_path) if f.endswith('.json')]
        total_files += len(files)
        print(f"  {scene_dir}: {len(files)} 个文件")
    
    print(f"\n总计: {len(scene_dirs)} 个场景目录，{total_files} 个标注文件")

if __name__ == "__main__":
    print("标注文件迁移脚本")
    print("=" * 50)
    
    # 确认操作
    response = input("确认要将标注文件迁移到按场景分组的结构吗？(y/N): ")
    if response.lower() != 'y':
        print("操作已取消")
        exit(0)
    
    # 执行迁移
    migrate_annotations()
    
    # 验证结果
    print("\n" + "="*50)
    verify_migration()
    
    print("\n迁移完成！新的标注保存功能将按场景创建子文件夹。")