#CRNN详解：端到端不定长文字识别模型

import torch
import torch.nn as nn

class VGG_FeatureExtractor(nn.Module):
    """
    经典的CRNN后端CNN特征提取器(裁剪版VGG)
    输入高度必须固定为32，宽度可以是可变的。
    """
    def __init__(self, input_channel=1, output_channel=512):
        super(VGG_FeatureExtractor, self).__init__()
        self.output_channel = output_channel
        
        # 定义核心卷积网络
        self.cnn = nn.Sequential(
            # Conv Block 1: 32xW -> 16xW/2
            nn.Conv2d(input_channel, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(True),
            nn.MaxPool2d(kernel_size=2, stride=2), # Output: 64 x 16 x W/2

            # Conv Block 2: 16xW/2 -> 8xW/4
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(True),
            nn.MaxPool2d(kernel_size=2, stride=2), # Output: 128 x 8 x W/4

            # Conv Block 3: 8xW/4 -> 4xW/4
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256), # BN在特征提取中很有效
            nn.ReLU(True),
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(True),
            # 关键：MaxPool在高减半，但在宽方向stride为1，不减半
            nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 1), padding=(0, 1)), # Output: 256 x 4 x (W/4+2)

            # Conv Block 4: 4xW/4 -> 1xW/4
            nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(True),
            # 关键：MaxPool高减半为1
            nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 1), padding=(0, 0)), # Output: 512 x 1 x Width_seq

            # Conv Block 5: 进一步融合特征
            nn.Conv2d(512, 512, kernel_size=2, stride=1, padding=0), # 高、宽都为2的卷积，进一步降维
            nn.ReLU(True)
            # 最终输出通道数：512
        )

    def forward(self, x):
        # 输入 x: (Batch, Input_channel, 32, W)
        conv = self.cnn(x)
        return conv

class BidirectionalLSTM(nn.Module):
    """
    单层双向LSTM
    """
    def __init__(self, input_size, hidden_size, output_size):
        super(BidirectionalLSTM, self).__init__()
        # PyTorch的LSTM: 输入 (Time, Batch, Input_size)
        self.rnn = nn.LSTM(input_size, hidden_size, bidirectional=True)
        # 双向输出是隐层大小的两倍，需要通过线性层降维到output_size
        self.embedding = nn.Linear(hidden_size * 2, output_size)

    def forward(self, x):
        # 输入 x: (Time, Batch, Input_size)
        recurrent, _ = self.rnn(x) # recurrent: (Time, Batch, Hidden*2)
        
        # 融合时间维和Batch维进行Linear处理
        T, B, H2 = recurrent.size()
        t_rec = recurrent.view(T * B, H2)
        
        output = self.embedding(t_rec) # (Time*Batch, Output_size)
        output = output.view(T, B, -1) # 复原 (Time, Batch, Output_size)
        
        return output

class CRNN(nn.Module):
    def __init__(self, img_h, nc, nclass, nh):
        """
        img_h: 输入图像高度(应为32)
        nc: 输入图像通道数(1 for gray, 3 for rgb)
        nclass: 字符类别数(必须包含blank标记，通常blank的index是0)
        nh: RNN的隐层神经元数量
        """
        super(CRNN, self).__init__()
        assert img_h % 16 == 0, 'img_h has to be a multiple of 16'
        
        # 1. CNN特征提取层
        self.cnn = VGG_FeatureExtractor(nc, 512)
        
        # 2. RNN序列建模层(两层BiLSTM)
        # 第一层: 将CNN通道数512映射到RNN隐层nh
        self.rnn1 = BidirectionalLSTM(512, nh, nh)
        # 第二层: 将第一层输出进一步建模，并映射到最终的字符类别数nclass
        self.rnn2 = BidirectionalLSTM(nh, nh, nclass)

    def forward(self, x):
        # 1. CNN层特征提取
        # Input x: (Batch, nc, img_h, W)
        conv = self.cnn(x)
        
        # 2. 特征序列化(Map-to-Sequence)
        # conv shape: (Batch, 512, 1, Width_seq)
        b, c, h, w = conv.size()
        assert h == 1, "The height of conv feature map must be 1"
        
        # Remove height dim: (B, C, 1, W_seq) -> (B, C, W_seq)
        conv = conv.squeeze(2)
        # Permute for RNN: (B, C, W_seq) -> (W_seq, B, C)
        conv = conv.permute(2, 0, 1) # (Time_steps, Batch, Input_size)
        
        # 3. RNN层序列预测
        rnn_out = self.rnn1(conv)
        rnn_out = self.rnn2(rnn_out)
        
        # 最终输出形状: (Time_steps, Batch, nclass)
        # 这也是PyTorch nn.CTCLoss需要的标准输入格式
        return rnn_out

# --- 测试模型输入输出 ---
if __name__ == "__main__":
    # 参数设置
    batch_size = 1
    input_channels = 1 # 灰度图
    img_h = 32
    img_w = 100        # 可变宽度
    nh = 256           # RNN隐层神经元数
    # nclass包括：空白符 + 字符集(如'a'-'z', '0'-'9')
    # 假设blank是0, 字母是1-26, 共27类
    nclass = 27 
    
    # 实例化模型
    model = CRNN(img_h, input_channels, nclass, nh)
    
    # 检查网络结构
    # print(model)

    # 模拟输入：1张灰度图, 32x100
    dummy_input = torch.randn(batch_size, input_channels, img_h, img_w)
    
    # 模型前向传播
    output = model(dummy_input)
    
    print(f"\n模型测试结果:")
    print(f"输入形状 (Batch, C, H, W): {dummy_input.shape}")
    print(f"输出形状 (Time_steps, Batch, nclass): {output.shape}")
    
    # 根据32x100的输入，CNN最终输出宽序列大约是24-26
    # 输出的时间步长应该大约是这个数
    time_steps, _, _ = output.shape
    assert time_steps > 0, "Error: Time_steps is 0"
    print("模型输出测试成功！")

import torch.nn.functional as F

# 1. 实例化损失函数
# blank=0表示在nclass中blank标记的索引是0
criterion = nn.CTCLoss(blank=0, reduction='mean')

# 2. 模型前向传播
# 假设input数据大小是(Batch, 1, 32, W)
model_output = model(images) # Shape: (Time_steps, B, nclass)

# 3. 计算Log-Probabilities
# CTC需要输入是对数概率
log_probs = F.log_softmax(model_output, dim=2)

# 4. 准备CTC Loss需要的参数
# T是序列的长度(Time_steps)，也就是RNN输出的第一维
input_lengths = torch.full((batch_size,), time_steps, dtype=torch.long)
# target_lengths是每张图中真实的标签长度(例如"hello"是5)
target_lengths = torch.tensor([len(t) for t in labels_encoded], dtype=torch.long)
# targets是将真实文本转为数字的序列，并平铺成一维

# 5. 计算损失
loss = criterion(log_probs, targets, input_lengths, target_lengths)

# 6. 反向传播与更新
loss.backward()
optimizer.step()