# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # 旷视AI智慧交通开源赛道-交通标志识别
#
# * 队伍：仙交小分队 
#
# * 初赛得分：
#
#
# 本次的比赛是交通标志检测，交通标志本身种类众多，大小不定，并且在交通复杂的十字路口场景下，由于光照、天气等因素的影响，使其被精确检测变得更加困难。通过反复实验，我们选择了今年旷视新出的YOLOX目标检测框架，megengine版本的yolox代码相对于pyotorch版本的bug较多，通过反复调试，我们最终的方案是：YOLOX_L + P6 + Focalloss + Inputsize2048 + 双线性插值上采样 + dataaug30，以下是整个调试的过程。

# ## 数据分析
# 根据官方提供的示例代码，我们主要是在预训练的模型上面做微调。
#
# 本次比赛提供的数据集，图片总数为3349张，长的均值为1680.0883845924157， 宽的均值为2227.9558077037923，训练集和验证集一共有2700张，初赛的测试集有580张。
#
# 详细的规模如下：
#
# ```
# (h, w, c): num
# (1200, 1600, 3): 941
#
# (1944, 2592, 3): 1259
#
# (2048, 2448, 3): 18
#
# (1080, 1920, 3): 32
#
# (1520, 2704, 3): 511
#
# (2048, 2048, 3): 582
# ```
#
# 数据一共包含5种类型的目标，分别是：红灯、直行标志、向左转弯标志、禁止驶入和禁止临时停车，目标的数量如下
#
# ```
# 0:red_tl--1465个
# 1:arr_s--1133个
# 2:arr_l--638个
# 3:no_driving_mark_allsort--622个
# 4:no_parking_mark--1142个
# ```
#
# 目前的数据集来看，第三类的目标比较少，还是和之前一样，毕竟是数据驱动的，我们最好还是能通过copy-paste这种方法在数据上做一些提升。
#
# 具体数据的例子如下，从下面的图中可以看出，目标基本都比较小，而且还存在很多遮挡的情况。
#
# ![](https://vehicle4cm.oss-cn-beijing.aliyuncs.com/typoraimgs/3279125,1bb5780005b1590b9.jpg)
#
# ![](https://vehicle4cm.oss-cn-beijing.aliyuncs.com/typoraimgs/image-20210916135049212.png)
#
#

# ## megengine-yolox的安装
#
# 比赛中要求使用的框架是megengine，其中yolox有两个实现版本，一个是[pytorch](https://github.com/Megvii-BaseDetection/YOLOX)版本的，一个是[megengine](https://github.com/MegEngine/YOLOX)版本的，其中pyotrch版本的代码相对于megengin版本的代码要相对更加完善一些，并且在pytorch中提供了大量的coco预训练模型，这些预训练模型也可以转化为megengine的版本做finetune使用。
#
# 首先需要安装相应的虚拟环境，本次的megstudio中提供了对应的环境，大家只需要在启动的时候选择megengine为1.4.0，python版本为3.7的开发环境即可，conda中名称为`xuan`的虚拟环境就是本次项目所要使用到的虚拟环境，安装分为三步。
#
# 第一步：安装YOLOX
#
# 注：因为在megstudio中已经事先安装好了megengine，所以需要将requirements中的注释，另外执行setup.py的时候会检查torch，所以还需要在requirements.txt中额外添加torch和torchvision。

# %cd /home/megstudio/workspace
# !git clone https://gitee.com/cmfighting/YOLOX.git

# %cd /home/megstudio/workspace/YOLOX
# !source activate xuan # 激活之后可以在命令行前面看到一个xuan的括号
# !pip install -r requirements.txt
# !python setup.py develop

第二步：安装pycocotools

# !pip install cython
# !pip install pycocotools

# 第三步：测试yolox是否可以使用
#
# 我们需要从yolox的megengine版本的官方地址下载官方提供的yolox-tiny的模型，执行下面的指令之后，如果可以在outputs的文件夹找到检测的正确结果，证明yolox已安装完成。
#

# %cd /home/megstudio/workspace/YOLOX
# !mkdir pretrained
# %cd pretrained
# !wget https://github.com/MegEngine/YOLOX/releases/download/0.0.1/yolox_tiny.pkl # 下载不了的话请大家手动下载并上传

# %cd /home/megstudio/workspace/YOLOX
# !python tools/demo.py image -n yolox-tiny -c pretrained/yolox_tiny.pkl --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 416 --save_result

# ## 训练交通标志检测的yolox-l模型
#
# 本次比赛提供的数据集是交通标志检测的数据集，因为本次比赛要求大家使用megengine作为我们的基本框架，megengine版本的yolox代码相对于pytorch版本的yolox代码不是很完整，整个训练的过程中存在很多小坑，包括缺少预训练模型，模型迁移过程中shape不一致，以及训练的时候可能会随机在某个时刻突然停止等，下面记录了整个训练的流程，最终训练好的模型也做了保存，如果大家没有足够多的算力卡来跑完训练的流程，可以直接跳到下面的验证过程。
#
# * 第一步：建立数据集的软链接
#
# 首先还是建立虚拟链接，以megstudio为例，数据保存的目录在`/home/megstudiodataset/dataset-2805/`目录下，其中annotations是标签文件，images放的是图片文件，如下图所示：
#
#   ![image-20210923205007420](https://vehicle4cm.oss-cn-beijing.aliyuncs.com/typoraimgs/image-20210923205007420.png)
#
#   
#   建立虚拟链接的代码如下：

# %cd /home/megstudio/workspace/YOLOX/datasets
# !mkdir traffic5
# !ln -s /home/megstudio/dataset/dataset-2805/images ./traffic5/train2017
# !ln -s /home/megstudio/dataset/dataset-2805/images ./traffic5/val2017
# !ln -s /home/megstudio/dataset/dataset-2805/images ./traffic5/test2017

# 建立annotations文件夹，然后把annntations的标注问价复制过来，改个名字
# %cd traffic5
# !mkdir annotations
# %cd annotations

# annotions文件要和coco的保持一致，因为这个数据集是只读的，所以我们只能新建文件之后把里面的内容复制粘贴过来（应该是有某些linux的命令可以使用，但是我太菜了），标注文件大致如下方所示：
#
# ![image-20210923205415437](https://vehicle4cm.oss-cn-beijing.aliyuncs.com/typoraimgs/image-20210923205415437.png)
#

# * 第二步：模型迁移
#
# 由于官方的megengine版本的yolox代码只提供了tiny版本的预训练模型，tiny版本的预训练模型速度虽然比较快，
# 但是精度相对较低，想要在比赛中拿到比较好的成绩的话，还是得用l或者x版本得预训练模型，这样就预训练的模型只能从官方的pytorch代码中进行转化。
#
# 这部分操作需要借助pytorch版本的yolox代码来完成，下载好yolox_l模型之后，通过
# [YOLOX/demo/MegEngine/python at main · Megvii-BaseDetection/YOLOX (github.com)](https://github.com/Megvii-BaseDetection/YOLOX/tree/main/demo/MegEngine/python)这里的教程完成模型的转化，转换完成的模型需要上传到pretrained目录下。
#

# +
* 第二步：注释数据归一化部分代码

pytorch转换后的模型在训练的时候是没有经过归一化处理的，官方的megengine版本代码是做了归一化处理的，如果直接进行训练，第一轮的map可能为0.0，
整体的效果也不会太好，这部分的逻辑在`yolox/data/data_augment.py`中，我们需要将179行到184行的代码注释掉，也即使下面的这段代码。

# +
#     image = image[:, :, ::-1]
#     image /= 255.0
#     if mean is not None:
#         image -= mean
#     if std is not None:
#         image /= std
# -

# * 第三步：修改`yolox/utils/checkpoint.py`
#
# 如果直接跑训练的代码的话，在模型迁移的时候中间会有大段地不匹配出现，分析发现可能是pyotrch转化过来地模型和megengine表示方式可能不太一样，最后我们通过reshape和try catch块处理了一下，
# 就能正确地完成整个迁移地过程了，修改之后地checkpoint.py文件如下

# +
# #!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
import os
import shutil
from loguru import logger
import megengine as mge

def load_ckpt(model, ckpt):
    model_state_dict = model.state_dict()
    load_dict = {}
    # v是期待的模型
    for key_model, v in model_state_dict.items():

        if key_model not in ckpt:
            logger.warning(
                "{} is not in the ckpt. Please double check and see if this is desired.".format(
                    key_model
                )
            )
            continue
        v_ckpt = ckpt[key_model]
        # 现在需要保证ckpt的能塞进去
        # print(key_model)
        # head.cls_preds.0.bias
        try:
            v_ckpt =v_ckpt.reshape(v.shape)
            if v.shape != v_ckpt.shape:
                logger.warning(
                    "ckpt Shape of {} in checkpoint is {}, while model shape of {} in model is {}.".format(
                        key_model, v_ckpt.shape, key_model, v.shape
                    )
                )
                continue
            load_dict[key_model] = v_ckpt
        except:
            logger.warning(
                "{} is not match".format(key_model)
                )
            # continue

    # for i in range(3):
    #     load_dict.pop("head.grids.{}".format(i))
    model.load_state_dict(load_dict, strict=False)
    return model

def save_checkpoint(state, is_best, save_dir, model_name=""):
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    filename = os.path.join(save_dir, model_name + "_ckpt.pkl")
    mge.save(state, filename)
    if is_best:
        best_filename = os.path.join(save_dir, "best_ckpt.pkl")
        shutil.copyfile(filename, best_filename)


# -

# * 第四步：添加P6层
#
# 因为我们考虑使用大地分辨率来训练模型，p6层只需要对p5做一次卷积就可以，相对于添加p2来说，不会对预训练地模型产生那么多地影响，所以我们需要对`yolox/models/yolo_pafpn.py`做修改，
# 添加p6层，在初始化中增加一个卷积层并在前向传播的最后对p5使用这个卷积，修改之后的代码如下：

# +
# #!/usr/bin/env python
# -*- encoding: utf-8 -*-
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.

import megengine.functional as F
import megengine.module as M
from .darknet import CSPDarknet
from .network_blocks import BaseConv, CSPLayer, DWConv, UpSample


class YOLOPAFPN(M.Module):
    """
    YOLOv3 model. Darknet 53 is the default backbone of this model.
    """
    def __init__(
        self, depth=1.0, width=1.0, in_features=("dark3", "dark4", "dark5"),
        in_channels=[256, 512, 1024], depthwise=False, act="silu",
    ):
        super().__init__()
        self.backbone = CSPDarknet(depth, width, depthwise=depthwise, act=act)
        self.in_features = in_features
        self.in_channels = in_channels
        Conv = DWConv if depthwise else BaseConv

        self.upsample = UpSample(scale_factor=2, mode="bilinear")
        self.lateral_conv0 = BaseConv(
            int(in_channels[2] * width), int(in_channels[1] * width), 1, 1, act=act
        )
        self.C3_p4 = CSPLayer(
            int(2 * in_channels[1] * width),
            int(in_channels[1] * width),
            round(3 * depth),
            False,
            depthwise=depthwise,
            act=act,
        )  # cat

        self.reduce_conv1 = BaseConv(
            int(in_channels[1] * width), int(in_channels[0] * width), 1, 1, act=act
        )
        self.C3_p3 = CSPLayer(
            int(2 * in_channels[0] * width),
            int(in_channels[0] * width),
            round(3 * depth),
            False,
            depthwise=depthwise,
            act=act,
        )

        # bottom-up conv
        self.bu_conv2 = Conv(
            int(in_channels[0] * width), int(in_channels[0] * width), 3, 2, act=act
        )
        self.C3_n3 = CSPLayer(
            int(2 * in_channels[0] * width),
            int(in_channels[1] * width),
            round(3 * depth),
            False,
            depthwise=depthwise,
            act=act,
        )

        # bottom-up conv
        self.bu_conv1 = Conv(
            int(in_channels[1] * width), int(in_channels[1] * width), 3, 2, act=act
        )
        self.C3_n4 = CSPLayer(
            int(2 * in_channels[1] * width),
            int(in_channels[2] * width),
            round(3 * depth),
            False,
            depthwise=depthwise,
            act=act,
        )
        
        self.conv_p6 = Conv(
            int(1024 * width), int(1024 * width), 3, 2, act=act
        )

    def forward(self, input):
        """
        Args:
            inputs: input images.

        Returns:
            Tuple[Tensor]: FPN feature.
        """

        #  backbone
        out_features = self.backbone(input)
        features = [out_features[f] for f in self.in_features]
        [x2, x1, x0] = features

        fpn_out0 = self.lateral_conv0(x0)  # 1024->512/32
        f_out0 = self.upsample(fpn_out0)  # 512/16
        f_out0 = F.concat([f_out0, x1], 1)  # 512->1024/16
        f_out0 = self.C3_p4(f_out0)  # 1024->512/16

        fpn_out1 = self.reduce_conv1(f_out0)  # 512->256/16
        f_out1 = self.upsample(fpn_out1)  # 256/8
        f_out1 = F.concat([f_out1, x2], 1)  # 256->512/8
        pan_out2 = self.C3_p3(f_out1)  # 512->256/8

        p_out1 = self.bu_conv2(pan_out2)  # 256->256/16
        p_out1 = F.concat([p_out1, fpn_out1], 1)  # 256->512/16
        pan_out1 = self.C3_n3(p_out1)  # 512->512/16

        p_out0 = self.bu_conv1(pan_out1)  # 512->512/32
        p_out0 = F.concat([p_out0, fpn_out0], 1)  # 512->1024/32
        pan_out0 = self.C3_n4(p_out0)  # 1024->1024/32
        
        pan_p6 = self.conv_p6(pan_out0)

        outputs = (pan_out2, pan_out1, pan_out0, pan_p6)
        return outputs



# -

# 修改完之后还需要在`yolo_head.py`和`exp/yolox_base.py`中做修改，和修改之后的pafpn匹配起来。

# * 第五步：修改为focalloss
#
# 通过实验发现focalloss能带来一部分的涨点，所以我们将原先的bceloss修改为了focalloss，具体的代码在`yolox/models/yolox_head.py`中, 如下:

# 添加focal
def focal_loss_discrite(self, pred, gt):
    pos_inds = F.equal(gt, 1).astype("float32")
    neg_inds = F.equal(gt, 0).astype("float32")
    pos_loss = F.log(pred+1e-5) * F.pow(1-pred, 2) * pos_inds * 0.75
    neg_loss = F.log(1-pred+1e-5) * F.pow(pred, 2) * neg_inds * 0.25
    loss = -(pos_loss + neg_loss)
    return loss
loss_obj = (
            # todo 修改为focalloss
            self.focal_loss_discrite(F.sigmoid(obj_preds).view(-1, 1), obj_targets)
            # self.bcewithlog_loss(obj_preds.view(-1, 1), obj_targets) # 原先的loss
        ).sum() / num_fg

# * 第六步：修改配置文件
#
# 修改后的配置文件如下，这里我们使用yolox_l作为我们的基础网络，考虑到小目标的问题，我们将inputsize调整到了2048，并将数据增强的轮数进行了翻倍

# #!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) Megvii, Inc. and its affiliates.
import os
from yolox.exp import Exp as MyExp
class Exp(MyExp):
    def __init__(self):
        super(Exp, self).__init__()
        self.input_size = (2048, 2048)  # (height, width)
        self.test_size = (2048, 2048)
        self.no_aug_epochs = 30
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
        # Define yourself dataset path
        self.data_dir = "datasets/traffic5"
        self.train_ann = "instances_train2017.json"
        self.val_ann = "instances_val2017.json"
        # self.val_ann = "instances_test2017.json" # 测试的时候解开注释
        self.num_classes = 5
        self.max_epoch = 300
        self.data_num_workers = 4
        self.eval_interval = 1

# * 第七步：将fpn中的上采样方式修改为双线性插值
#
# ```
# self.upsample = UpSample(scale_factor=2, mode="bilinear")
# ```
#

# * 第八步：开始训练
#
# 训练是在megstudio的v100显卡上完成的，如果你选择的显卡显存比较下，可以将inputsize的大小调整为640。
#
# 训练的代码如下：
#
# 会有两处报错，第一次报错的时候需要将trainer.py 250行的['model']去掉
# 第二次报错的时候在yolo_base.py中添加self.data_dir = "datasets/traffic5"并在100和129行传入

# %cd /home/megstudio/workspace/YOLOX
# !python tools/train.py -f exps/example/yolox_l.py -d 1 -b 4 -c pretrained/yolox_l_mge.pkl

# ## 模型验证
#
# 训练好的的模型在YOLOX_outputs/yolox_l/best_ckpt.pkl 目录下，执行下列的指令就可以完成验证的过程。

# %cd /home/megstudio/workspace/YOLOX
# !source activate xuan
# !python tools/eval.py -f exps/example/yolox_l.py -c YOLOX_outputs/yolox_l/best_ckpt.pkl -b 1 -d 1 --conf 0.01

# ## 生成results.json
#
# 接下来要生成测试图片的结果，首先需要解开`yolox_l.py`18行的注释解开，并在`coco_evaluator.py`的167行解开注释，在/home/megstudio/workspace/submit下生成results.json文件

# %cd /home/megstudio/workspace/YOLOX
# !source activate xuan
# !python tools/eval.py -f exps/example/yolox_l.py -c YOLOX_outputs/yolox_l/best_ckpt.pkl -b 1 -d 1 --conf 0.001

# ## 评价测试集
# 直接在原先数据集的基础上更改测试集即可，建立软链接
#

# %cd /home/megstudio/workspace/YOLOX/datasets
# !ln -s /home/megstudio/dataset/traffic5/images ./traffic5/test2017

# %cd /home/megstudio/workspace/YOLOX
# !source activate xuan # 激活之后可以在命令行前面看到一个xuan的括号
# !pip install -r requirements.txt
# !python setup.py develop
# !pip install pycocotools
# # !python tools/eval.py -f exps/example/yolox_l.py -c YOLOX_outputs/yolox_l/best_ckpt.pkl -b 1 -d 1 --conf 0.01

# 重新试试在验证集上的效果，保证不出错
# %cd /home/megstudio/workspace/YOLOX
# !source activate xuan
# !python tools/eval.py -f exps/example/yolox_l.py -c YOLOX_outputs/yolox_l/best_ckpt.pkl -b 1 -d 1 --fuse --conf 0.1

# 重新试试在验证集上的效果，保证不出错
# %cd /home/megstudio/workspace/YOLOX
# !source activate xuan
# !python tools/eval.py -f exps/example/yolox_l.py -c YOLOX_outputs/yolox_l/best_ckpt.pkl -b 1 -d 1 --fuse --conf 0.1

StarcoderdataJupyterScriptsDedupFiltered