————>第一部分:训练自定义模型
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
!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'
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
类,并设置 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(...))
计算训练集文件夹和验证集文件夹中的文件数量,并将结果打印出来,显示训练集和验证集中的图像文件数量。
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
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)。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
文件的内容,以确认文件是否正确创建并包含了所需的类别信息。
预训练模型列表:
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
。执行完整段代码后,模型将开始训练,并根据配置文件中的设置进行优化和保存模型参数。
# ————>第二部分:推理
%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)
最后设置为实际的输出图像路径。
进行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"])
定义函数 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阈值进行处理。
检查是否检测到目标,如果没有检测到目标,则返回空的边界框、类别和得分列表。从检测结果中提取边界框信息。将边界框坐标按比例缩放回原始图像尺寸。提取类别和得分信息。将边界框、类别和得分列表作为函数的输出返回。
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函数在图像上绘制矩形框和类别标签。矩形框的坐标由边界框的左上角和右下角确定。类别标签包括类别名称和对应的得分百分比。绘制了边界框和标签的图像作为函数的输出返回。
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))
这段代码执行以下操作:
cv2.imread
函数读取测试图像,并将其存储在变量img
中。yolox_inference
函数对图像进行YOLOX推断,得到边界框(bboxes
)、类别(bbclasses
)和得分(scores
)。draw_yolox_predictions
函数,将边界框、类别和得分绘制在图像上,得到输出图像out_image
。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()