在 COCO数据集上训练 YOLOX （代码详解）

在 COCO数据集上训练 YOLOX （代码详解）

————>第一部分：训练自定义模型

库导入、cuda版本

import warnings
warnings.filterwarnings("ignore")

import ast
import os
import json
import pandas as pd
import torch
import importlib
import cv2 

from shutil import copyfile
from tqdm.notebook import tqdm
tqdm.pandas()
from sklearn.model_selection import GroupKFold
from PIL import Image
from string import Template
from IPython.display import display

TRAIN_PATH = '/input/tensorflow-great-barrier-reef'

# check Torch and CUDA version
print(f"Torch: {torch.__version__}")
!nvcc --version

（一）安装yolox

!git clone https://github.com/Megvii-BaseDetection/YOLOX -q

%cd YOLOX
!pip install -U pip && pip install -r requirements.txt
!pip install -v -e . 

!pip install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

（二）数据集COTS的准备

2.1 数据集及注释

def get_bbox(annots):
    bboxes = [list(annot.values()) for annot in annots]
    return bboxes

def get_path(row):
    row['image_path'] = f'{TRAIN_PATH}/train_images/video_{row.video_id}/{row.video_frame}.jpg'
    return row

• 函数 get_bbox ：将字典的值转换为列表，并将列表组成的列表存储在变量 bboxes 中
• 函数 get_path 通过使用 f-string 将 TRAIN_PATH、row.video_id 和 row.video_frame 的值组合成字符串，并将其赋值给 row['image_path'] 。
• 在处理图像数据时使用的。函数 get_bbox 用于从图像注释中提取边界框信息，函数 get_path 用于根据图像的视频ID和帧数生成图像路径。

df = pd.read_csv("/input/tensorflow-great-barrier-reef/train.csv")
df.head(5)

# Taken only annotated photos
df["num_bbox"] = df['annotations'].apply(lambda x: str.count(x, 'x'))
df_train = df[df["num_bbox"]>0]

#Annotations 
df_train['annotations'] = df_train['annotations'].progress_apply(lambda x: ast.literal_eval(x))
df_train['bboxes'] = df_train.annotations.progress_apply(get_bbox)

#Images resolution
df_train["width"] = 1280
df_train["height"] = 720

#Path of images
df_train = df_train.progress_apply(get_path, axis=1)

kf = GroupKFold(n_splits = 5) 
df_train = df_train.reset_index(drop=True)
df_train['fold'] = -1
for fold, (train_idx, val_idx) in enumerate(kf.split(df_train, y = df_train.video_id.tolist(), groups=df_train.sequence)):
    df_train.loc[val_idx, 'fold'] = fold

df_train.head(5)

• 使用 GroupKFold 进行数据集的分割，并为每个样本分配一个折叠（fold）编号。
  -通过实例化 GroupKFold 类，并设置 n_splits 参数为 5，创建了名为 kf 的交叉验证对象。
• 接下来，通过使用 df_train.reset_index(drop=True) ，对 df_train 数据框进行
• 添加了一个名为 fold' 的新列到 df_train 数据框，并将所有行的值初始化为 -1。将用于存储每个样本所属的折叠编号。
• 通过使用 enumerate 函数遍历，其中 df_train 是要划分的数据集，y 是要划分的目标变量， groups 是指定样本分组的列。
• 在每次循环中， train_idx 是训练集的索引， val_idx 是验证集的索引， fold 是当前的折叠编号。
• 最后，通过使用 df_train.loc[val_idx, 'fold'] = fold ，将验证集对应的行的 fold 列的值设置为当前的折叠编号，完成了对数据集的折叠划分。

HOME_DIR = '/kaggle/working/' 
DATASET_PATH = 'dataset/images'

!mkdir {HOME_DIR}dataset
!mkdir {HOME_DIR}{DATASET_PATH}
!mkdir {HOME_DIR}{DATASET_PATH}/train2017
!mkdir {HOME_DIR}{DATASET_PATH}/val2017
!mkdir {HOME_DIR}{DATASET_PATH}/annotations
SELECTED_FOLD = 4

for i in tqdm(range(len(df_train))):
    row = df_train.loc[i]
    if row.fold != SELECTED_FOLD:
        copyfile(f'{row.image_path}', f'{HOME_DIR}{DATASET_PATH}/train2017/{row.image_id}.jpg')
    else:
        copyfile(f'{row.image_path}', f'{HOME_DIR}{DATASET_PATH}/val2017/{row.image_id}.jpg') 
        
print(f'Number of training files: {len(os.listdir(f"{HOME_DIR}{DATASET_PATH}/train2017/"))}')
print(f'Number of validation files: {len(os.listdir(f"{HOME_DIR}{DATASET_PATH}/val2017/"))}')

将数据集中的图像文件按照折叠划分的结果复制到不同的文件夹中，以用于训练和验证。

定义一个主目录路径 HOME_DIR 和数据集路径 DATASET_PATH。使用 !mkdir 命令创建了几个子目录，包括 train2017、val2017 和 annotations，用于存储训练集图像、验证集图像和注释文件。

选择一个特定的折叠编号 SELECTED_FOLD，用于指定要在验证集中使用的折叠。

使用循环遍历 df_train 数据框中的每一行，对于每个样本，根据其折叠编号将图像文件复制到相应的训练集或验证集文件夹中。
在循环中，使用 copyfile 函数将源文件复制到目标文件路径。

最后，通过使用 len(os.listdir(...)) 计算训练集文件夹和验证集文件夹中的文件数量，并将结果打印出来，显示训练集和验证集中的图像文件数量。

2.2 创建COCO文件注释

def save_annot_json(json_annotation, filename):
    with open(filename, 'w') as f:
        output_json = json.dumps(json_annotation)
        f.write(output_json)
annotion_id = 0

定义函数 save_annot_json 和变量 annotation_id。

函数 save_annot_json 用于将 JSON 格式的注释数据保存到文件中。
接受两个参数：json_annotation 是要保存的注释数据，filename 是要保存到的文件名。

函数使用 open 函数以写入模式打开文件，并将文件对象赋值给变量 f。然后，通过调用 json.dumps 函数将注释数据转换为 JSON 字符串格式，并将结果赋值给变量 output_json。使用文件对象的 write 方法将 JSON 字符串写入文件。

变量 annotation_id 是一个整数，用于标识注释的 ID。

def dataset2coco(df, dest_path):
    
    global annotion_id
    
    annotations_json = {
        "info": [],
        "licenses": [],
        "categories": [],
        "images": [],
        "annotations": []
    }
    
    info = {
        "year": "2021",
        "version": "1",
        "description": "COTS dataset - COCO format",
        "contributor": "",
        "url": "https://kaggle.com",
        "date_created": "2021-11-30T15:01:26+00:00"
    }
    annotations_json["info"].append(info)
    
    lic = {
            "id": 1,
            "url": "",
            "name": "Unknown"
        }
    annotations_json["licenses"].append(lic)

    classes = {"id": 0, "name": "starfish", "supercategory": "none"}

    annotations_json["categories"].append(classes)

    
    for ann_row in df.itertuples():
            
        images = {
            "id": ann_row[0],
            "license": 1,
            "file_name": ann_row.image_id + '.jpg',
            "height": ann_row.height,
            "width": ann_row.width,
            "date_captured": "2021-11-30T15:01:26+00:00"
        }
        
        annotations_json["images"].append(images)
        
        bbox_list = ann_row.bboxes
        
        for bbox in bbox_list:
            b_width = bbox[2]
            b_height = bbox[3]
            
            # some boxes in COTS are outside the image height and width
            if (bbox[0] + bbox[2] > 1280):
                b_width = bbox[0] - 1280 
            if (bbox[1] + bbox[3] > 720):
                b_height = bbox[1] - 720 
                
            image_annotations = {
                "id": annotion_id,
                "image_id": ann_row[0],
                "category_id": 0,
                "bbox": [bbox[0], bbox[1], b_width, b_height],
                "area": bbox[2] * bbox[3],
                "segmentation": [],
                "iscrowd": 0
            }
            
            annotion_id += 1
            annotations_json["annotations"].append(image_annotations)
        
        
    print(f"Dataset COTS annotation to COCO json format completed! Files: {len(df)}")
    return annotations_json

实现函数 dataset2coco，用于将数据集的注释信息转换为 COCO 格式的 JSON 文件。

接受两个参数：df 包含图像和注释信息；dest_path 是要保存生成的 COCO JSON 文件的目标路径。

在函数中，定义了一个全局变量 annotation_id，用于给每个注释分配唯一的标识符。创建了一个空的字典 annotations_json，用于存储最终的 COCO JSON 数据。

构建 JSON 数据的各个部分，包括 “info”、“licenses”、“categories”、“images” 和 “annotations”。在这些部分中，添加了相应的信息，如数据集的描述、许可证信息、类别信息等。

使用 df.itertuples() 迭代数据集的每一行。对于每一行，创建了一个表示图像的字典，并将其添加到 “images” 部分中。同时，获取注释框的列表，并遍历每个注释框。根据注释框的信息，创建一个表示注释的字典，并将其添加到 “annotations” 部分中。

在创建注释时，还对注释框的坐标进行了一些处理，确保注释框的宽度和高度不超过图像的宽度和高度。

最后，打印一条消息，表示数据集的注释已成功转换为 COCO 格式的 JSON 数据，并返回生成的 JSON 数据。

# Convert COTS dataset to JSON COCO
train_annot_json = dataset2coco(df_train[df_train.fold != SELECTED_FOLD], f"{HOME_DIR}{DATASET_PATH}/train2017/")
val_annot_json = dataset2coco(df_train[df_train.fold == SELECTED_FOLD], f"{HOME_DIR}{DATASET_PATH}/val2017/")

# Save converted annotations
save_annot_json(train_annot_json, f"{HOME_DIR}{DATASET_PATH}/annotations/train.json")
save_annot_json(val_annot_json, f"{HOME_DIR}{DATASET_PATH}/annotations/valid.json")

调用 dataset2coco 函数将训练集中不包含选定折叠的部分转换为 COCO 格式的 JSON 数据，并指定保存路径为训练集文件夹下。生成的 JSON 数据被赋值给变量 train_annot_json。

调用 dataset2coco 函数将训练集中包含选定折叠的部分转换为 COCO 格式的 JSON 数据，并指定保存路径为验证集文件夹下。生成的 JSON 数据被赋值给变量 val_annot_json。

调用 save_annot_json 函数，将训练集的 COCO JSON 注释保存到指定路径下的 train.json 文件中。

调用 save_annot_json 函数，将验证集的 COCO JSON 注释保存到指定路径下的 valid.json 文件中。

（三） 准备配置文件

# Choose model for your experiments NANO or YOLOX-S (you can adapt for other model type)

NANO = False

3.1 YOLOX-S 实验配置文件

config_file_template = '''

#!/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.depth = 0.33
        self.width = 0.50
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
        
        # Define yourself dataset path
        self.data_dir = "/working/dataset/images"
        self.train_ann = "train.json"
        self.val_ann = "valid.json"

        self.num_classes = 1

        self.max_epoch = $max_epoch
        self.data_num_workers = 2
        self.eval_interval = 1
        
        self.mosaic_prob = 1.0
        self.mixup_prob = 1.0
        self.hsv_prob = 1.0
        self.flip_prob = 0.5
        self.no_aug_epochs = 2
        
        self.input_size = (960, 960)
        self.mosaic_scale = (0.5, 1.5)
        self.random_size = (10, 20)
        self.test_size = (960, 960)
'''

模型配置文件的模板，用于定义实验的配置参数。

模板中定义了一个继承自 yolox.exp.Exp 的类 Exp，用于配置实验参数。

在 Exp 类的 __init__ 方法中，初始化了一些基本的实验参数，如模型的深度、宽度、实验名称等。

定义数据集的路径和注释文件名。需要根据实际情况修改 data_dir、train_ann 和 val_ann 的值，以指向正确的数据集路径和注释文件。

其他的参数设置包括类别数量、最大训练轮数、数据加载的进程数量、评估间隔等。设置数据增强的相关参数，如 Mosaic、Mixup、HSV 变换的概率、翻转概率等。

input_size 表示输入图像的大小，mosaic_scale 表示 Mosaic 数据增强的尺度范围，random_size 表示随机缩放的尺度范围，test_size 表示测试图像的大小。

配置文件中的 $max_epoch 是一个占位符，表示最大训练轮数的值需要在具体使用时进行替换。

这些是模型配置文件中的一些参数和设置，它们的含义如下：

• self.num_classes: 表示数据集中的类别数量。在这个例子中，设置为 1，表示只有一个类别（starfish）。
• self.max_epoch: 表示训练的最大轮数。在实际使用时，你需要将 $max_epoch 替换为具体的数值，例如 100 或其他适当的值。
• self.data_num_workers: 表示数据加载时使用的进程数量。在这个例子中，设置为 2。
• self.eval_interval: 表示评估模型的间隔轮数。在这个例子中，每训练 1 轮进行一次评估。
• self.mosaic_prob: 表示应用 Mosaic 数据增强的概率。在这个例子中，设置为 1.0，表示始终应用 Mosaic。
• self.mixup_prob: 表示应用 Mixup 数据增强的概率。在这个例子中，设置为 1.0，表示始终应用 Mixup。
• self.hsv_prob: 表示应用 HSV 变换的概率。在这个例子中，设置为 1.0，表示始终应用 HSV 变换。
• self.flip_prob: 表示图像翻转的概率。在这个例子中，设置为 0.5，表示有 50% 的概率进行图像翻转。
• self.no_aug_epochs: 表示不应用数据增强的轮数。在这个例子中，设置为 2，表示前两轮不进行数据增强。
• self.input_size: 表示模型输入的图像大小。在这个例子中，设置为 (960, 960)。
• self.mosaic_scale: 表示 Mosaic 数据增强的尺度范围。在这个例子中，设置为 (0.5, 1.5)，表示 Mosaic 数据增强会将图像尺度缩放到原始尺度的 0.5 到 1.5 倍之间。
• self.random_size: 表示随机缩放的尺度范围。在这个例子中，设置为 (10, 20)，表示随机缩放的尺度为原始尺度的 10% 到 20% 之间。
• self.test_size: 表示测试图像的大小。在这个例子中，设置为 (960, 960)。

3.2 YOLOX-NANO 配置文件

if NANO:
    config_file_template = '''

#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) Megvii, Inc. and its affiliates.

import os

import torch.nn as nn

from yolox.exp import Exp as MyExp


class Exp(MyExp):
    def __init__(self):
        super(Exp, self).__init__()
        self.depth = 0.33
        self.width = 0.25
        self.input_size = (416, 416)
        self.mosaic_scale = (0.5, 1.5)
        self.random_size = (10, 20)
        self.test_size = (416, 416)
        self.exp_name = os.path.split(
            os.path.realpath(__file__))[1].split(".")[0]
        self.enable_mixup = False

        # Define yourself dataset path
        self.data_dir = "/working/dataset/images"
        self.train_ann = "train.json"
        self.val_ann = "valid.json"

        self.num_classes = 1

        self.max_epoch = $max_epoch
        self.data_num_workers = 2
        self.eval_interval = 1

    def get_model(self, sublinear=False):
        def init_yolo(M):
            for m in M.modules():
                if isinstance(m, nn.BatchNorm2d):
                    m.eps = 1e-3
                    m.momentum = 0.03

        if "model" not in self.__dict__:
            from yolox.models import YOLOX, YOLOPAFPN, YOLOXHead
            in_channels = [256, 512, 1024]
            # NANO model use depthwise = True, which is main difference.
            backbone = YOLOPAFPN(self.depth,
                                 self.width,
                                 in_channels=in_channels,
                                 depthwise=True)
            head = YOLOXHead(self.num_classes,
                             self.width,
                             in_channels=in_channels,
                             depthwise=True)
            self.model = YOLOX(backbone, head)

        self.model.apply(init_yolo)
        self.model.head.initialize_biases(1e-2)
        return self.model

'''

如果选择的是 NANO 模型，那么会使用以下的配置信息：

• self.depth = 0.33: NANO 模型的深度设置为 0.33。
• self.width = 0.25: NANO 模型的宽度设置为 0.25。
• self.input_size = (416, 416): NANO 模型的输入图像大小设置为 (416, 416)。
• self.mosaic_scale = (0.5, 1.5): NANO 模型的 Mosaic 数据增强尺度范围设置为 (0.5, 1.5)。
• self.random_size = (10, 20): NANO 模型的随机缩放尺度范围设置为 (10, 20)。
• self.test_size = (416, 416): NANO 模型的测试图像大小设置为 (416, 416)。
• self.enable_mixup = False: NANO 模型的 Mixup 数据增强设置为 False。

创建模型，并对模型的参数进行初始化操作。返回的模型实例可以在训练和推理阶段使用。

定义 get_model 方法，用于获取模型的实例。

• init_yolo 函数用于初始化 YOLO 模型中的 BatchNormalization 层的参数。
• 如果 model 不在当前实例的属性字典中，即模型尚未被创建，则会执行以下步骤来创建模型：
  • 从 yolox.models 模块中导入 YOLOX、YOLOPAFPN 和 YOLOXHead。
  • 设置输入通道数 in_channels 为 [256, 512, 1024]。
  • 使用 YOLOPAFPN 创建骨干网络 backbone，其中使用了 NANO 模型特有的 depthwise=True 参数。
  • 使用 YOLOXHead 创建模型的头部 head，其中使用了 NANO 模型特有的 depthwise=True 参数。
  • 使用 YOLOX 将骨干网络和头部连接起来，创建最终的模型实例 self.model。
• 调用 init_yolo 函数来初始化模型的 BatchNormalization 层的参数。
• 调用 initialize_biases 方法来初始化模型头部的偏置项。
• 返回模型实例 self.model。

PIPELINE_CONFIG_PATH='cots_config.py'

pipeline = Template(config_file_template).substitute(max_epoch = 20)

with open(PIPELINE_CONFIG_PATH, 'w') as f:
    f.write(pipeline)

根据模板 config_file_template 生成一个管道配置文件，并将其写入到 PIPELINE_CONFIG_PATH 指定的文件中。

使用 Template 类对模板进行处理，将其中的 $max_epoch 替换为实际的值（ 20）。生成的配置文件的内容保存在变量 pipeline 中。

使用 open 函数以写入模式打开 PIPELINE_CONFIG_PATH 文件，并将 pipeline 的内容写入文件中，完成配置文件的创建和保存操作。

# ./yolox/data/datasets/voc_classes.py

voc_cls = '''
VOC_CLASSES = (
  "starfish",
)
'''
with open('./yolox/data/datasets/voc_classes.py', 'w') as f:
    f.write(voc_cls)

# ./yolox/data/datasets/coco_classes.py

coco_cls = '''
COCO_CLASSES = (
  "starfish",
)
'''
with open('./yolox/data/datasets/coco_classes.py', 'w') as f:
    f.write(coco_cls)

# check if everything is ok    
!more ./yolox/data/datasets/coco_classes.py

创建两个类别文件 voc_classes.py 和 coco_classes.py，用于指定数据集的类别信息。

使用 voc_cls 存储了 VOC 数据集的类别信息，其中只包含一个类别 “starfish”。使用 open 函数以写入模式打开 ./yolox/data/datasets/voc_classes.py 文件，并将 voc_cls 的内容写入文件中，完成 VOC 类别文件的创建和保存操作。

使用变量 coco_cls 存储了 COCO 数据集的类别信息，同样只包含一个类别 “starfish”。再次使用 open 函数以写入模式打开 ./yolox/data/datasets/coco_classes.py 文件，并将 coco_cls 的内容写入文件中，完成 COCO 类别文件的创建和保存操作。

使用 !more 命令检查 ./yolox/data/datasets/coco_classes.py 文件的内容，以确认文件是否正确创建并包含了所需的类别信息。

（四）下载预训练权重

预训练模型列表：

• YOLOX-s
• YOLOX-m
• YOLOX-nano for inference speed (!)
• etc.

sh = 'wget https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_s.pth'
MODEL_FILE = 'yolox_s.pth'

if NANO:
    sh = '''
    wget https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_nano.pth
    '''
    MODEL_FILE = 'yolox_nano.pth'

with open('script.sh', 'w') as file:
  file.write(sh)

!bash script.sh

根据变量 NANO 的值，选择相应的权重文件的下载链接和文件名。
将下载链接保存到 script.sh 文件中，并使用 !bash 命令来执行 script.sh 脚本，从链接中下载相应的权重文件。

执行完整段代码后，会根据 NANO 的值下载相应的预训练权重文件，并保存到当前工作目录中。

（五）训练模型

!cp ./tools/train.py ./
!python train.py \
    -f cots_config.py \
    -d 1 \
    -b 32 \
    --fp16 \
    -o \
    -c {MODEL_FILE}   # Remember to chenge this line if you take different model eg. yolo_nano.pth, yolox_s.pth or yolox_m.pth

执行模型的训练操作。
使用 python train.py 命令来启动训练过程，传递了一些参数：

• -f cots_config.py 指定了配置文件的路径为 cots_config.py。
• -d 1 指定了使用的 GPU 设备编号为 1。
• -b 32 指定了每个批次的图像数量为 32。
• --fp16 启用混合精度训练，使用半精度浮点数进行计算。
• -o 启用优化器。
• -c {MODEL_FILE} 指定了使用的预训练权重文件，{MODEL_FILE} 会被替换为实际的文件名，例如 yolox_s.pth。

执行完整段代码后，模型将开始训练，并根据配置文件中的设置进行优化和保存模型参数。
# ————>第二部分：推理

（六）运行推理

6.1 使用 YOLOX 工具进行推理

%cp ../../input/yolox-fix-for-demo-inference/demo.py tools/demo.py

TEST_IMAGE_PATH = "/working/dataset/images/val2017/0-4614.jpg"
MODEL_PATH = "./YOLOX_outputs/cots_config/best_ckpt.pth"

!python tools/demo.py image \
    -f cots_config.py \
    -c {MODEL_PATH} \
    --path {TEST_IMAGE_PATH} \
    --conf 0.1 \
    --nms 0.45 \
    --tsize 960 \
    --save_result \
    --device gpu

这段代码执行了模型的推理操作，用于在测试图像上进行目标检测。
使用 TEST_IMAGE_PATH 变量指定了测试图像的路径，该图像将用于进行目标检测。

使用 MODEL_PATH 变量指定了模型的权重文件路径，该文件将用于加载训练好的模型参数。

使用 python tools/demo.py image 命令启动推理过程，传递了一些参数：

• -f cots_config.py 指定了配置文件的路径为 cots_config.py。
• -c {MODEL_PATH} 指定了使用的模型权重文件的路径，{MODEL_PATH} 会被替换为实际的文件路径。
• --path {TEST_IMAGE_PATH} 指定了测试图像的路径，{TEST_IMAGE_PATH} 会被替换为实际的图像路径。
• --conf 0.1 指定了置信度阈值为 0.1，低于该阈值的检测结果将被过滤。
• --nms 0.45 指定了非极大值抑制的阈值为 0.45，用于抑制重叠的检测框。
• --tsize 960 指定了图像的目标尺寸为 960x960 像素。
• --save_result 启用保存检测结果的功能，结果将保存在 YOLOX_outputs 目录中。
• --device gpu 指定了使用 GPU 进行推理。

OUTPUT_IMAGE_PATH = "./YOLOX_outputs/cots_config/vis_res/0-4614.jpg" 
Image.open(OUTPUT_IMAGE_PATH)

最后设置为实际的输出图像路径。

6.2 使用自定义脚本进行推理（用于 COTS 推理部分）

6.2.1 设置模型

进行YOLOX模型推理所需的实验配置、模型参数和推理参数、加载YOLOX模型。

from yolox.utils import postprocess
from yolox.data.data_augment import ValTransform

COCO_CLASSES = (
  "starfish",
)

# get YOLOX experiment
current_exp = importlib.import_module('cots_config')
exp = current_exp.Exp()

# set inference parameters
test_size = (960, 960)
num_classes = 1
confthre = 0.1
nmsthre = 0.45


# get YOLOX model
model = exp.get_model()
model.cuda()
model.eval()

# get custom trained checkpoint
ckpt_file = "./YOLOX_outputs/cots_config/best_ckpt.pth"
ckpt = torch.load(ckpt_file, map_location="cpu")
model.load_state_dict(ckpt["model"])

6.2.2 INFERENCE BBOXES

定义函数 yolox_inference，用于对输入图像进行YOLOX推理并返回检测到的边界框、类别和得分。

def yolox_inference(img, model, test_size): 
    bboxes = []
    bbclasses = []
    scores = []
    
    preproc = ValTransform(legacy = False)

    tensor_img, _ = preproc(img, None, test_size)
    tensor_img = torch.from_numpy(tensor_img).unsqueeze(0)
    tensor_img = tensor_img.float()
    tensor_img = tensor_img.cuda()

    with torch.no_grad():
        outputs = model(tensor_img)
        outputs = postprocess(
                    outputs, num_classes, confthre,
                    nmsthre, class_agnostic=True
                )

    if outputs[0] is None:
        return [], [], []
    
    outputs = outputs[0].cpu()
    bboxes = outputs[:, 0:4]

    bboxes /= min(test_size[0] / img.shape[0], test_size[1] / img.shape[1])
    bbclasses = outputs[:, 6]
    scores = outputs[:, 4] * outputs[:, 5]
    
    return bboxes, bbclasses, scores

定义函数yolox_inference，它接受三个参数：img（输入图像），model（YOLOX模型），test_size（测试尺寸）。
同时创建了空列表bboxes、bbclasses和scores，用于存储检测结果。

创建了一个ValTransform对象preproc，用于对输入图像进行预处理转换。

使用preproc对象对输入图像进行预处理，将其转换为模型输入的格式，并返回处理后的图像和一个空标签（在推理阶段不需要标签）。

将处理后的图像转换为PyTorch张量，并进行类型转换和将张量移动到GPU上。

使用模型对图像进行推理，得到输出outputs，然后通过postprocess函数对输出进行后处理，根据设定的置信度阈值和NMS阈值进行处理。

检查是否检测到目标，如果没有检测到目标，则返回空的边界框、类别和得分列表。从检测结果中提取边界框信息。将边界框坐标按比例缩放回原始图像尺寸。提取类别和得分信息。将边界框、类别和得分列表作为函数的输出返回。

6.2.3 绘制结果

def draw_yolox_predictions(img, bboxes, scores, bbclasses, confthre, classes_dict):
    for i in range(len(bboxes)):
            box = bboxes[i]
            cls_id = int(bbclasses[i])
            score = scores[i]
            if score < confthre:
                continue
            x0 = int(box[0])
            y0 = int(box[1])
            x1 = int(box[2])
            y1 = int(box[3])

            cv2.rectangle(img, (x0, y0), (x1, y1), (0, 255, 0), 2)
            cv2.putText(img, '{}:{:.1f}%'.format(classes_dict[cls_id], score * 100), (x0, y0 - 3), cv2.FONT_HERSHEY_PLAIN, 0.8, (0,255,0), thickness = 1)
    return img

绘制YOLOX检测结果的函数draw_yolox_predictions。

定义了一个函数draw_yolox_predictions，它接受参数img（输入图像），bboxes（边界框列表），scores（得分列表），bbclasses（类别列表），confthre（置信度阈值），classes_dict（类别字典）。

遍历每个边界框，并提取相应的边界框坐标、类别ID和得分。如果得分低于置信度阈值，则继续下一个循环。然后，将边界框坐标转换为整数类型。

使用OpenCV函数在图像上绘制矩形框和类别标签。矩形框的坐标由边界框的左上角和右下角确定。类别标签包括类别名称和对应的得分百分比。绘制了边界框和标签的图像作为函数的输出返回。

6.2.4 整合

TEST_IMAGE_PATH = "/working/dataset/images/val2017/0-4614.jpg"
img = cv2.imread(TEST_IMAGE_PATH)

# Get predictions
bboxes, bbclasses, scores = yolox_inference(img, model, test_size)

# Draw predictions
out_image = draw_yolox_predictions(img, bboxes, scores, bbclasses, confthre, COCO_CLASSES)

# Since we load image using OpenCV we have to convert it 
out_image = cv2.cvtColor(out_image, cv2.COLOR_BGR2RGB)
display(Image.fromarray(out_image))

这段代码执行以下操作：

1. 通过cv2.imread函数读取测试图像，并将其存储在变量img中。
2. 使用yolox_inference函数对图像进行YOLOX推断，得到边界框（bboxes）、类别（bbclasses）和得分（scores）。
3. 调用draw_yolox_predictions函数，将边界框、类别和得分绘制在图像上，得到输出图像out_image。
4. 由于使用OpenCV加载图像，需要将输出图像从BGR颜色空间转换为RGB颜色空间。
5. 最后，使用display函数显示输出图像。

（七）评估

import greatbarrierreef

env = greatbarrierreef.make_env()   # initialize the environment
iter_test = env.iter_test()  

submission_dict = {
    'id': [],
    'prediction_string': [],
}

for (image_np, sample_prediction_df) in iter_test:
 
    bboxes, bbclasses, scores = yolox_inference(image_np, model, test_size)
    
    predictions = []
    for i in range(len(bboxes)):
        box = bboxes[i]
        cls_id = int(bbclasses[i])
        score = scores[i]
        if score < confthre:
            continue
        x_min = int(box[0])
        y_min = int(box[1])
        x_max = int(box[2])
        y_max = int(box[3])
        
        bbox_width = x_max - x_min
        bbox_height = y_max - y_min
        
        predictions.append('{:.2f} {} {} {} {}'.format(score, x_min, y_min, bbox_width, bbox_height))
    
    prediction_str = ' '.join(predictions)
    sample_prediction_df['annotations'] = prediction_str
    env.predict(sample_prediction_df)

    print('Prediction:', prediction_str)

对一批测试图像进行目标检测，并将检测结果格式化为预测字符串。将预测字符串与对应的样本预测数据框结合起来，并通过env.predict函数提交给环境进行处理。

创建submission_dict的字典，其中包含两个键值对：id和prediction_string。这个字典将用于存储最终的提交结果。

通过一个for循环遍历iter_test。每次迭代，获得image_np和sample_prediction_df

使用yolox_inference函数对image_np进行推理，该函数可能是一个使用YOLOX模型进行目标检测的函数。它返回检测到的边界框（bboxes）、类别标签（bbclasses）和置信度分数（scores）。

使用 for 循环遍历检测到的边界框。对于每个边界框，将获取其坐标、类别标签和置信度分数。如果置信度分数低于confthre可能是一个阈值），则跳过该边界框。

计算边界框的宽度和高度，并将其与置信度分数一起格式化为字符串。将该字符串添加到predictions列表中。

在完成所有边界框的处理后，使用空格将predictions列表中的所有字符串连接起来，形成一个预测字符串prediction_str。

将prediction_str赋值给sample_prediction_df中的一个名为annotations的列，用于存储预测结果。调用env.predict函数，并将sample_prediction_df作为参数传递给它，以将预测结果发送到环境中进行处理。同时打印出预测字符串prediction_str。

sub_df = pd.read_csv('submission.csv')
sub_df.head()