# -*- coding: utf-8 -*-
# @Time    : 2019/8/23 21:57
# @Author  : zhoujun
from addict import Dict
from torch import nn
import torch.nn.functional as F
 
from models.backbone import build_backbone
from models.neck import build_neck
from models.head import build_head
 
 
class Model(nn.Module):
    def __init__(self, model_config: dict):
        """
        PANnet
        :param model_config: 模型配置
        """
        super().__init__()
        model_config = Dict(model_config)  # 一个配置字典，用于指定模型的结构。转换为 Dict 对象，方便通过属性访问。
        # 从配置字典中提取 backbone、neck 和 head 的类型，并分别从配置中删除它们。
        backbone_type = model_config.backbone.pop('type')
        neck_type = model_config.neck.pop('type')
        head_type = model_config.head.pop('type')
        # 使用 build_backbone、build_neck 和 build_head 函数构建模型的不同部分。
        self.backbone = build_backbone(backbone_type, **model_config.backbone)
        self.neck = build_neck(neck_type, in_channels=self.backbone.out_channels, **model_config.neck)
        self.head = build_head(head_type, in_channels=self.neck.out_channels, **model_config.head)
        # self.name 存储了模型的名字，由 backbone、neck 和 head 的类型组合而成。
        self.name = f'{backbone_type}_{neck_type}_{head_type}'
 
    def forward(self, x):
        _, _, H, W = x.size()  # x 是输入的张量，尺寸为 [batch_size, channels, height, width]。
        backbone_out = self.backbone(x)  # 使用 self.backbone 处理输入，得到 backbone_out。
        neck_out = self.neck(backbone_out)  # 使用 self.neck 处理 backbone_out，得到 neck_out。
        y = self.head(neck_out)  # 使用 self.head 处理 neck_out，得到最终输出 y。
        y = F.interpolate(y, size=(H, W), mode='bilinear', align_corners=True)  # F.interpolate 用于将输出 y 的尺寸调整为与输入 x 相同，以保持空间分辨率一致。
        return y  # 最终返回调整后的 y。
 
 
if __name__ == '__main__':
    import torch
 
    device = torch.device('cpu')  # 创建一个 CPU 设备（可以改为 GPU）。
    x = torch.zeros(2, 3, 640, 640).to(device)  # 创建一个大小为 [2, 3, 640, 640] 的零张量作为输入。
 
    model_config = {
        'backbone': {'type': 'resnest50', 'pretrained': True, "in_channels": 3},
        'neck': {'type': 'FPN', 'inner_channels': 256},  # 分割头，FPN or FPEM_FFM
        'head': {'type': 'DBHead', 'out_channels': 2, 'k': 50},
    }
    model = Model(model_config=model_config).to(device)  # 初始化 Model 对象，传入配置字典。
    import time  # 测量前向传播的时间。
 
    tic = time.time()
    y = model(x)
    # 打印出前向传播的时间、输出的形状、模型的名称和模型的详细结构。
    print(time.time() - tic)
    print(y.shape)
    print(model.name)
    print(model)
    #（可选）将模型的状态字典保存到文件 PAN.pth。
    # torch.save(model.state_dict(), 'PAN.pth')

from __future__ import print_function
 
import argparse
import os
 
import anyconfig
 
 
def init_args():
    parser = argparse.ArgumentParser(description='DBNet.pytorch')
    parser.add_argument('--config_file', default='config/open_dataset_resnet18_FPN_DBhead_polyLR.yaml', type=str)
    parser.add_argument('--local_rank', dest='local_rank', default=0, type=int, help='Use distributed training')
 
    args = parser.parse_args()
    return args
 
 
def main(config):
    # 导入必要的模块，包括模型构建、损失函数、数据加载器、训练器、后处理和评估指标。
    import torch
    from models import build_model, build_loss
    from data_loader import get_dataloader
    from trainer import Trainer
    from post_processing import get_post_processing
    from utils import get_metric
    # 检查是否有多个 GPU，如果有，则初始化分布式训练环境。
    if torch.cuda.device_count() > 1:
        torch.cuda.set_device(args.local_rank)
        torch.distributed.init_process_group(backend="nccl", init_method="env://", world_size=torch.cuda.device_count(), rank=args.local_rank)
        config['distributed'] = True
    else:
        config['distributed'] = False
    config['local_rank'] = args.local_rank
 
    # 根据配置文件加载训练和验证数据加载器。
    train_loader = get_dataloader(config['dataset']['train'], config['distributed'])
    assert train_loader is not None
    if 'validate' in config['dataset']:
        validate_loader = get_dataloader(config['dataset']['validate'], False)
    else:
        validate_loader = None
 
    # 构建损失函数并将其移动到 GPU。
    criterion = build_loss(config['loss']).cuda()
 
    # 配置模型的输入通道数（彩色图像为 3 通道，灰度图像为 1 通道）。
    config['arch']['backbone']['in_channels'] = 3 if config['dataset']['train']['dataset']['args']['img_mode'] != 'GRAY' else 1
    model = build_model(config['arch'])
 
    # 构建模型、后处理函数和评估指标。
    post_p = get_post_processing(config['post_processing'])
    metric = get_metric(config['metric'])
 
    # 创建 Trainer 对象并开始训练。
    trainer = Trainer(config=config,
                      model=model,
                      criterion=criterion,
                      train_loader=train_loader,
                      post_process=post_p,
                      metric_cls=metric,
                      validate_loader=validate_loader)
    trainer.train()
 
 
if __name__ == '__main__':
    # 处理模块路径，确保当前目录和上级目录在 Python 路径中。
    import sys
    import pathlib
    __dir__ = pathlib.Path(os.path.abspath(__file__))
    sys.path.append(str(__dir__))
    sys.path.append(str(__dir__.parent.parent))
    # project = 'DBNet.pytorch'  # 工作项目根目录
    # sys.path.append(os.getcwd().split(project)[0] + project)
 
    # 使用 anyconfig 读取配置文件，并解析基础配置。
    from utils import parse_config
 
    args = init_args()
    assert os.path.exists(args.config_file)
    config = anyconfig.load(open(args.config_file, 'rb'))
    if 'base' in config:
        config = parse_config(config)
    # 调用 main 函数开始模型训练。
    main(config)