ocr-training/scripts/configure_training.py

import os
import sys
import urllib.request
import tarfile
import yaml

if sys.stdout and sys.stdout.encoding != 'utf-8':
    try:
        sys.stdout.reconfigure(encoding='utf-8')
    except AttributeError:
        pass

MODEL_URL = "https://paddleocr.bj.bcebos.com/PP-OCRv4/chinese/ch_PP-OCRv4_rec_train.tar"
PRETRAIN_DIR_NAME = "ch_PP-OCRv4_rec_train"

def download_progress_hook(block_num, block_size, total_size):
    percent = int(block_num * block_size * 100 / total_size)
    percent = min(100, percent)
    sys.stdout.write(f"\r📥 Đang tải base model: {percent}% [{block_num * block_size}/{total_size} bytes]")
    sys.stdout.flush()

def download_and_extract_base_model(target_dir):
    os.makedirs(target_dir, exist_ok=True)
    tar_path = os.path.join(target_dir, "ch_PP-OCRv4_rec_train.tar")
    extracted_model_path = os.path.join(target_dir, PRETRAIN_DIR_NAME)
    
    if os.path.exists(tar_path) and os.path.getsize(tar_path) == 0:
        print("⚠️  Phát hiện tệp tar tải dở dang (0 bytes). Đang xóa để tải lại...")
        os.remove(tar_path)

    if not os.path.exists(tar_path) and not os.path.exists(extracted_model_path):
        print(f"📥 Bắt đầu tải base model PP-OCRv4 từ: {MODEL_URL}")
        try:
            urllib.request.urlretrieve(MODEL_URL, tar_path, reporthook=download_progress_hook)
            print("\n✅ Tải hoàn tất.")
        except Exception as e:
            print(f"\n❌ Lỗi khi tải mô hình: {e}")
            if os.path.exists(tar_path):
                os.remove(tar_path)
            return None
    else:
        print("ℹ️  Tệp mô hình base đã tồn tại hoặc đã được giải nén. Bỏ qua tải xuống.")

    if os.path.exists(tar_path) and not os.path.exists(extracted_model_path):
        print(f"📦 Bắt đầu giải nén tệp {tar_path}...")
        try:
            with tarfile.open(tar_path, "r") as tar:
                tar.extractall(path=target_dir)
            print("✅ Giải nén hoàn tất.")
            os.remove(tar_path)
            print("🗑️  Đã dọn dẹp tệp nén .tar.")
        except Exception as e:
            print(f"❌ Lỗi khi giải nén mô hình: {e}")
            if os.path.exists(tar_path):
                os.remove(tar_path)
            return None
            
    student_model_prefix = os.path.join(extracted_model_path, "best_accuracy")
    print(f"🎯 Đường dẫn base model học sinh (student model): {student_model_prefix}")
    return student_model_prefix

def configure_training(template_path, output_config_path, data_dir, dict_path, pretrained_model_prefix, 
                       is_kaggle=None, max_text_length=80, image_shape=[3, 48, 1024], batch_size=64, 
                       model_head_type='MultiHead', epoch_num=200, save_model_dir=None, train_label_name='train_label.txt'):
    """
    Cấu hình YAML cho V25 với các tham số độ dài tối đa, kích thước ảnh, batch size, loại model head.
    """
    if is_kaggle is None:
        is_kaggle = os.path.exists('/kaggle/working')

    if not os.path.exists(template_path):
        print(f"❌ Không tìm thấy tệp cấu hình mẫu tại: {template_path}")
        return False
        
    print(f"🛠️  Đang đọc tệp cấu hình mẫu từ: {template_path}...")
    with open(template_path, 'r', encoding='utf-8') as f:
        config = yaml.safe_load(f)

    # 1. Global config
    config['Global'] = config.get('Global', {})
    
    # Auto GPU check
    has_gpu = False
    try:
        import paddle
        has_gpu = paddle.is_compiled_with_cuda() and paddle.device.cuda.device_count() > 0
    except ImportError:
        has_gpu = True
            
    config['Global']['use_gpu'] = has_gpu
    config['Global']['pretrained_model'] = pretrained_model_prefix.replace('\\', '/')

    if save_model_dir is None:
        output_dir = os.path.join(data_dir, 'output/rec_cham').replace('\\', '/')
    else:
        output_dir = save_model_dir.replace('\\', '/')
    
    config['Global']['save_res_path'] = output_dir
    config['Global']['save_model_dir'] = output_dir
    config['Global']['character_dict_path'] = dict_path.replace('\\', '/')
    config['Global']['epoch_num'] = epoch_num
    config['Global']['save_epoch_step'] = min(10, epoch_num)
    config['Global']['eval_batch_step'] = [0, 1000] if not is_kaggle else [0, 2000]
    config['Global']['use_space_char'] = True
    config['Global']['max_text_length'] = max_text_length
    config['Global']['d2s_train_image_shape'] = image_shape

    # 2. Config datasets and loaders
    for split in ['Train', 'Eval']:
        if split in config:
            dataset = config[split].get('dataset', {})
            dataset['data_dir'] = os.path.join(data_dir, 'cham_synthetic_images').replace('\\', '/')
            
            # Default label list
            label_file = train_label_name if split == 'Train' else 'val_clean_short_label.txt'
            dataset['label_file_list'] = [os.path.join(data_dir, f'cham_synthetic_images/{label_file}').replace('\\', '/')]
            
            config[split]['dataset'] = dataset
            
            # Loader
            loader = config[split].get('loader', {})
            loader['batch_size_per_card'] = batch_size
            loader['num_workers'] = 2
            config[split]['loader'] = loader

            # Train sampler
            if split == 'Train' and 'sampler' in config[split]:
                sampler = config[split]['sampler']
                if 'first_bs' in sampler:
                    sampler['first_bs'] = batch_size
                # Match scales to the target image size
                img_c, img_h, img_w = image_shape
                sampler['scales'] = [[img_w, img_h - 16], [img_w, img_h], [img_w, img_h + 16]]

    # 3. Model Head & Loss modification
    img_c, img_h, img_w = image_shape
    
    if model_head_type == 'CTCHead':
        # CTC-only setup
        config['Architecture']['Head'] = {
            'name': 'CTCHead',
            'Neck': {
                'name': 'svtr',
                'dims': 120,
                'depth': 2,
                'hidden_dims': 120,
                'kernel_size': [1, 3],
                'use_guide': True
            },
            'Head': {
                'fc_decay': 0.00001
            }
        }
        config['Loss'] = {
            'name': 'CTCLoss'
        }
        config['PostProcess'] = {
            'name': 'CTCLabelDecode'
        }
        
        # Modify transforms for Train
        train_transforms = []
        for t in config['Train']['dataset'].get('transforms', []):
            t_name = list(t.keys())[0]
            if t_name == 'RecConAug':
                t['RecConAug']['image_shape'] = [img_h, img_w, img_c]
                t['RecConAug']['max_text_length'] = max_text_length
                train_transforms.append(t)
            elif t_name == 'MultiLabelEncode':
                # Replace with CTCLabelEncode
                train_transforms.append({'CTCLabelEncode': None})
            elif t_name == 'KeepKeys':
                train_transforms.append({'KeepKeys': {'keep_keys': ['image', 'label', 'length']}})
            else:
                train_transforms.append(t)
        config['Train']['dataset']['transforms'] = train_transforms
        
        # Modify transforms for Eval
        eval_transforms = []
        for t in config['Eval']['dataset'].get('transforms', []):
            t_name = list(t.keys())[0]
            if t_name == 'MultiLabelEncode':
                eval_transforms.append({'CTCLabelEncode': None})
            elif t_name == 'RecResizeImg':
                t['RecResizeImg']['image_shape'] = [img_c, img_h, img_w]
                eval_transforms.append(t)
            elif t_name == 'KeepKeys':
                eval_transforms.append({'KeepKeys': {'keep_keys': ['image', 'label', 'length']}})
            else:
                eval_transforms.append(t)
        config['Eval']['dataset']['transforms'] = eval_transforms

    else:
        # MultiHead (CTC + NRTR) setup
        config['Architecture']['Head'] = {
            'name': 'MultiHead',
            'head_list': [
                {
                    'CTCHead': {
                        'Neck': {
                            'name': 'svtr',
                            'dims': 120,
                            'depth': 2,
                            'hidden_dims': 120,
                            'kernel_size': [1, 3],
                            'use_guide': True
                        },
                        'Head': {
                            'fc_decay': 0.00001
                        }
                    }
                },
                {
                    'NRTRHead': {
                        'nrtr_dim': 384,
                        'max_text_length': max_text_length
                    }
                }
            ]
        }
        config['Loss'] = {
            'name': 'MultiLoss',
            'loss_config_list': [
                {'CTCLoss': None},
                {'NRTRLoss': None}
            ]
        }
        config['PostProcess'] = {
            'name': 'CTCLabelDecode'
        }
        
        # Modify transforms for Train
        train_transforms = []
        for t in config['Train']['dataset'].get('transforms', []):
            t_name = list(t.keys())[0]
            if t_name == 'RecConAug':
                t['RecConAug']['image_shape'] = [img_h, img_w, img_c]
                t['RecConAug']['max_text_length'] = max_text_length
                train_transforms.append(t)
            elif t_name == 'MultiLabelEncode':
                t['MultiLabelEncode']['gtc_encode'] = 'NRTRLabelEncode'
                train_transforms.append(t)
            elif t_name == 'KeepKeys':
                train_transforms.append({'KeepKeys': {'keep_keys': ['image', 'label_ctc', 'label_gtc', 'length', 'valid_ratio']}})
            else:
                train_transforms.append(t)
        config['Train']['dataset']['transforms'] = train_transforms
        
        # Modify transforms for Eval
        eval_transforms = []
        for t in config['Eval']['dataset'].get('transforms', []):
            t_name = list(t.keys())[0]
            if t_name == 'MultiLabelEncode':
                t['MultiLabelEncode']['gtc_encode'] = 'NRTRLabelEncode'
                eval_transforms.append(t)
            elif t_name == 'RecResizeImg':
                t['RecResizeImg']['image_shape'] = [img_c, img_h, img_w]
                eval_transforms.append(t)
            elif t_name == 'KeepKeys':
                eval_transforms.append({'KeepKeys': {'keep_keys': ['image', 'label_ctc', 'label_gtc', 'length', 'valid_ratio']}})
            else:
                eval_transforms.append(t)
        config['Eval']['dataset']['transforms'] = eval_transforms

    # Save out
    os.makedirs(os.path.dirname(output_config_path), exist_ok=True)
    with open(output_config_path, 'w', encoding='utf-8') as f:
        yaml.safe_dump(config, f, allow_unicode=True, default_flow_style=False, sort_keys=False)
        
    print(f"✅ Đã tạo cấu hình V25 ({model_head_type}, {image_shape}, BS={batch_size}) tại: {output_config_path}")
    return True

if __name__ == '__main__':
    is_kaggle = os.path.exists('/kaggle/working')
    
    if is_kaggle:
        PROJECT_ROOT = "/kaggle/working/paddleocr_cham_finetune"
        paddleocr_root = "/kaggle/working/PaddleOCR"
    else:
        PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        paddleocr_root = os.path.join(PROJECT_ROOT, "PaddleOCR")
    
    config_template = os.path.join(paddleocr_root, "configs/rec/PP-OCRv4/PP-OCRv4_mobile_rec.yml")
    output_config = os.path.join(PROJECT_ROOT, "configs", "rec_cham_config.yml")
    data_root = os.path.join(PROJECT_ROOT, "data")
    dict_file = os.path.join(data_root, "cham_dict.txt")
    
    pretrain_target_dir = os.path.join(data_root, "pretrain_models")
    model_prefix = download_and_extract_base_model(pretrain_target_dir)
    
    if model_prefix and os.path.exists(config_template):
        configure_training(config_template, output_config, data_root, dict_file, model_prefix, is_kaggle)