【代码解读】RRNet: A Hybrid Detector for Object Detection in Drone-captured Images
文章目录
-
- 1. train.py
- 2. DistributedWrapper类
-
- 2.1 init函数
- 2.2 train函数
- 2.3 dist_training_process函数
- 3. RRNetOperator类
-
- 3.1 init函数
-
- 3.1.1 make_dataloader函数
- 3.2 training_process函数
-
- 3.2.1 criterion函数
- 4. RRNet类（网络模型类）
-
- 4.1 init函数
-
- 4.1.1 get_backbone函数
- 4.1.2 CenterNetDetector类
- 4.1.3 FasterRCNNDetector类
- 4.2 forward函数
-
- 4.2.1 forward_stage1函数
- 4.2.2 transform_bbox函数
- 4.2.3 forward_stage2函数
- 5. DronesDET类（数据集类）
-
- 5.1 init函数
-
- 5.1.1 self.transforms组合类
-
- 5.1.1.1 FillDuck类
- 5.2 __getitem__函数
1. train.py
首先我们将代码从GitHub上下载下来：代码地址
找到程序的主入口train.py这个类，可以看到这个类比较简单，大部分是引用其他类。具体每一个类的定义可以从不同小节中查看
from configs.rrnet_config import Config
from operators.distributed_wrapper import DistributedWrapper
from operators.rrnet_operator import RRNetOperator
if __name__ == '__main__':
dis_operator = DistributedWrapper(Config, RRNetOperator) 详见 2 节
dis_operator.train()
print('Training is Done!')
2. DistributedWrapper类
2.1 init函数
首先来看这个类的初始化函数
def __init__(self, cfg, operator_class):
"""
This is a wrapper class for distributed training.
:param cfg: configuration.
:param operator_class: We use this class to construct the operator for training and evaluating.
"""
self.cfg = cfg
self.operator_class = operator_class
这是一个用于分布式训练的包装器（Wrapper）类。它用于在分布式环境下进行训练。
构造函数中的参数说明如下：
cfg: 表示配置参数，用于设置训练过程中的各种参数和超参数。
operator_class: 这是一个类（Class），用于构造训练和评估操作符（Operator）
2.2 train函数
def train(self):
"""
Start multiprocessing training.
"""
self.setup_distributed_params()
mp.spawn(self.dist_training_process, nprocs=self.cfg.Distributed.ngpus_per_node,
args=(self.cfg.Distributed.ngpus_per_node, self.cfg))
mp.spawn 函数用于启动多个训练进程，并在每个进程中调用 self.dist_training_process 方法
nprocs 参数表示启动的进程数，即用于分布式训练的GPU数量（或进程数量）
args 参数是传递给每个进程的参数，这里传递了 self.cfg.Distributed.ngpus_per_node 和 self.cfg。
2.3 dist_training_process函数
def dist_training_process(self, gpu, ngpus_per_node, cfg):
operator = self.init_operator(gpu, ngpus_per_node, cfg)
operator.training_process()
来看一下 init_operator 函数
def init_operator(self, gpu, ngpus_per_node, cfg):
"""
Create distributed model operator.
:param gpu: gpu id.
:param ngpus_per_node: to calculate the real rank.
:param cfg: configuration.
:return: model operator.
"""
cfg.Distributed.gpu_id = gpu
print("=> Use GPU: {}".format(gpu))
# I. Init distributed process group.
cfg.Distributed.rank = cfg.Distributed.rank * ngpus_per_node + gpu
dist.init_process_group(backend='nccl', init_method=cfg.Distributed.dist_url,
world_size=cfg.Distributed.world_size, rank=cfg.Distributed.rank)
torch.cuda.set_device(gpu)
# II. Init operator.
return self.operator_class(cfg)
首先将当前进程的GPU编号 gpu 赋值给配置参数 cfg.Distributed.gpu_id，用于指定当前进程使用的GPU
然后，根据当前进程的GPU编号和 ngpus_per_node 计算当前进程的真实排名（rank），赋值给配置参数 cfg.Distributed.rank。排名是用于在分布式训练中标识不同进程的标识符，每个进程都有唯一的排名。
接下来，通过调用 dist.init_process_group 方法初始化分布式进程组
随后，通过 torch.cuda.set_device(gpu) 将当前进程的GPU设备设置为 gpu，以确保模型和数据存储在正确的GPU上。
最后，通过调用 self.operator_class(cfg) 创建并初始化模型操作符，并将其返回。
init_operator 的返回值是RRNetOperator类，紧接着调用operator.training_process()进行训练，所以需要查看RRNetOperator的定义（详见 3 节）。
3. RRNetOperator类
3.1 init函数
def __init__(self, cfg):
self.cfg = cfg
model = RRNet(cfg).cuda(cfg.Distributed.gpu_id)
model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
self.optimizer = optim.Adam(model.parameters(), lr=cfg.Train.lr)
self.lr_sch = optim.lr_scheduler.MultiStepLR(self.optimizer, milestones=cfg.Train.lr_milestones, gamma=0.1)
self.training_loader, self.validation_loader = make_dataloader(cfg, collate_fn='rrnet')
super(RRNetOperator, self).__init__(cfg=self.cfg, model=model, lr_sch=self.lr_sch)
# TODO: change it to our class
self.hm_focal_loss = FocalLossHM()
self.l1_loss = RegL1Loss()
self.main_proc_flag = cfg.Distributed.gpu_id == 0
初始化 RRNet 模型，并将其移动到 cfg.Distributed.gpu_id 指定的GPU上 （详见 4 节）
将模型中的 BatchNorm 层转换为同步 BatchNorm，以便在分布式训练中使用
初始化 Adam 优化器，用于更新模型参数。
初始化学习率调度器，用于调整优化器的学习率
初始化训练数据加载器和验证数据加载器，用于加载训练和验证数据 （详见3.1.1）
调用父类 BaseOperator 的构造函数，传递配置参数、模型和学习率调度器。
初始化热图的 Focal Loss，用于计算热图的损失函数。
初始化回归损失函数，用于计算目标的回归损失
判断当前进程是否为主进程（即 GPU 编号为 0 的进程），如果是主进程，则设置 self.main_proc_flag 为 True，否则为 False。
3.1.1 make_dataloader函数
datasets = {
'drones_det': DronesDET
}
def make_dataloader(cfg, collate_fn=None):
if cfg.dataset not in datasets:
raise NotImplementedError
train_dataset = datasets[cfg.dataset](root_dir=cfg.data_root, transforms=cfg.Train.transforms, split='train',
with_road_map=cfg.Train.with_road) （详见 5 节）
val_dataset = datasets[cfg.dataset](root_dir=cfg.data_root, transforms=cfg.Val.transforms, split='val')
if collate_fn is 'ctnet':
collate_fn = train_dataset.collate_fn_ctnet
elif collate_fn is 'rrnet':
collate_fn = train_dataset.collate_fn_ctnet
else:
collate_fn = train_dataset.collate_fn
train_loader = _Dataloader(DataLoader(train_dataset,
batch_size=cfg.Train.batch_size, num_workers=cfg.Train.num_workers,
sampler=cfg.Train.sampler(train_dataset) if cfg.Train.sampler else None,
pin_memory=True, collate_fn=collate_fn,
shuffle=True if cfg.Train.sampler is None else False))
val_loader = DataLoader(val_dataset,
batch_size=cfg.Val.batch_size, num_workers=cfg.Val.num_workers,
sampler=cfg.Val.sampler(val_dataset) if cfg.Val.sampler else None,
pin_memory=True, collate_fn=train_dataset.collate_fn,
shuffle=True if cfg.Val.sampler is None else False)
return train_loader, val_loader
根据配置参数 cfg.dataset 确定数据集的名称，并检查数据集是否在 datasets 字典中注册
根据配置参数创建训练和验证数据集 train_dataset 和 val_dataset
根据 collate_fn 的值确定使用哪个数据集的 collate_fn
如果 collate_fn 为 'ctnet' 或 'rrnet'：则使用相应数据集的 collate_fn_ctnet 方法
否则使用数据集的默认 collate_fn 方法
创建训练数据加载器 train_loader 和验证数据加载器 val_loader
最后，返回创建的训练和验证数据加载器 train_loader 和 val_loader
3.2 training_process函数
def training_process(self):
if self.main_proc_flag:
logger = Logger(self.cfg)
self.model.train()
total_loss = 0
total_hm_loss = 0
total_wh_loss = 0
total_off_loss = 0
total_s2_reg_loss = 0
for step in range(self.cfg.Train.iter_num):
self.lr_sch.step()
self.optimizer.zero_grad()
try:
imgs, annos, gt_hms, gt_whs, gt_inds, gt_offsets, gt_reg_masks, names = self.training_loader.get_batch()
targets = gt_hms, gt_whs, gt_inds, gt_offsets, gt_reg_masks, annos
except RuntimeError as e:
if 'out of memory' in str(e):
print("WARNING: ran out of memory with exception at step {}.".format(step))
continue
outs = self.model(imgs)
targets = gt_hms, gt_whs, gt_inds, gt_offsets, gt_reg_masks, annos
hm_loss, wh_loss, offset_loss, s2_reg_loss = self.criterion(outs, targets)
if step < 2000:
s2_factor = 0
else:
s2_factor = 1
loss = hm_loss + (0.1 * wh_loss) + offset_loss + s2_reg_loss*s2_factor
loss.backward()
self.optimizer.step()
total_loss += float(loss)
total_hm_loss += float(hm_loss)
total_wh_loss += float(wh_loss)
total_off_loss += float(offset_loss)
total_s2_reg_loss += float(s2_reg_loss)
if self.main_proc_flag:
if step % self.cfg.Train.print_interval == self.cfg.Train.print_interval - 1:
# Loss
for param_group in self.optimizer.param_groups:
lr = param_group['lr']
log_data = {'scalar': {
'train/total_loss': total_loss / self.cfg.Train.print_interval,
'train/hm_loss': total_hm_loss / self.cfg.Train.print_interval,
'train/wh_loss': total_wh_loss / self.cfg.Train.print_interval,
'train/off_loss': total_off_loss / self.cfg.Train.print_interval,
'train/s2_reg_loss': total_s2_reg_loss / self.cfg.Train.print_interval,
'train/lr': lr
}}
# Generate bboxs
s1_pred_bbox, s2_pred_bbox = self.generate_bbox(outs, batch_idx=0)
# Visualization
img = (denormalize(imgs[0].cpu()).permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
# Do nms
s2_pred_bbox = self._ext_nms(s2_pred_bbox)
#
s1_pred_on_img = visualize(img.copy(), s1_pred_bbox, xywh=True, with_score=True)
s2_pred_on_img = visualize(img.copy(), s2_pred_bbox, xywh=True, with_score=True)
gt_img = visualize(img.copy(), annos[0, :, :6], xywh=False)
s1_pred_on_img = torch.from_numpy(s1_pred_on_img).permute(2, 0, 1).unsqueeze(0).float() / 255.
s2_pred_on_img = torch.from_numpy(s2_pred_on_img).permute(2, 0, 1).unsqueeze(0).float() / 255.
gt_on_img = torch.from_numpy(gt_img).permute(2, 0, 1).unsqueeze(0).float() / 255.
log_data['imgs'] = {'Train': [s1_pred_on_img, s2_pred_on_img, gt_on_img]}
logger.log(log_data, step)
total_loss = 0
total_hm_loss = 0
total_wh_loss = 0
total_off_loss = 0
total_s2_reg_loss = 0
if step % self.cfg.Train.checkpoint_interval == self.cfg.Train.checkpoint_interval - 1 or \
step == self.cfg.Train.iter_num - 1:
self.save_ckp(self.model.module, step, logger.log_dir)
判断当前进程是否是主进程，如果是则初始化一个记录器，用于记录训练过程和指标。
将模型设置为训练模式
初始化变量以跟踪训练过程中的总损失和不同损失组件 total_loss，total_hm_loss，total_wh_loss，total_off_loss，total_s2_reg_loss
循环遍历训练步骤（iter_num 是总训练步数）：
self.lr_sch.step()：使用学习率调度器调整学习率。
self.optimizer.zero_grad()：在反向传播之前将所有模型参数的梯度清零。
尝试从训练数据加载器中加载一个批次的训练数据：
self.training_loader.get_batch()：获取一个训练数据批次，包括图像、注释、gt热图、gt宽高、gt索引、gt偏移量、gt区域掩码和图像名称
如果数据加载过程中出现 "out of memory" 错误，捕获错误并跳过下一个训练步骤
通过模型进行前向传播，以获取给定输入图像imgs的预测结果outs （详见4.2节）
将gt_hms, gt_whs, gt_inds, gt_offsets, gt_reg_masks, annos赋值为targets
self.criterion(outs, targets)：计算损失，包括热图损失（hm_loss）、宽高损失（wh_loss）、偏移量损失（offset_loss）和 s2 回归损失（s2_reg_loss）（详见3.3.1节）
在前2000个训练步之前，将 s2_factor 设置为 0，之后设置为 1。它是应用于 s2 回归损失的缩放因子
将损失组件组合在一起以计算用于反向传播的总损失（loss）。
loss.backward()：计算损失相对于模型参数的梯度。
self.optimizer.step()：使用计算得到的梯度更新模型参数。
更新当前迭代的总损失和各个损失组件，包括：total_loss，total_hm_loss，total_wh_loss，total_off_loss，total_s2_reg_loss
如果当前进程是主进程，并且当前步数是打印间隔的最后一步（print_interval 是打印间隔），则执行以下操作：
为每个参数组获取学习率，并保存到 lr 中。
创建一个字典 log_data，用于存储要记录的数据，包括总损失和各个损失组件的平均值以及学习率。
生成预测的边界框 s1_pred_bbox 和 s2_pred_bbox。
将图像从张量转换为NumPy数组，用于可视化。
执行非最大抑制（NMS）算法，筛选出 s2_pred_bbox 中的重叠边界框。
用 visualize 函数，将预测的边界框绘制在图像上，并将结果存储在 s1_pred_on_img 和 s2_pred_on_img 中。
将原始注释（ground truth）绘制在图像上，结果存储在 gt_img 中。
将图像转换回PyTorch张量，并进行相应的归一化操作。
创建一个字典 log_data['imgs'] 来存储生成的图像。这些图像将在日志中记录。
将损失组件的计数器重置为零，以便下一个打印间隔时重新计算平均值。
如果当前步数是保存检查点的间隔的最后一步，或者当前步数是训练的最后一步，则执行以下操作：
调用 self.save_ckp 函数保存模型的检查点
3.2.1 criterion函数
def criterion(self, outs, targets):
s1_hms, s1_whs, s1_offsets, s2_reg, bxyxy, scores, _ = outs
gt_hms, gt_whs, gt_inds, gt_offsets, gt_reg_masks, gt_annos = targets
bs = s1_hms[0].size(0)
hm_loss = 0
wh_loss = 0
off_loss = 0
# I. Stage 1
for s in range(self.cfg.Model.num_stacks):
s1_hm = s1_hms[s]
s1_wh = s1_whs[s]
s1_offset = s1_offsets[s]
s1_hm = torch.clamp(torch.sigmoid(s1_hm), min=1e-4, max=1-1e-4)
# Heatmap Loss
hm_loss += self.hm_focal_loss(s1_hm, gt_hms) / self.cfg.Model.num_stacks
# WH Loss
wh_loss += self.l1_loss(s1_wh, gt_reg_masks, gt_inds, gt_whs) / self.cfg.Model.num_stacks
# OffSet Loss
off_loss += self.l1_loss(s1_offset, gt_reg_masks, gt_inds, gt_offsets) / self.cfg.Model.num_stacks
# II. Stage2 Loss
s2_reg_loss = 0
# Calculate IOU between prediction and bbox
# 1. Transform bbox.
gt_annos[:, :, 2:4] += gt_annos[:, :, 0:2]
for b_idx in range(bs):
batch_flag = bxyxy[:, 0] == b_idx
bbox = bxyxy[batch_flag][:, 1:]
gt_anno = gt_annos[b_idx]
iou = torchvision.ops.box_iou(bbox*self.cfg.Train.scale_factor, gt_anno[:, :4])
max_iou, max_idx = torch.max(iou, dim=1)
pos_idx = max_iou > 0.5
# 2. Regression Loss
if pos_idx.sum() == 0:
pos_idx = torch.zeros_like(max_iou, device=max_iou.device).byte()
pos_idx[0] = 1
pos_factor = 0
else:
pos_factor = 1
gt_reg = self.generate_bbox_target(bbox[pos_idx, :]*self.cfg.Train.scale_factor, gt_anno[max_idx[pos_idx], :4])
s2_reg_loss += F.smooth_l1_loss(s2_reg[batch_flag][pos_idx], gt_reg) * pos_factor / bs
return hm_loss, wh_loss, off_loss, s2_reg_loss
将outs解包为各个阶段的预测结果
将targets解包为真实的标签信息
获取batch size
初始化heatmap、WH和Offset的损失为0
循环遍历网络输出的每个阶段：
获取当前阶段的heatmap、WH和Offset预测结果
对当前阶段的heatmap进行sigmoid激活函数并进行范围截断，避免出现取log时的溢出和计算NaN
计算heatmap损失，使用Focal Loss作为损失函数，并将每个阶段的heatmap损失累加到hm_loss中
计算WH损失，使用平滑L1损失函数，并将每个阶段的WH损失累加到wh_loss中
计算Offset损失，使用平滑L1损失函数，并将每个阶段的Offset损失累加到off_loss中
初始化Stage2的回归损失为0
将真实边界框的坐标从(x_min, y_min, w, h)形式转换为(x_min, y_min, x_max, y_max)形式
循环遍历batch中的每个样本：
从bbox的第一列中得到当前样本的标识
获取当前样本对应的预测边界框
获取当前样本的真实边界框
计算预测边界框和真实边界框之间的IoU
找到每个预测边界框与真实边界框最匹配的IoU和对应的真实边界框索引
找到IoU大于0.5的预测边界框的索引（表示匹配的边界框）
如果没有匹配的边界框，则选择一个预测边界框作为匹配，以确保至少有一个匹配的边界框
并将pos_factor设置为0表示没有匹配的边界框，否则设置为1表示有至少一个匹配的边界框
生成匹配的预测边界框和对应的真实边界框的回归目标
使用平滑L1损失函数计算回归损失，并将每个样本的回归损失累加到s2_reg_loss中。
返回第一阶段的heatmap损失hm_loss，WH损失wh_loss，Offset损失off_loss
和第二阶段的回归损失s2_reg_loss作为损失函数的输出。
4. RRNet类（网络模型类）
4.1 init函数
def __init__(self, cfg):
super(RRNet, self).__init__()
self.num_stacks = cfg.Model.num_stacks
self.num_classes = cfg.num_classes
self.nms_type = cfg.Model.nms_type_for_stage1
self.nms_per_class = cfg.Model.nms_per_class_for_stage1
self.backbone = get_backbone(cfg.Model.backbone, num_stacks=self.num_stacks) 详见4.1.1
self.hm = CenterNetDetector(planes=self.num_classes, num_stacks=self.num_stacks, hm=True) 详见4.1.2
self.wh = CenterNetWHDetector(planes=1, num_stacks=self.num_stacks)
self.offset_reg = CenterNetDetector(planes=2, num_stacks=self.num_stacks)
self.head_detector = FasterRCNNDetector() 详见4.1.3
4.1.1 get_backbone函数
根据配置文件我们可以知道，model的backbone是hourglass
def hourglass_net(num_stacks=2):
"""
Make Hourglass Net.
:param num_stacks: number of stacked blocks.
:return: model
"""
model = HourglassNet(num_stacks=num_stacks)
model.load_state_dict(torch.load('./hourglass.pth'), strict=False)
return model
4.1.2 CenterNetDetector类
class CenterNetDetector(nn.Module):
def __init__(self, planes, hm=True, num_stacks=2):
super(CenterNetDetector, self).__init__()
self.hm = hm
self.num_stacks = num_stacks
self.detect_layer = nn.ModuleList([nn.Sequential(
BasicCov(3, 256, 256, with_bn=False),
# BasicCov(3, 40 * (2 ** _), 256, with_bn=False),
nn.Conv2d(256, planes, (1, 1))
) for _ in range(self.num_stacks)
])
if self.hm:
for heat in self.detect_layer:
heat[-1].bias.data.fill_(-2.19)
def forward(self, input, index):
output = self.detect_layer[index](input)
return output
在__init__方法中，设置了一些属性：
self.hm: 一个布尔值，表示是否生成热图（heatmap）的预测。如果hm=True，则需要生成热图的预测，否则不需要。
self.num_stacks: 表示堆叠的数量。该属性用于确定需要生成多少个堆叠的预测结果。
创建了一个nn.ModuleList，其中每个元素是一个包含几个层的nn.Sequential对象。
对于每个堆叠，nn.Sequential中包含：
一个BasicCov层，这是一个自定义的卷积层，输入通道为3，输出通道为256。
一个nn.Conv2d层，用于将256通道的特征图输出到指定的planes通道。这里默认为num_classes
如果self.hm为True，则对所有的nn.Conv2d层的bias进行初始化
在forward方法中，输入input和索引index，然后调用相应堆叠的detect_layer，并将input传递给它，得到输出output。该输出表示对应堆叠的检测器的预测结果。
4.1.3 FasterRCNNDetector类
class FasterRCNNDetector(nn.Module):
def __init__(self):
super(FasterRCNNDetector, self).__init__()
self.top_layer = Bottleneck(inplanes=256, planes=64)
self.regressor = nn.Conv2d(256, 4, kernel_size=1)
def forward(self, feat):
feat = self.top_layer(feat)
feat = F.adaptive_avg_pool2d(feat, 1)
reg = self.regressor(feat)
reg = reg.view(reg.size(0), reg.size(1))
return reg
在__init__方法中，创建了两个成员变量：
self.top_layer：表示Faster R-CNN中的顶层特征层。这里采用了Bottleneck作为顶层特征层。
Bottleneck是一个自定义的卷积层，其参数inplanes=256表示输入通道数为256，planes=64表示输出通道数为64。
self.regressor：表示回归层，用于预测目标框的边界框坐标。
nn.Conv2d(256, 4, kernel_size=1)定义了一个卷积层，输入通道数为256，输出通道数为4，即每个目标框有4个边界坐标。
在forward方法中，输入feat是从CenterNet中传递过来的特征图。
首先，将feat传递给self.top_layer，得到顶层特征层feat。
对feat进行自适应平均池化（adaptive average pooling）操作，将其尺寸调整为1x1，以得到一个固定大小的特征向量。
将特征向量传递给self.regressor，进行回归操作，得到目标框的边界框坐标预测
将预测结果展平为(batch_size, 4)的形状，其中4表示每个目标框的边界框坐标信息
返回边界框坐标预测reg
4.2 forward函数
def forward(self, x, k=1500):
# I. Forward Backbone
pre_feat = self.backbone(x)
# II. Forward Stage 1 to generate heatmap, wh and offset.
hms, whs, offsets = self.forward_stage1(pre_feat) 详见4.2.1
# III. Generate the true xywh for Stage 1.
bboxs = self.transform_bbox(hms[-1], whs[-1], offsets[-1], k) # (bs, k, 6) 详见4.2.2
# IV. Stage 2.
bxyxys = []
scores = []
clses = []
for b_idx in range(bboxs.size(0)):
# Do nms
bbox = bboxs[b_idx]
bbox = self.nms(bbox)
xyxy = bbox[:, :4]
scores.append(bbox[:, 4])
clses.append(bbox[:, 5])
batch_idx = torch.ones((xyxy.size(0), 1), device=xyxy.device) * b_idx
bxyxy = torch.cat((batch_idx, xyxy), dim=1)
bxyxys.append(bxyxy)
bxyxys = torch.cat(bxyxys, dim=0)
scores = torch.cat(scores, dim=0)
clses = torch.cat(clses, dim=0)
# Generate the ROIAlign features.
roi_feat = torchvision.ops.roi_align(torch.relu(pre_feat[-1]), bxyxys, (3, 3))
# Forward Stage 2 to predict and wh offset.
stage2_reg = self.forward_stage2(roi_feat) 详见4.2.3
return hms, whs, offsets, stage2_reg, bxyxys, scores, clses
首先，通过self.backbone(x)调用网络的backbone部分来对输入x进行前向传播，得到pre_feat。
然后，调用self.forward_stage1(pre_feat)来将pre_feat传递给Stage 1，以生成预测的热图（heatmap）、宽高（wh）和偏移（offsets）。这些预测存储在hms、whs和offsets变量中。
接下来，通过调用self.transform_bbox(hms[-1], whs[-1], offsets[-1], k)
对Stage 1的输出进行后处理，以生成真实的边界框坐标。这些边界框存储在变量bboxs中。
然后，对每个边界框进行非极大值抑制（NMS），以去除冗余的预测框。
处理后的边界框存储在变量bxyxys中，其中包含边界框的坐标（xyxy）、得分和类别
使用torchvision.ops.roi_align函数，将pre_feat[-1]和bxyxys作为输入，生成ROIAlign特征roi_feat。
最后，将roi_feat传递给Stage 2，即调用self.forward_stage2(roi_feat)，以预测边界框的宽高和偏移
将预测结果以元组的形式返回：hms、whs、offsets、stage2_reg、bxyxys、scores和clses。
4.2.1 forward_stage1函数
def forward_stage1(self, feats):
hms = []
whs = []
offsets = []
for i in range(self.num_stacks):
feat = feats[i]
feat = torch.relu(feat)
hm = self.hm(feat, i)
wh = self.wh(feat, i)
offset = self.offset_reg(feat, i)
hms.append(hm)
whs.append(wh)
offsets.append(offset)
return hms, whs, offsets
创建三个空列表：hms、whs和offsets
用for循环遍历feats中的每个特征图，并进行以下操作：
通过torch.relu(feat)将特征图进行ReLU激活。
将ReLU激活后的特征图传递给self.hm，并传递堆叠的索引i，得到热图预测hm。
将ReLU激活后的特征图传递给self.wh，并传递堆叠的索引i，得到宽高预测wh。
将ReLU激活后的特征图传递给self.offset_reg，并传递堆叠的索引i，得到偏移预测offset。
将每个堆叠的热图、宽高和偏移预测分别添加到对应的列表hms、whs和offsets中
最后，将三个列表hms、whs和offsets作为结果返回，这些列表分别包含了不同堆叠的热图、宽高和偏移预测结果
4.2.2 transform_bbox函数
def transform_bbox(self, hm, wh, offset, k=250):
batchsize, cls_num, h, w = hm.size()
hm = torch.sigmoid(hm)
scores, inds, clses, ys, xs = self._topk(hm, k)
offset = self._transpose_and_gather_feat(offset, inds)
offset = offset.view(batchsize, k, 2)
xs = xs.view(batchsize, k, 1) + offset[:, :, 0:1]
ys = ys.view(batchsize, k, 1) + offset[:, :, 1:2]
wh = self._transpose_and_gather_feat(wh, inds).clamp(min=0)
wh = wh.view(batchsize, k, 2)
clses = clses.view(batchsize, k, 1).float()
scores = scores.view(batchsize, k, 1)
pred_x = (xs - wh[..., 0:1] / 2)
pred_y = (ys - wh[..., 1:2] / 2)
pred_w = wh[..., 0:1]
pred_h = wh[..., 1:2]
pred = torch.cat([pred_x, pred_y, pred_w + pred_x, pred_h + pred_y, scores, clses], dim=2)
return pred
对热图hm应用Sigmoid激活函数，将其转换为概率值，表示每个像素点是目标的概率。
调用_topk函数，从热图中选取前k个最高概率的像素点，并获取这些像素点的坐标、类别、分数等信息。这个函数用于筛选预测结果。
对偏移offset进行变换和采样，将其应用到对应的高分概率像素点的坐标上，得到修正后的目标中心点坐标。
对宽高wh进行变换和采样，将其应用到对应的高分概率像素点上，并取值大于等于零，确保预测的宽高是非负的。
将预测的中心点坐标和宽高信息拼接在一起，形成最终的边界框预测结果。
返回包含边界框预测信息的pred
4.2.3 forward_stage2函数
def forward_stage2(self, feats,):
stage2_reg = self.head_detector(feats)
return stage2_reg
5. DronesDET类（数据集类）
5.1 init函数
def __init__(self, root_dir, transforms=None, split='train', with_road_map=False):
'''
:param root_dir: root of annotations and image dirs
:param transform: Optional transform to be applied
on a sample.
'''
# get the csv
self.images_dir = os.path.join(root_dir, split, 'images')
self.annotations_dir = os.path.join(root_dir, split, 'annotations')
self.roadmap_dir = os.path.join(root_dir, split, 'roadmap')
mdf = os.listdir(self.images_dir)
restr = r'\w+?(?=(.jpg))'
for index, mm in enumerate(mdf):
mdf[index] = re.match(restr, mm).group()
self.mdf = mdf
self.transforms = transforms
self.with_road_map = with_road_map
根据root_dir和split参数构建了指向'images'目录的路径。
根据root_dir和split参数构建了指向'annotations'目录的路径
根据root_dir和split参数构建了指向'roadmap'目录的路径
列出了'images'目录中的所有文件，并将它们赋值给变量mdf。
定义了一个正则表达式模式。用于匹配文件名中的字母数字字符（和下划线）
定义一个循环，它遍历mdf列表中的每个元素
使用re.match函数将正则表达式模式（restr）应用于当前文件名（mm），提取文件名中的字母数字部分（不包括'.jpg'扩展名），并将其重新赋值给mdf列表的对应索引。
循环结束后，将只包含文件名（不带'.jpg'）的修改后的mdf列表赋值给实例变量self.mdf。
将传递给构造方法的transforms参数赋值给实例变量self.transforms
将传递给构造方法的with_road_map参数赋值给实例变量self.with_road_map （这里默认是true）
5.1.1 self.transforms组合类
查看self.transforms的具体定义
Config.Train.transforms = Compose([
MultiScale(scale=(1, 1.15, 1.25, 1.35, 1.5)),
ToTensor(),
MaskIgnore(Config.Train.mean),
FillDuck(),
HorizontalFlip(),
RandomCrop(Config.Train.crop_size),
Normalize(Config.Train.mean, Config.Train.std),
ToHeatmap(scale_factor=Config.Train.scale_factor)
])
MultiScale是一个多尺度缩放转换。它将图像按照指定的尺度因子进行多次缩放，以增加训练数据的多样性
ToTensor将图像和注释数据转换为张量形式
MaskIgnore是一个mask忽略转换。它使用指定的均值（Config.Train.mean）来标记忽略区域
FillDuck这是一个填充“Duck”的转换 （论文中的数据增强，详见5.1.1.1）
HorizontalFlip这是一个水平翻转转换。它以一定的概率水平翻转图像，从而增加数据的多样性。
RandomCrop(Config.Train.crop_size)是一个随机裁剪转换。它将图像随机裁剪到指定的尺寸
Normalize(Config.Train.mean, Config.Train.std)是一个图像归一化转换。它将图像像素值标准化为均值为Config.Train.mean，标准差为Config.Train.std的数据
ToHeatmap(scale_factor=Config.Train.scale_factor)是一个转换，将图像数据转换为热图（heatmap）数据。热图常用于一些特定的目标检测或姿态估计任务，用于标记目标的位置或关键点。
5.1.1.1 FillDuck类
class FillDuck(object):
def __init__(self, cls_list=(1, 2, 3, 7, 8, 10), factor=0.00005):
self.cls_list = torch.tensor(cls_list).unsqueeze(0)
self.factor = factor
def __call__(self, data):
return F.fill_duck(data, self.cls_list, self.factor)
cls_list 是一个包含需要填充的目标类别的列表，默认包含类别 1、2、3、7、8 和 10。（论文中提到的类别）
factor 是一个填充因子，用于控制填充的程度，默认为 0.00005
接下来来看fill_duck的具体定义
def fill_duck(data, cls_list, factor):
try:
img, annos, roadmap = data
# I. Get valid area.
valid_idx = roadmap.view(-1)
idx = torch.nonzero(valid_idx).view(-1)
if idx.size(0) == 0:
return img, annos
xs = idx % roadmap.size(1)
ys = idx // roadmap.size(1)
coor = torch.stack((xs, ys), dim=1)
annos_cls = annos[:, 5]
从data中解包出图像、注释和roadmap数据，分别赋值给img、annos和roadmap。
将roadmap数据展平为一维张量，valid_idx中的元素是原始roadmap图像中每个像素的值。
通过torch.nonzero函数找到valid_idx中非零元素的索引，即有效区域的索引。然后使用view(-1)将索引展平为一维张量。
如果有效区域中的像素数量为0（即没有有效区域）
则直接返回原始图像和注释数据，不进行后续的处理。
计算有效区域中每个像素的x坐标
计算有效区域中每个像素的y坐标
将x坐标和y坐标合并为一个坐标张量coor，其中每一行包含一个有效像素的(x, y)坐标。
从注释数据annos中提取出目标类别信息
# II Calculate scale factor for depth.
people_flag = annos_cls == 1
people_bbox = annos[people_flag, :4]
if people_bbox.size(0) != 0:
people_diag = people_bbox[:, 2:4].pow(2).sum(dim=1).sqrt()
topk = min(3, people_diag.size(0))
max_diag, max_idx = torch.topk(people_diag, k=topk)
min_diag, min_idx = torch.topk(people_diag, k=1, largest=False)
y_diff = people_bbox[max_idx, 1] - people_bbox[min_idx, 1]
scale_factor = ((max_diag - min_diag) / (y_diff.abs() + 1e-5)).mean()
else:
scale_factor = 1
创建了一个布尔索引，用于选择目标类别为1的目标
使用布尔索引people_flag来选择目标类别为1的目标的边界框信息，用people_flag选择出这些目标的前4列，即包含边界框的左上角坐标和右下角坐标的信息。
判断是否存在目标类别为1的目标
计算目标类别为1的目标框的对角线长度
取其右下角坐标减去左上角坐标得到边界框的宽和高，然后使用勾股定理计算对角线长度。
确定了最大尺度因子的计算个数
找到目标类别为1的目标中，对角线长度最大的k个目标，并返回它们的对角线长度和对应的索引。
找到目标类别为1的目标中，对角线长度最小的1个目标，并返回它的对角线长度和对应的索引。
算了目标类别为1的目标中，对角线长度最大和最小的目标的上下边界之间的差值。
计算目标类别为1的目标的尺度因子。
它通过最大和最小对角线长度之间的差值除以上下边界之间的差值得到尺度因子，并取平均值作为最终的尺度因子。
如果目标类别为1的目标不存在（即people_bbox.size(0) == 0），则尺度因子设为1，表示不进行尺度变换。
# III. For relation class.
people_flag = annos_cls == 2
people_select_annos = annos[people_flag, :]
relation_flag = torch.zeros_like(annos_cls).byte()
if people_select_annos.size(0) != 0:
iou = bbox_iou(people_select_annos[:, :4], annos[:, :4], x1y1x2y2=False)
if iou.size(1) > 2:
max_v, max_i = torch.topk(iou, dim=1, k=2)
flag = max_v[:, 1] > 0
max_i = max_i[flag, :]
people_idx = max_i[:, 0]
vechile_idx = max_i[:, 1]
relation_flag[people_idx] = 1
relation_flag[vechile_idx] = 1
创建了一个布尔索引，用于选择目标类别为2的目标
使用布尔索引people_flag来选择目标类别为2的目标的所有信息
创建了一个与annos_cls形状相同的零张量relation_flag ，并将其转换为布尔型
判断是否存在目标类别为2的目标
计算目标类别为2的目标与所有目标之间的IOU（交并比）
判断IOU矩阵的列数是否大于2
找到IOU矩阵中每行的最大和次大的值，并返回它们的值和索引
创建一个布尔索引，用于选择次大的IOU值大于0的行
使用布尔索引flag来选择满足条件的行
分别提取次大IOU值对应的行的第一个索引和第二个索引
将人目标的索引和其他与人目标有关系的目标的索引设置为1
# IV. Calculate aug N.
cls = cls_list.repeat(annos.size(0), 1)
normal_flag = (cls == annos_cls.unsqueeze(1).repeat(1, cls.size(1)).long()).sum(dim=1) > 0
normal_flag = normal_flag * (1 - relation_flag)
total_n = max(int(factor * valid_idx.sum()), 5)
relation_n = relation_flag.float().sum() / 2
normal_n = normal_flag.float().sum()
if relation_n + normal_n == 0:
return img, annos
r_n = int(relation_n / (relation_n + normal_n) * total_n)
n_n = total_n - r_n
将目标类别列表cls_list重复annos.size(0)次，生成一个形状为(annos.size(0), len(cls_list))的张量cls
通过布尔索引生成一个标记向量normal_flag，用于标记目标是否为普通（normal）目标
根据normal_flag和relation_flag的取值，对普通目标的标记向量进行进一步调整
计算总样本数，用于控制数据增强的采样数量
计算关系目标的数量
计算普通目标的数量
判断关系目标和普通目标的数量之和是否为0。如果为0，表示没有需要采样的目标，直接返回原始图像和注释数据
计算关系目标的采样数量
计算普通目标的采样数量
# V. Fill image
paste_idx = torch.randint(low=0, high=coor.size(0), size=(total_n,))
paste_coors = coor[paste_idx]
new_annos = []
# 1. Sample normal object.
if n_n != 0:
normal_annos = annos[normal_flag, :]
sample_idx = torch.randint(low=0, high=normal_annos.size(0), size=(n_n,))
sample_annos = normal_annos[sample_idx]
for i, anno in enumerate(sample_annos):
paste_coor = paste_coors[i].float()
# Apply depth scale.
anno_ct_y = anno[1] + anno[3] / 2
diff = (anno_ct_y - paste_coor[1]).abs() * scale_factor
anno_diag = (anno[2].pow(2) + anno[3].pow(2)).sqrt()
if anno_ct_y > paste_coor[1]:
# Do reduce.
factor = 1 - diff / anno_diag
else:
factor = 1 + diff / anno_diag
cropped_obj = img[:, int(anno[1]):int(anno[1]+anno[3]), int(anno[0]):int(anno[0]+anno[2])]
factor = factor.clamp(min=0.5, max=2)
cropped_obj = F.interpolate(
cropped_obj.unsqueeze(0),
scale_factor=float(factor),
mode='bilinear',
align_corners=True
)[0]
obj_h, obj_w = cropped_obj.size()[-2:]
paste_coor[0] -= obj_w / 2
paste_coor[1] -= obj_h / 2
paste_coor[0] = paste_coor[0].clamp(min=1, max=img.size(2)-obj_w - 1)
paste_coor[1] = paste_coor[1].clamp(min=1, max=img.size(1)-obj_h - 1)
img[:, int(paste_coor[1]):int(paste_coor[1]+obj_h),
int(paste_coor[0]):int(paste_coor[0]+obj_w)] = cropped_obj
new_annos.append(torch.tensor([[int(paste_coor[0]), int(paste_coor[1]), int(obj_w), int(obj_h), anno[4], anno[5], anno[6], anno[7]]]))
生成一个随机索引paste_idx，用于从坐标张量coor中随机采样total_n个坐标。
使用随机索引paste_idx从坐标张量coor中选取对应的坐标，得到paste_coors，即采样得到的随机坐标。
创建一个空列表new_annos，用于存储生成的新的目标注释
判断是否需要对普通目标进行采样
使用布尔索引normal_flag，选择普通目标的注释数据
生成一个随机索引sample_idx，用于从普通目标的注释数据中随机采样n_n个目标。
使用随机索引sample_idx从普通目标的注释数据中选取对应的目标
for循环，遍历随机采样得到的普通目标的注释数据
获取当前目标的随机坐标，将其转换为浮点数类型
计算目标的中心y坐标
计算目标中心y坐标与随机坐标y的差值，并乘以尺度因子scale_factor，用于调整目标的尺度。
计算目标边界框的对角线长度
如果目标中心y坐标大于随机坐标y，说明随机坐标位于目标下方，此时将尺度因子设为1减去差值与对角线长度比例的值。
如果目标中心y坐标小于随机坐标y，说明随机坐标位于目标上方，此时将尺度因子设为1加上差值与对角线长度比例的值。
从原始图像img中裁剪出目标的图像块
将尺度因子限制在0.5到2之间，避免过大或过小的尺度变换
使用双线性插值对目标图像块进行尺度变换
获取经过尺度变换后的目标图像块的高度和宽度
将随机坐标paste_coor的x和y分别减去目标图像块的宽度和高度的一半，将随机坐标对准到目标图像块的中心。
将随机坐标的x和y限制在图像的有效范围内，避免出现坐标越界
将经过尺度变换后的目标图像块插入到原始图像img中的随机坐标位置处
将当前增强后的目标的信息添加到new_annos列表中
# 2. Sample Relation Object.
if r_n != 0:
people_annos = annos[people_idx, :]
vechile_annos = annos[vechile_idx, :]
sample_idx = torch.randint(low=0, high=people_annos.size(0), size=(r_n,))
sample_people_annos = people_annos[sample_idx]
sample_vechile_annos = vechile_annos[sample_idx]
sample_people_annos[:, 2:4] += sample_people_annos[:, 0:2]
sample_vechile_annos[:, 2:4] += sample_vechile_annos[:, 0:2]
for i in range(r_n):
paste_coor = paste_coors[i + n_n].float()
people_anno = sample_people_annos[i]
vechile_anno = sample_vechile_annos[i]
min_x = int(min(people_anno[0], vechile_anno[0]))
min_y = int(min(people_anno[1], vechile_anno[1]))
max_x = int(max(people_anno[2], vechile_anno[2]))
max_y = int(max(people_anno[3], vechile_anno[3]))
# Apply depth scale.
anno_ct_y = (min_y + max_y) / 2
diff = (anno_ct_y - paste_coor[1]).abs() * scale_factor
anno_diag = math.sqrt((max_x-min_x)**2 + (max_y-min_y)**2)
if anno_ct_y > paste_coor[1]:
# Do reduce.
factor = 1 - diff / anno_diag
else:
factor = 1 + diff / anno_diag
cropped_obj = img[:, min_y:max_y, min_x:max_x]
factor = factor.clamp(min=0.5, max=2)
cropped_obj = F.interpolate(
cropped_obj.unsqueeze(0),
scale_factor=float(factor),
mode='bilinear',
align_corners=True
)[0]
obj_h, obj_w = cropped_obj.size()[-2:]
paste_coor[0] -= obj_w / 2
paste_coor[1] -= obj_h / 2
paste_coor[0] = paste_coor[0].clamp(min=1, max=img.size(2)-obj_w - 1)
paste_coor[1] = paste_coor[1].clamp(min=1, max=img.size(1)-obj_h - 1)
img[:, int(paste_coor[1]):int(paste_coor[1]+obj_h),
int(paste_coor[0]):int(paste_coor[0]+obj_w)] = cropped_obj
x_bias = min_x - paste_coor[0]
y_bias = min_y - paste_coor[1]
new_people = people_anno
new_people[2:4] -= new_people[0:2]
new_people[2:4] *= factor
new_people[0] -= x_bias
new_people[1] -= y_bias
new_vechile = vechile_anno
new_vechile[2:4] -= new_vechile[0:2]
new_vechile[2:4] *= factor
new_vechile[0] -= x_bias
new_vechile[1] -= y_bias
new_annos.append(new_people.unsqueeze(0).floor())
new_annos.append(new_vechile.unsqueeze(0).floor())
new_annos = torch.cat(new_annos)
annos = torch.cat((annos, new_annos))
判断是否需要对关系目标进行采样
使用索引people_idx和vechile_idx分别从原始目标注释数据中选择关系目标和与之相关的目标
生成一个随机索引sample_idx，用于从关系目标的注释数据中随机采样r_n个目标
使用随机索引sample_idx从关系目标和与之相关的目标的注释数据中选取对应的目标。
将目标的边界框坐标转换为(x_min, y_min, x_max, y_max)的形式。
遍历关系目标的采样结果
获取当前关系目标的随机坐标，将其转换为浮点数类型
分别获取当前关系目标和与之相关的目标的注释数据
别计算当前目标的左上角x和y坐标
分别计算当前目标的右下角x和y坐标
计算目标的中心y坐标，并计算其与随机坐标y的差值，并乘以尺度因子scale_factor。
计算目标的对角线长度，用于后续计算尺度变换的缩放因子
根据目标的中心y坐标和随机坐标y的关系来选择尺度变换的因子
如果目标的中心y坐标大于随机坐标y，说明随机坐标位于目标下方，此时将尺度因子设为1减去差值与对角线长度比例的值。
如果目标的中心y坐标小于随机坐标y，说明随机坐标位于目标上方，此时将尺度因子设为1加上差值与对角线长度比例的值。
从原始图像img中裁剪出包含目标的图像块
行将尺度因子限制在0.5到2之间，避免过大或过小的尺度变换
使用双线性插值对目标图像块进行尺度变换
获取缩放后的目标图像块的高度和宽度
将随机坐标paste_coor的x和y分别减去目标图像块的宽度和高度的一半，将随机坐标对准到目标图像块的中心。
将随机坐标的x和y限制在图像的有效范围内，避免出现坐标越界
将经过尺度变换后的目标图像块插入到原始图像img中的随机坐标位置处，完成数据增强的操作。
分别计算目标图像块左上角相对于随机坐标的x和y偏移量
分别创建新的张量new_people和new_vechile，用于存储经过尺度变换和偏移后的目标注释信息。
将目标的右下角坐标转换为宽度和高度
将目标的宽度和高度乘以尺度因子，完成尺度变换
将目标的左上角坐标加上x和y偏移量，完成位置偏移
将经过尺度变换和偏移后的人和车辆目标的注释信息添加到new_annos列表中。
使用torch.cat()函数将所有增强后的目标注释信息拼接成一个张量，形状为(N, 8)，N是增强后的目标数量。
将原始目标注释信息和增强后的目标注释信息拼接在一起，形成最终的目标注释信息。
返回增强后的图像img和增强后的目标注释信息annos
return img, annos
except:
return data[0], data[1]
5.2 __getitem__函数
def __getitem__(self, item):
name = self.mdf[item]
img_name = os.path.join(self.images_dir, '{}.jpg'.format(name))
txt_name = os.path.join(self.annotations_dir, '{}.txt'.format(name))
# read image
image = Image.open(img_name).convert("RGB")
# read annotation
annotation = pd.read_csv(txt_name, header=None)
annotation = np.array(annotation)[:, :8]
annotation = annotation[annotation[:, 5] != 11]
# read road segmentation
roadmap = None
if self.with_road_map:
roadmap_name = os.path.join(self.roadmap_dir, '{}.jpg'.format(name))
roadmap = cv2.imread(roadmap_name)
sample = (image, annotation, roadmap)
if self.transforms:
sample = self.transforms(sample)
return sample + (name,)
根据传入的item索引，从self.mdf列表中获取相应的文件名（不包括'.jpg'扩展名）
构建了图像文件的完整路径，用于读取图像数据（加入了后缀名jpg）
构建了注释文件的完整路径，用于读取注释数据
使用PIL库打开图像文件，然后将其转换为RGB格式。Image.open()用于读取图像数据。
使用Pandas库从注释文件中读取CSV格式的注释数据
将读取的注释数据转换为NumPy数组，并保留前8列数据
筛选掉注释中第5列等于11的行。这可能是为了排除某个特定的类别。
创建一个变量roadmap并初始化为None
判断self.with_road_map是否为True：
如果数据集包含roadmap数据，这一行构建了roadmap图像文件的完整路径，用于读取roadmap数据。
数据集包含roadmap数据，则使用OpenCV库读取roadmap图像数据
将图像、注释和roadmap数据（如果有的话）打包成一个元组，并赋值给变量sample
检查self.transforms是否存在（非None）。
如果存在，说明数据集已经定义了数据变换（数据增强等），则将sample应用到这些变换上
将打包好的样本元组返回，并附加文件名（不包括'.jpg'）作为元组的最后一个元素。这样，样本数据和对应的文件名就一并返回了。