#PointPillars检测模型训练(无 config)

class PointPillarsModel:

    task_name = "pointpillars_kitti_car" # 任务名（任意）

    @classmethod
    def model(cls):
        """模型结构"""
        model = cls._build_pp_model(cls, is_deploy=False)
        return model

    @classmethod
    def deploy_model(cls):
        """deploy model, 用于编译"""
        deploy_model = cls._build_pp_model(cls, is_deploy=True)
        return deploy_model

    @classmethod
    def deploy_inputs(cls):
        """deploy inputs, 用于编译"""
        deploy_inputs = dict(  # noqa C408
            points=[
                torch.randn(150000, 4),
            ],
        )
        return deploy_inputs

    @classmethod
    def float_model(cls, pretrain_ckpt=None):
        """浮点模型，加载参数"""
        model = cls.model()
        if pretrain_ckpt:
            ckpt_loader = LoadCheckpoint(
                pretrain_ckpt,
                allow_miss=False,
                ignore_extra=False,
                verbose=True,
            )
            model = ckpt_loader(model)

        return model

    @classmethod
    def qat_model(cls, pre_step_ckpt=None, pretrain_ckpt=None):
        """QAT 模型"""

        float_model = cls.float_model(pre_step_ckpt)
        qat_model = Float2QAT()(float_model)

        if pretrain_ckpt:
            ckpt_loader = LoadCheckpoint(
                pretrain_ckpt,
                allow_miss=False,
                ignore_extra=False,
                verbose=True,
            )
            qat_model = ckpt_loader(qat_model)

        return qat_model

    @classmethod
    def int_infer_model(cls, pre_step_ckpt=None, pretrain_ckpt=None, use_deploy=False):
        if use_deploy:  # for training
            model = cls.deploy_model()  # float model
        else:           # for prediction
            model = cls.model()
        model = Float2QAT()(model)  # qat model

        if pre_step_ckpt:  # load QAT checkpoint
            qat_ckpt_loader = LoadCheckpoint(
                pre_step_ckpt,
                allow_miss=False,
                ignore_extra=False,
                verbose=True,
            )
            model = qat_ckpt_loader(model)  # qat model with state_dict

        int_model = QAT2Quantize()(model)

        if pretrain_ckpt:
            int_ckpt_loader = LoadCheckpoint(
                pretrain_ckpt,
                allow_miss=False,
                 ignore_extra=False,
                verbose=True,
            )
            int_model = int_ckpt_loader(model)

        return int_model

    def _build_pp_model(self, is_deploy=False):
        """模型结构实现"""
        # Voxelization cfg
        pc_range = [0, -39.68, -3, 69.12, 39.68, 1]
        voxel_size = [0.16, 0.16, 4.0]
        max_points_in_voxel = 100
        max_voxels_num = 12000
        class_names = ["Car"]

        def get_feature_map_size(point_cloud_range, voxel_size):
            point_cloud_range = np.array(point_cloud_range, dtype=np.float32)
            voxel_size = np.array(voxel_size, dtype=np.float32)
            grid_size = (
                point_cloud_range[3:] - point_cloud_range[:3]
            ) / voxel_size
            grid_size = np.round(grid_size).astype(np.int64)
            return grid_size

        model = PointPillarsDetector(
            feature_map_shape=get_feature_map_size(pc_range, voxel_size),
            is_deploy=is_deploy,
            pre_process=PointPillarsPreProcess(
                pc_range=pc_range,
                voxel_size=voxel_size,
                max_voxels_num=max_voxels_num,
                max_points_in_voxel=max_points_in_voxel,
            ),
            reader=PillarFeatureNet(
                num_input_features=4,
                num_filters=(64,),
                with_distance=False,
                pool_size=(1, max_points_in_voxel),
                voxel_size=voxel_size,
                pc_range=pc_range,
                bn_kwargs=None,
                quantize=True,
                use_4dim=True,
                use_conv=True,
            ),
            backbone=PointPillarScatter(
                num_input_features=64,
                use_horizon_pillar_scatter=True,
                quantize=True,
            ),
            neck=SECONDNeck(
                in_feature_channel=64,
                down_layer_nums=[3, 5, 5],
                down_layer_strides=[2, 2, 2],
                down_layer_channels=[64, 128, 256],
                up_layer_strides=[1, 2, 4],
                up_layer_channels=[128, 128, 128],
                bn_kwargs=None,
                quantize=True,
            ),
            head=PointPillarsHead(
                num_classes=len(class_names),
                in_channels=sum([128, 128, 128]),
                use_direction_classifier=True,
            ),
            anchor_generator=Anchor3DGeneratorStride(
                anchor_sizes=[[1.6, 3.9, 1.56]],  # noqa B006
                anchor_strides=[[0.32, 0.32, 0.0]],  # noqa B006
                anchor_offsets=[[0.16, -39.52, -1.78]],  # noqa B006
                rotations=[[0, 1.57]],  # noqa B006
                class_names=class_names,
                match_thresholds=[0.6],
                unmatch_thresholds=[0.45],
            ),
            targets=LidarTargetAssigner(
                box_coder=GroundBox3dCoder(n_dim=7),
                class_names=class_names,
                positive_fraction=-1,
            ),
            loss=PointPillarsLoss(
                num_classes=len(class_names),
                loss_cls=FocalLossV2(
                    alpha=0.25,
                    gamma=2.0,
                    from_logits=False,
                    reduction="none",
                    loss_weight=1.0,
                ),
                loss_bbox=SmoothL1Loss(
                    beta=1 / 9.0,
                    reduction="none",
                    loss_weight=2.0,
                ),
                loss_dir=CrossEntropyLoss(
                    use_sigmoid=False,
                    reduction="none",
                    loss_weight=0.2,
                ),
            ),
            postprocess=PointPillarsPostProcess(
                num_classes=len(class_names),
                box_coder=GroundBox3dCoder(n_dim=7),
                use_direction_classifier=True,
                num_direction_bins=2,
                # test_cfg
                use_rotate_nms=False,
                nms_pre_max_size=1000,
                nms_post_max_size=300,
                nms_iou_threshold=0.5,
                score_threshold=0.4,
                post_center_limit_range=[0, -39.68, -5, 69.12, 39.68, 5],
                max_per_img=100,
            ),
        )

        if is_deploy:
            model.anchor_generator = None
            model.targets = None
            model.loss = None
            model.postprocess = None

        return model

class DataHelper:

    data_dir = "./tmp_data/kitti3d"     # 数据根目录
    train_batch_size = 2                # 训练 batch_size
    val_batch_size = 1                  # 评测 batch_size

    @classmethod
    def train_data_loader(cls):
        """train dataloader 用于训练"""
        return cls.build_dataloader(cls, is_training=True)

    @classmethod
    def val_data_loader(cls):
        """val dataloader 用于评测验证"""
        return cls.build_dataloader(cls, is_training=False)

    def build_dataloader(self, is_training=True):
        """构建 dataloader 的具体实现"""
        transforms = self.build_transforms(self, self.data_dir, is_training)

        split_dir = "train_lmdb" if is_training else "val_lmdb"
        dataset = Kitti3D(
            data_path=os.path.join(self.data_dir, split_dir),
            transforms=transforms,
        )

        if is_training:
            sampler = torch.utils.data.DistributedSampler(dataset)
        else:
            sampler = None

        dataloader = torch.utils.data.DataLoader(
            dataset=dataset,
            batch_size=self.train_batch_size if is_training else self.val_batch_size,
            sampler=sampler,
            shuffle=False,
            num_workers=1,
            pin_memory=True,
            collate_fn=hat.data.collates.collate_kitti3d,
        )

        return dataloader

    def build_transforms(self, data_dir, is_training=True):
        """transforms 实现"""

        class_names = ["Car"]
        pc_range = [0, -39.68, -3, 69.12, 39.68, 1]

        if is_training:
            transforms = torchvision.transforms.Compose(
                [
                    ObjectSample(
                        class_names=class_names,
                        remove_points_after_sample=False,
                        db_sampler=DataBaseSampler(
                            enable=True,
                            root_path=data_dir,
                            db_info_path=os.path.join(
                                data_dir, "kitti3d_dbinfos_train.pkl"
                            ),  # noqa E501
                            sample_groups=[dict(Car=15)],  # noqa C408
                            db_prep_steps=[  # noqa C408
                                dict(  # noqa C408
                                    type="DBFilterByDifficulty",
                                    filter_by_difficulty=[-1],
                                ),
                                dict(  # noqa C408
                                    type="DBFilterByMinNumPoint",
                                    filter_by_min_num_points=dict(  # noqa C408
                                        Car=5,
                                    ),
                                ),
                            ],
                            global_random_rotation_range_per_object=[0, 0],
                            rate=1.0,
                        ),
                    ),
                    ObjectNoise(
                        gt_rotation_noise=[-0.15707963267, 0.15707963267],
                        gt_loc_noise_std=[0.25, 0.25, 0.25],
                        global_random_rot_range=[0, 0],
                        num_try=100,
                        class_names=class_names,
                    ),
                    PointRandomFlip(probability=0.5),
                    PointGlobalRotation(rotation=[-0.78539816, 0.78539816]),
                    PointGlobalScaling(min_scale=0.95, max_scale=1.05),
                    ShufflePoints(True),
                    ObjectRangeFilter(point_cloud_range=pc_range),
                    LidarReformat(),
                ]
            )
        else:
            transforms = torchvision.transforms.Compose([Reformat()])
        return transforms

class PointPillarsModel:
    ...

    @classmethod
    def optimizer(cls, model, stage):
        """训练时用到的 optimizer 设置"""
        if stage == "float":
            optimizer = torch.optim.AdamW(
                params=model.parameters(),
                betas=(0.95, 0.99),
                lr=2e-4,
                weight_decay=0.01,
            )
        elif stage == "qat":
            optimizer = torch.optim.SGD(
                params=model.parameters(),
                lr=2e-4,
                momentum=0.9,
                weight_decay=0.0,
            )
        else:
            optimizer = None

        return optimizer

    @classmethod
    def lr_schedule(cls, stage):
        """训练时用到的学习率调整策略"""
        if stage == "float":
            lr_updater = CyclicLrUpdater(
                target_ratio=(10, 1e-4),
                cyclic_times=1,
                step_ratio_up=0.4,
                step_log_interval=50,
            )

        elif stage == "qat":
            lr_updater = CyclicLrUpdater(
                target_ratio=(10, 1e-4),
                cyclic_times=1,
                step_ratio_up=0.4,
                step_log_interval=50,
            )
        else:
            lr_updater = None

        return lr_updater

    @classmethod
    def train_metrics(cls):
        """训练时打印 Log，只显示 loss 即可"""
        return LossShow()

    @classmethod
    def val_metrics(cls):
        """模型评测的 Metric"""
        class_names = ["Car"]
        val_metrics = Kitti3DMetricDet(
            compute_aos=True,
            current_classes=class_names,
            difficultys=[0, 1, 2],
        )
        return val_metrics

model.fuse_model()
model.set_qconfig()
horizon.quantization.prepare_qat(model, inplace=True)