%matplotlib inline
import torch
from d2l import torch as d2l

img = d2l.plt.imread('01_Data/03_catdog.jpg')
h, w = img.shape[:2]
print(h, w) # 打印导入图片的高、宽
print(img.shape) # 最后一个元素为通道数，通道数为3

# 探测小目标
# 调用display_anchors函数，输入参数是特征图的宽度（fmap_w=4）、高度（fmap_h=4）以及锚框的尺寸（s=[0.15]）。
# 这行代码将在特征图上生成并显示锚框，特征图的尺寸是4x4，锚框的尺寸是0.15（相对于特征图的尺寸）。
# 对于每个特征图的像素，都会生成一个以该像素为中心，尺寸为0.15的锚框。
display_anchors(fmap_w=4,fmap_h=4,s=[0.15])

# 将特征图的高度和宽度减小一半，然后使用较大的锚框来检测较大的目标
display_anchors(fmap_w=2,fmap_h=2,s=[0.4])

# 将特征图的高度和宽度较小一半，然后将锚框的尺度增加到0.8
display_anchors(fmap_w=1,fmap_h=1,s=[0.8])

# 连接多尺度的预测
# 定义一个函数forward，输入参数是输入数据x和一个网络层block。
def forward(x, block):
    # 返回通过网络层block处理输入数据x后的结果。
    return block(x)

# 使用全0的输入数据和一个类别预测层，进行前向传播。这里输入数据的形状是(2,8,20,20)，表示有2张图片，每张图片有8个通道，尺寸是20x20。
# 类别预测层的参数是8,5,10，表示输入通道数是8，锚框数是5，类别数是10。
Y1 = forward(torch.zeros((2,8,20,20)),cls_predictor(8,5,10))
# 使用另一个全0的输入数据和一个类别预测层，进行前向传播。这里输入数据的形状是(2,16,10,10)，表示有2张图片，每张图片有16个通道，尺寸是10x10。
# 类别预测层的参数是16,3,10，表示输入通道数是16，锚框数是3，类别数是10。
Y2 = forward(torch.zeros((2,16,10,10)),cls_predictor(16,3,10))
# 打印Y1的形状，这里是(2, 55, 20, 20)，表示有2张图片，每张图片有55个通道，尺寸是20x20。这里的55是因为每个锚框需要预测10+1个类别的概率，共有5个锚框，所以通道数是5*(10+1)=55。
print(Y1.shape) # 5 * (10 + 1) = 55, 55为每一个像素输出的所有锚框数对应的类别
print(Y2.shape) # 2为小批量里面有多少图片

# 定义一个函数flatten_pred，输入参数是预测结果pred。
def flatten_pred(pred):
    # pred.permute(0,2,3,1)通道数放在最后
    # start_dim=1把pred.permute(0,2,3,1)后面三个向量拉成一个向量，这样就变成一个2D矩阵
    # 2D矩阵的高为批量大小，宽为每个图片的所有
    # 如果pred.permute(0,2,3,1)不改的话，对于同一像素的类别预测flatten后会相隔较远
    # 返回拉平后的预测结果。首先使用permute方法调整通道数的位置，然后使用flatten方法将除第一维（批量大小）之外的所有维度拉平。
    return torch.flatten(pred.permute(0,2,3,1),start_dim=1) 

# 定义一个函数concat_preds，输入参数是多个预测结果的列表preds。
def concat_preds(preds):
    #[print(flatten_pred(p).shape) for p in preds]
    # 返回连接后的预测结果。首先使用列表推导式和flatten_pred函数拉平每个预测结果，然后使用torch.cat方法在dim=1的维度上连接所有的预测结果。
    return torch.cat([flatten_pred(p) for p in preds], dim=1)

# 打印连接后的预测结果的形状。这里输入的是两个预测结果Y1和Y2，输出的形状是(2, 35300)，表示有2张图片，每张图片有25300个预测结果。
print(concat_preds([Y1, Y2]).shape)

# 高和宽减半块
# 定义一个函数down_sample_blk，输入参数是输入通道数(in_channels)和输出通道数(out_channels)。
def down_sample_blk(in_channels, out_channels):
    # 创建一个空列表blk，用于存储网络层。
    blk = []
    # 循环2次，每次添加2个卷积层、1个批量归一化层和1个ReLU激活函数。
    for _ in range(2):
        # 添加一个卷积层，输入通道数是in_channels，输出通道数是out_channels，卷积核大小是3，填充大小是1。
        blk.append(nn.Conv2d(in_channels,out_channels,kernel_size=3,padding=1))     
        # 添加一个批量归一化层，输入通道数是out_channels。
        blk.append(nn.BatchNorm2d(out_channels))
        # 添加一个ReLU激活函数。
        blk.append(nn.ReLU())
        # 更新输入通道数为当前的输出通道数。
        in_channels = out_channels
    # 添加一个最大池化层，池化核大小是2，这可以将输入数据的高度和宽度减半。
    blk.append(nn.MaxPool2d(2))
    # 返回一个顺序容器，其中包含了所有的网络层。
    return nn.Sequential(*blk)

# 打印通过下采样模块处理后的输出数据的形状。这里输入的是全0的数据，形状是(2,3,20,20)，表示有2张图片，每张图片有3个通道，尺寸是20x20。
# 下采样模块的参数是3和10，表示输入通道数是3，输出通道数是10。输出的形状是(2,10,10,10)，表示有2张图片，每张图片有10个通道，尺寸是10x10。
print(forward(torch.zeros((2,3,20,20)),down_sample_blk(3,10)).shape)

# 基本网络块
# 定义一个函数base_net，没有输入参数。
def base_net(): # 从原始图片抽特征，到第一次featuremao产生锚框，中间的那一截叫base_net  
    # 创建一个空列表blk，用于存储网络层。
    blk = []
    # 定义一个列表num_filters，表示每个下采样模块的输入通道数和输出通道数。
    num_filters = [3,16,32,64] # 通道数由3到16，再到32，再到64
    # 循环遍历num_filters列表，每次添加一个下采样模块。
    for i in range(len(num_filters) - 1):
        # 添加一个下采样模块，输入通道数是num_filters[i]，输出通道数是num_filters[i+1]。
        blk.append(down_sample_blk(num_filters[i],num_filters[i+1]))
    # 返回一个顺序容器，其中包含了所有的网络层。
    return nn.Sequential(*blk)

# 打印通过基础网络模块处理后的输出数据的形状。这里输入的是全0的数据，形状是(2,3,256,256)，表示有2张图片，每张图片有3个通道，尺寸是256x256。
# 输出的形状是(2,64,32,32)，表示有2张图片，每张图片有64个通道，尺寸是32x32。
print(forward(torch.zeros((2,3,256,256)),base_net()).shape)

# 创建一个模型实例，然后使用它执行前向计算
# 创建一个 TinySSD 类的实例，类别数量为 1。
net = TinySSD(num_classes=1)
# 创建一个全零张量，形状为（32，3，256，256），代表有 32 张 3 通道的 256x256 图像。
X = torch.zeros((32,3,256,256))
# 使用模型进行前向计算，输入是创建的全零张量，输出是锚框、类别预测和边界框预测。
anchors, cls_preds, bbox_preds = net(X)
# 打印锚框的形状，每一张图片都有5444个锚框。
print('output anchors:',anchors.shape) 
# 打印类别预测的形状，32为批量大小，2为类别1加上背景为2。
print('output class preds:',cls_preds.shape) 
# 打印边界框预测的形状，每张图片有5444个锚框，每个锚框有4个坐标，所以是5444*4=21776。
print('output bbox preds:',bbox_preds.shape) 

# 读取香蕉检测数据集
# 设置批量大小为32。
batch_size = 32
# 使用 d2l 的 load_data_bananas 函数加载香蕉检测数据集，只使用训练数据集（train_iter）。
# "_" 是一个常用的占位符，表示我们不关心的变量（在这里，我们不关心测试数据集）。
train_iter, _ = d2l.load_data_bananas(batch_size)

# 训练模型
# 设置训练周期数和计时器
num_epochs, timer = 20, d2l.Timer()
# 创建一个动画对象，用于动态显示训练过程中的类别错误和边界框的平均绝对误差
animator = d2l.Animator(xlabel='epoch',xlim=[1,num_epochs],
                       legend=['class error','bbox mae'])
# 将网络移动到相应的设备（如GPU）
net = net.to(device)
# 开始训练周期
for epoch in range(num_epochs):
    # 创建一个累加器用于存储训练过程中的指标
    metric = d2l.Accumulator(4)
    # 设置网络为训练模式
    net.train()
    # 遍历训练数据
    for features, target in train_iter:
        # 计时开始
        timer.start()
        # 清零梯度
        trainer.zero_grad()
        # 将特征和目标移动到相应的设备
        X, Y = features.to(device), target.to(device)
        # 前向传播，得到锚框、类别预测和边界框预测
        anchors, cls_preds, bbox_preds = net(X)
        # 获取边界框标签、掩码和类别标签
        bbox_labels, bbox_masks, cls_labels = d2l.multibox_target(anchors, Y)   
        # 计算损失
        l = calc_loss(cls_preds, cls_labels, bbox_preds, bbox_labels, bbox_masks)  
        # 反向传播，计算梯度
        l.mean().backward()
        # 更新参数
        trainer.step()
        # 更新指标
        metric.add(cls_eval(cls_preds,cls_labels),cls_labels.numel(),
                  bbox_eval(bbox_preds,bbox_labels,bbox_masks),
                  bbox_labels.numel())
    # 计算类别错误率和边界框平均绝对误差
    cls_err, bbox_mae = 1 - metric[0] / metric[1], metric[2] / metric[3]
    # 更新动画
    animator.add(epoch + 1, (cls_err, bbox_mae))

# 打印最后的类别错误率和边界框平均绝对误差
print(f'class err {cls_err:.2e}, bbox mae {bbox_mae:.2e}')
# 打印训练速度和设备信息
print(f'{len(train_iter.dataset) / timer.stop():.1f} examples/sec on'
     f'{str(device)}')

# 筛选所有置信度不低于0.9的边界框，做为最终输出
# 定义显示边界框的函数
def display(img, output, threshold):
    # 设置图像大小
    d2l.set_figsize((5,5))
    # 展示图像
    fig = d2l.plt.imshow(img)
    # 遍历预测结果
    for row in output:
        # 获取置信度
        score = float(row[1])
        # 如果置信度低于阈值，则忽略该预测结果
        if score < threshold:
            continue
        # 获取图像的高度和宽度
        h, w = img.shape[0:2]
        # 获取预测的边界框，并将其转换为实际像素坐标
        bbox = [row[2:6] * torch.tensor((w, h, w, h), device=row.device)]
        # 显示边界框
        d2l.show_bboxes(fig.axes,bbox,'%.2f' % score,'w')

# 显示置信度不低于0.9的边界框
display(img,output.cpu(),threshold=0.9)

# 定义平滑L1损失函数
def smooth_11(data, scalar):
    out = []
    # 遍历输入数据
    for i in data:
        # 对于绝对值小于1/scalar^2的数据，采用平方损失
        if abs(i) < 1 / (scalar**2):
            out.append(((scalar * i)**2) / 2)
        # 对于绝对值大于1/scalar^2的数据，采用绝对值损失
        else:
            out.append(abs(i) - 0.5 / (scalar**2))
    # 返回平滑L1损失
    return torch.tensor(out)

# 定义不同的sigma值
sigmas = [10, 1, 0.5]
# 定义不同的线型
lines = ['-','--','-.']
# 定义x值范围
x = torch.arange(-2,2,0.1)
# 设置图像大小
d2l.set_figsize()

# 遍历不同的sigma值
for l, s in zip(lines,sigmas):
    # 计算平滑L1损失
    y = smooth_11(x,scalar=s)
    # 绘制损失函数曲线
    d2l.plt.plot(x,y,l,label='sigma=%.1f' % s)

# 显示图例
d2l.plt.legend()

# 定义焦点损失函数
def focal_loss(gamma, x):
    # 根据公式计算焦点损失
    return -(1 - x)**gamma * torch.log(x)

# 定义x值范围
x = torch.arange(0.01, 1, 0.01)
# 遍历不同的gamma值
for l, gamma in zip(lines, [0, 1, 5]):
    # 计算焦点损失
    y = d2l.plt.plot(x, focal_loss(gamma, x), l, label='gamma=%.1f' % gamma)   

# 显示图例
d2l.plt.legend()