小样本计数网络FamNet(Learning To Count Everything)

undo_try2024-04-16 13:12

小样本计数网络FamNet(Learning To Count Everything)

  • 大多数计数方法都仅仅针对一类特定的物体，如人群计数、汽车计数、动物计数等。一些方法可以进行多类物体的计数，但是training set中的类别和test set中的类别必须是相同的。

  • 为了增加计数方法的可拓展性，Viresh Ranjan在2021年CVPR上提出了小样本计数网络(Few-Shot Counting, FSC) FamNet。

  • 论文地址：Learning To Count Everything

1 FamNet网络概述

  • 小样本计数如下图所示，给定一张或几张support描述计数的物体类别，并给定一张待计数的图像query，FSC希望计算出该类别的物体在query中出现的个数。
  • 除了在训练集中出现的类别 (称为base classes)，在测试阶段，FSC还需要处理完全没有见过的类别 (称为novel classes)。

从今天的角度来看，小样本计数的baselines主要类别如下：

1.1 FamNet网络架构

  • 绝大多数计数的方法都只能针对一个类别的情况，造成这种情况的原因有两个：

    • 计数需要很细致的标注
    • 甚至没有包含多个类别的数据集

  • 为了解决第一个问题，将计数看成是few-shot回归任务。

    • 提出了FamNet架构来解决few-shot counting的任务，在测试阶段引入了一个few-shot adaptation scheme提升了模型的性能
    • 值得注意的是，本文在测试过程中选择的是和训练过程中完全不同的类别

  • 为了解决第二个问题自己建立了一个数据集(如下图)：FSC-147（a few-shot counting dataset，有147个类别，超过6000张图片）

    • 本文收集了超过6000张图片，共147个类别的图像，主要由厨房餐具、办公文具、信纸、交通工具、动物等组成。
    • 数量从7-3731，平均每幅图56个目标。每个目标实例用近似中心的点来标注。
    • 另外，三个目标实例随机挑选出来加上矩形框作为目标样例。
    • 注意一下数据集的划分：训练集：3659张，89类；验证集：1286张，29类；测试集：1190张，29类，总计147类，可以发现训练集和测试集的类别不一致。

下图是FamNet的网络架构：

  • 输入：RGB图像以及想要计数的物体的几个sampler bounding boxes(下图中红框标注的鱼，只需要3个左右的bbox信息)；

  • 输出：密度图（和输入图像尺寸是一样的，只有1个channel），根据密度图可以得到想要的object的数量。

  • FamNet主要由两个核心模块组成：一个是多尺度的特征提取模块(绿色矩形)，一个是密度图预测的模块(绿色矩形)。

多尺度的特征提取模块

  • 为了处理大范围的物体类别，本文用在ImageNet上预训练了的网络进行特征提取，并且在训练时候冻结参数，不进行微调。

  • 具体来讲，FamNet用的是ImageNet上pretrained过的ResNet-50的前4个block进行多尺度特征提取，而且这些block的权重在FamNet训练时是不改变的。

  • 用ResNet-50的第3个和第4个block得到的特征图表征输入图像，对这些特征图用ROI pooing(紫色矩阵框)得到exempler的多尺度特征。

    • 在FamNet的源码中，ROI pooing实现其实非常简单，直接用的Pytorch中的插值函数实现。

    • ROI pooing是在Faster RCNN里提出的。Faster RCNN中将RPN生成的候选框投影到特征图上获得相应的特征矩阵后，由于每个特征矩阵的大小不一定是相同的，于是有一个ROI Pooling层，将这些特征矩阵缩放到相同尺度的特征图，Faster RCNN用的是7×7，然后再将特征图展平通过一系列全连接层得到预测结果。

密度图预测模块

  • 为了使模型不是针对特定类的，密度图预测模块的设计是需要类无关的。
  • 那么显然，为了类无关的就不能直接把提取到的图像特征送入密度图预测模块。作者使用了一个correlation map表征exempler特征和输入的图像特征的相似性(上图灰色矩形框)，然后再做density prediction，这里作者使用卷积操作来计算exempler特征和输入图像特征的相似性。
  • 同样，由于object可能会有尺度变化，因此作者把exempler的大小进行缩放（0.9-1.1），并生成多个correlation maps，这些特征图最后被concat后送入密度图预测模块。

1.2 FamNet中的损失函数

  • 在模型训练的时候，使用的损失函数是常规的MSE损失(只有密度图预测模块参与训练)。

  • 在模型预测或测试的时候，用的是Min-Count loss和Perturbation loss的加权和。

  • 可能你会好奇，预测时候还需要计算损失吗？

    • 这是因为训练阶段，我们利用的只是样本外观特征信息。
    • 在推理阶段，作者提出一种新方法(如下图推理加入Adaptation可以提高预测精度，源码中默认训练100个epochs)用于提升估计模块的精度，可以充分利用Bounding boxes样本的位置信息。

最小计数损失(Min-Count loss)
L M i n C o u n t = ∑ b ∈ B m a x ( 0 , 1 − ∣ ∣ Z b ∣ ∣ 1 ) B 为给定所有的样本框， Z b 表示在密度图 Z 上位置 b 处的裁剪图 , ∣ ∣ Z b ∣ ∣ 1 为 Z b 所有值的求和 L_{MinCount}=\sum_{b\in B}max(0, 1-||Z_b||_1) \\ B为给定所有的样本框，Z_b表示在密度图Z上位置b处的裁剪图,||Z_b||_1为Z_b所有值的求和\\ LMinCount=b∈B∑max(0,1−∣∣Zb∣∣1)B为给定所有的样本框，Zb表示在密度图Z上位置b处的裁剪图,∣∣Zb∣∣1为Zb所有值的求和

看了这个损失函数的数学表达式，你可能比较懵逼，那我们直接来看代码。

  • 通过下面代码，我们知道MinCountLoss其实就是为了约束每一个exempler的框内至少有一个物体
  • 当 ∣ ∣ Z b ∣ ∣ 1 > 1 ||Z_b||1>1 ∣∣Zb∣∣1>1时， L M i n C o u n t = 0 L{MinCount}=0 LMinCount=0。因此，最小计数损失只会惩罚那些没有判断出exempler区域内目标大于1的情况。
python 复制代码 
# 主要是为了约束每一个exempler的框内至少有一个物体
def MincountLoss(output, boxes, use_gpu=True):
    ones = torch.ones(1)
    if use_gpu:
        ones = ones.cuda()
    Loss = 0.
    if boxes.shape[1] > 1:
        boxes = boxes.squeeze()
        # 遍历每一个提供的bbox
        for tempBoxes in boxes.squeeze():
            y1 = int(tempBoxes[1])
            y2 = int(tempBoxes[3])
            x1 = int(tempBoxes[2])
            x2 = int(tempBoxes[4])
            # 将bbox区域内密度求和
            X = output[:, :, y1:y2, x1:x2].sum()
            # 小于1就计算loss，目的是确保每一个exempler的框内至少有一个物体
            if X.item() <= 1:
                Loss += F.mse_loss(X, ones)
    else:
        boxes = boxes.squeeze()
        y1 = int(boxes[1])
        y2 = int(boxes[3])
        x1 = int(boxes[2])
        x2 = int(boxes[4])
        X = output[:, :, y1:y2, x1:x2].sum()
        if X.item() <= 1:
            Loss += F.mse_loss(X,ones)  
    return Loss

扰动损失Perturbation loss

  • 密度图 Z Z Z本质上为exemplers和图像的关联响应（卷积）图。

  • 而exemplers周围的密度值理想情况下也会是一个高斯分布，所以不满足这个分布的情况就会有Loss。

  • G h × w G_{h×w} Gh×w为尺寸是 h × w h×w h×w的2D高斯分布图。

  • 预测最终所用的联合适应损失(adaptation loss)是Min-Count loss和Perturbation loss的加权和。

python 复制代码 
def PerturbationLoss(output, boxes, sigma=8, use_gpu=True):
    Loss = 0.
    if boxes.shape[1] > 1:
        boxes = boxes.squeeze()
        # 遍历每一个提供的bbox
        for tempBoxes in boxes.squeeze():
            y1 = int(tempBoxes[1])
            y2 = int(tempBoxes[3])
            x1 = int(tempBoxes[2])
            x2 = int(tempBoxes[4])
            out = output[:, :, y1:y2, x1:x2]
            # 作者认为预测密度图会围绕bbox呈Gaussian分布，所以不满足这个分布的情况就会有Loss
            GaussKernel = matlab_style_gauss2D(shape=(out.shape[2], out.shape[3]), sigma=sigma)
            GaussKernel = torch.from_numpy(GaussKernel).float()
            if use_gpu:
                GaussKernel = GaussKernel.cuda()
            Loss += F.mse_loss(out.squeeze(), GaussKernel)
    else:
        boxes = boxes.squeeze()
        y1 = int(boxes[1])
        y2 = int(boxes[3])
        x1 = int(boxes[2])
        x2 = int(boxes[4])
        out = output[:, :, y1:y2, x1:x2]
        Gauss = matlab_style_gauss2D(shape=(out.shape[2], out.shape[3]), sigma=sigma)
        GaussKernel = torch.from_numpy(Gauss).float()
        if use_gpu:
            GaussKernel = GaussKernel.cuda()
        Loss += F.mse_loss(out.squeeze(), GaussKernel)
    return Loss

实验部分不再介绍，可以参考原文，只贴出一个与其他少样本方法的比较。

2 FamNet源码分析

源码来自官方作者实现：GitHub - cvlab-stonybrook/LearningToCountEverything

下面我们通过源码，来对FamNet有更深的理解。

下面代码在源码基础上增加了中文注释并进行了微调，仓库地址：小样本计数网络FamNet

2.1 模型创建

  • 利用在ImageNet数据集上预训练的ResNet-50模型的layer3、layer4作特征提取
  • 密度预测模块CountRegressor由5个卷积和3个上采样层（上采样2倍）组成
python 复制代码 
# LearningToCountEverything/model.py
class Resnet50FPN(nn.Module):
    def __init__(self):
        super(Resnet50FPN, self).__init__()
        self.resnet = torchvision.models.resnet50(pretrained=True)
        children = list(self.resnet.children())
        self.conv1 = nn.Sequential(*children[:4])
        self.conv2 = children[4]
        self.conv3 = children[5]
        self.conv4 = children[6]

    def forward(self, im_data):
        feat = OrderedDict()
        feat_map = self.conv1(im_data)
        feat_map = self.conv2(feat_map)
        feat_map3 = self.conv3(feat_map)
        feat_map4 = self.conv4(feat_map3)
        feat['map3'] = feat_map3
        feat['map4'] = feat_map4
        return feat

class CountRegressor(nn.Module):
    def __init__(self, input_channels, pool='mean'):
        super(CountRegressor, self).__init__()
        self.pool = pool
        # 经过三次上采样，将feature map采样到原始图像大小
        self.regressor = nn.Sequential(
            nn.Conv2d(in_channels=input_channels, out_channels=196, kernel_size=7, padding=3),
            nn.ReLU(),
            nn.UpsamplingBilinear2d(scale_factor=2),
            nn.Conv2d(in_channels=196, out_channels=128, kernel_size=5, padding=2),
            nn.ReLU(),
            nn.UpsamplingBilinear2d(scale_factor=2),
            nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.UpsamplingBilinear2d(scale_factor=2),
            nn.Conv2d(in_channels=64, out_channels=32, kernel_size=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=32, out_channels=1, kernel_size=1),  # 输出通道为1
            nn.ReLU(),
        )

    def forward(self, im):
        num_sample = im.shape[0]
        if num_sample == 1:
            # im.squeeze(0)先压缩通道0
            output = self.regressor(im.squeeze(0))
            if self.pool == 'mean':
                # 取多个bbox所产生的correlation map的平均值
                output = torch.mean(output, dim=(0), keepdim=True)
                return output
            elif self.pool == 'max':
                output, _ = torch.max(output, 0, keepdim=True)
                return output
        else:
            for i in range(0, num_sample):
                output = self.regressor(im[i])
                if self.pool == 'mean':
                    output = torch.mean(output, dim=(0), keepdim=True)
                elif self.pool == 'max':
                    output, _ = torch.max(output, 0, keepdim=True)
                if i == 0:
                    Output = output
                else:
                    Output = torch.cat((Output, output), dim=0)
            return Output

2.2 模型训练

  • 训练过程中采用MSE损失函数
  • 多尺度特征提取模块使用Resnet50FPN、密度图预测模块使用CountRegressor
  • 训练过程，ResNet-50中block的权重是不改变的，只更新CountRegressor的参数
python 复制代码 
# LearningToCountEverything/train.py
def train(args):
    device = torch.device(args.gpu if torch.cuda.is_available() else "cpu")

    # MSE损失函数
    criterion = nn.MSELoss().to(device)

    # 1、多尺度特征提取模块
    resnet50_conv = Resnet50FPN().to(device)
    resnet50_conv.eval()
    # 2、密度图预测模块
    # input_channels=6，这是因为map3和map4分别产生三个尺度（0.9、1.0、1.1）的correlation map
    regressor = CountRegressor(input_channels=6, pool='mean').to(device)
    weights_normal_init(regressor, dev=0.001)
    regressor.train()


    # 3、这里冻结Resnet50FPN，只更新CountRegressor的参数
    optimizer = optim.Adam(regressor.parameters(), lr=args.learning_rate)

    best_mae, best_rmse = 1e7, 1e7
    stats = list()
    for epoch in range(0, args.epochs):
        regressor.train()
        # 训练1个epoch
        train_loss, train_mae, train_rmse = train_one_epoch(args, resnet50_conv, optimizer, regressor, criterion, device)
        regressor.eval()
        # 模型评估
        val_mae, val_rmse = eval(args, resnet50_conv, regressor, device)
        # 将训练过程中的loss等信息保存到文件中
        stats.append((train_loss, train_mae, train_rmse, val_mae, val_rmse))
        stats_file = join(args.output_dir, "stats" + ".txt")
        with open(stats_file, 'w') as f:
            for s in stats:
                f.write("%s\n" % ','.join([str(x) for x in s]))
        if best_mae < val_mae:
            best_mae = val_mae
            best_rmse = val_rmse
            model_name = args.output_dir + '/' + f"FamNet_{best_mae}.pth"
            torch.save(regressor.state_dict(), model_name)
        print("Epoch {}, Avg. Epoch Loss: {} Train MAE: {} Train RMSE: {} Val MAE: {} Val RMSE: {} Best Val MAE: {} Best Val RMSE: {} ".format(
                epoch + 1, stats[-1][0], stats[-1][1], stats[-1][2], stats[-1][3], stats[-1][4], best_mae, best_rmse))
  • train_one_epoch如下所示
  • 训练过程中，需要对图像大小进行预处理
    • 训练过程中同时需对密度图进行处理，预测或测试不需对密度图进行处理。
    • 如果高度和宽度均小于指定值，则不进行调整大小。
    • 如果图像的宽度或高度超过指定值，则调整图像大小，使得：
      • 新高度和新宽度的最大值不超过指定值
      • 新高度和新宽度可被8整除
      • 保持纵横比
  • 核心函数为extract_features，用来获取不同尺度的correlation map，然后将6个correlation map进行concat，输入到密度图预测模块regressor
python 复制代码 
# LearningToCountEverything/train.py
# 模型训练
def train_one_epoch(args, resnet50_conv, optimizer, regressor, criterion, device):
    print("Training on FSC147 train set data")
    # 加载数据集及相关信息
    annotations, data_split, im_dir, gt_dir = load_data(args)
    # 训练数据集中图片名称, 如['7.jpg',...]
    im_ids = data_split['train']
    random.shuffle(im_ids)
    train_mae = 0
    train_rmse = 0
    train_loss = 0
    cnt = 0

    pbar = tqdm(im_ids)
    for im_id in pbar:
        cnt += 1
        anno = annotations[im_id]
        bboxes = anno['box_examples_coordinates']
        # 1、获取每幅图片上少量的exemplar object的bbox信息
        rects = list()
        for bbox in bboxes:
            x1 = bbox[0][0]
            y1 = bbox[0][1]
            x2 = bbox[2][0]
            y2 = bbox[2][1]
            rects.append([y1, x1, y2, x2])
        # 2、加载图片
        image = Image.open('{}/{}'.format(im_dir, im_id))
        image.load()
        # 3、加载密度图
        density_path = gt_dir + '/' + im_id.split(".jpg")[0] + ".npy"
        density = np.load(density_path).astype('float32')
        # 4、装入到字典中
        sample = {'image': image,
                  'lines_boxes': rects,
                  'gt_density': density}
        # 训练过程中，对图像进行预处理
        sample = TransformTrain(sample)
        image, boxes, gt_density = sample['image'].to(device), sample['boxes'].to(device), sample['gt_density'].to(device)

        with torch.no_grad():
            # 获取不同尺度的correlation map，并concat在一起
            features = extract_features(resnet50_conv, image.unsqueeze(0), boxes.unsqueeze(0), MAPS, Scales)
        features.requires_grad = True
        optimizer.zero_grad()
        # 将6个correlation map进行concat，然后输入 密度图预测模块
        output = regressor(features)

        # if image size isn't divisible by 8, gt size is slightly different from output size
        if output.shape[2] != gt_density.shape[2] or output.shape[3] != gt_density.shape[3]:
            orig_count = gt_density.sum().detach().item()
            gt_density = F.interpolate(gt_density, size=(output.shape[2], output.shape[3]), mode='bilinear')
            new_count = gt_density.sum().detach().item()
            if new_count > 0:
                # 保证gt_density缩放后的count和orig_count相同
                gt_density = gt_density * (orig_count / new_count)
        # 计算mse loss
        loss = criterion(output, gt_density)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()
        pred_cnt = torch.sum(output).item()
        gt_cnt = torch.sum(gt_density).item()
        cnt_err = abs(pred_cnt - gt_cnt)
        # 计算训练过程中mae和rmse
        train_mae += cnt_err
        train_rmse += cnt_err ** 2
        pbar.set_description('真实cnt: {:6.1f}, 预测cnt: {:6.1f}, 错误个数: {:6.1f}. Current MAE: {:5.2f}, RMSE: {:5.2f}'.format(gt_cnt, pred_cnt,
                                                                        abs(pred_cnt - gt_cnt), train_mae/cnt, (train_rmse/cnt)**0.5))
    # 计算训练1个epoch的train_loss、train_mae、train_rmse
    train_loss = train_loss / len(im_ids)
    train_mae = (train_mae / len(im_ids))
    train_rmse = (train_rmse / len(im_ids))**0.5
    return train_loss, train_mae, train_rmse

核心函数为extract_features

  • 使用了一个correlation map表征exempler特征和输入的图像特征的相似性(利用卷积实现)
  • 使用双线性插值实现ROI Pooling操作，使得每一张图片中的exemplars的高和宽统一
  • 由于object可能会有尺度变化，因此作者把exemplers的大小进行缩放（0.9、1.0、1.1）
  • 生成多个correlation maps(map3和map4分别产生三个尺度的correlation map)，这些特征图最后被concat，最后送入密度图预测模块
python 复制代码 
# LearningToCountEverything/utils.py
def extract_features(feature_model, image, boxes, feat_map_keys=['map3','map4'], exemplar_scales=[0.9, 1.1]):
    """
      1、使用了一个correlation map表征exempler特征和输入的图像特征的相似性(利用卷积实现)
      2、由于object可能会有尺度变化，因此作者把exempler的大小进行缩放（0.9、1.1）
      3、生成多个correlation maps(map3和map4分别产生三个尺度的correlation map)，这些特征图最后被concat送入密度图预测模块
    """
    N, M = image.shape[0], boxes.shape[2] # 图片个数及exampler的个数

    # 使用ResNet-50的第三个和第四个block得到的特征图表征输入图像
    # Getting features for the image N * C * H * W
    Image_features = feature_model(image)

    # Getting features for the examples (N*M) * C * h * w
    for ix in range(0, N):
        # boxes = boxes.squeeze(0)
        boxes = boxes[ix][0]  # 一张图像中exampler的bbox信息
        cnter = 0
        for keys in feat_map_keys:
            image_features = Image_features[keys][ix].unsqueeze(0)
            if keys == 'map1' or keys == 'map2':
                Scaling = 4.0
            elif keys == 'map3':
                Scaling = 8.0
            elif keys == 'map4':
                Scaling = 16.0
            else:
                Scaling = 32.0
            # exempler的bbox信息缩放到特征图尺度
            boxes_scaled = boxes / Scaling
            boxes_scaled[:, 1:3] = torch.floor(boxes_scaled[:, 1:3])
            boxes_scaled[:, 3:5] = torch.ceil(boxes_scaled[:, 3:5])
            # make the end indices exclusive
            boxes_scaled[:, 3:5] = boxes_scaled[:, 3:5] + 1
            feat_h, feat_w = image_features.shape[-2], image_features.shape[-1]
            # make sure exemplars don't go out of bound
            boxes_scaled[:, 1:3] = torch.clamp_min(boxes_scaled[:, 1:3], 0)
            boxes_scaled[:, 3] = torch.clamp_max(boxes_scaled[:, 3], feat_h)
            boxes_scaled[:, 4] = torch.clamp_max(boxes_scaled[:, 4], feat_w)            
            box_hs = boxes_scaled[:, 3] - boxes_scaled[:, 1]
            box_ws = boxes_scaled[:, 4] - boxes_scaled[:, 2]            
            # 获取一张图片中exemplars在特征图上最大的高和宽
            max_h = math.ceil(max(box_hs))
            max_w = math.ceil(max(box_ws))            
            for j in range(0, M):
                y1, x1 = int(boxes_scaled[j, 1]), int(boxes_scaled[j, 2])
                y2, x2 = int(boxes_scaled[j, 3]), int(boxes_scaled[j, 4])
                # print(y1,y2,x1,x2,max_h,max_w)
                if j == 0:
                    # 获取exemplars的相应特征图
                    examples_features = image_features[:, :, y1:y2, x1:x2]
                    if examples_features.shape[2] != max_h or examples_features.shape[3] != max_w:
                        # 双线性插值填充(目的：每一张图片中的exemplars的高和宽统一)
                        examples_features = F.interpolate(examples_features, size=(max_h, max_w), mode='bilinear')
                else:
                    feat = image_features[:, :, y1:y2, x1:x2]
                    if feat.shape[2] != max_h or feat.shape[3] != max_w:
                        # 双线性插值填充(目的：每一张图片中的exemplars的高和宽统一)
                        feat = F.interpolate(feat, size=(max_h, max_w), mode='bilinear')
                    # concat
                    examples_features = torch.cat((examples_features, feat), dim=0)
            """
                Convolving example features over image features
                【使用卷积计算相似】
                这里把examples_features当作卷积核 在 image_features上进行卷积，得到correlation map
            """
            h, w = examples_features.shape[2], examples_features.shape[3]
            # input shape(minibatch, in_channels, iH, iW)
            # weight shape(out_channels, in_channels/groups, kH, kW)
            # out_channels就是一张图片中bbox的个数
            features = F.conv2d(
                    input=F.pad(input=image_features, pad=((int(w/2)), int((w-1)/2), int(h/2), int((h-1)/2))), # 进行填充，使卷积后correlation map高宽不变
                    weight=examples_features
                )
            combined = features.permute([1, 0, 2, 3])
            # computing features for scales 0.9 and 1.1
            # 考虑到不同 scale 的 object 我们会对其进行缩放
            # 缩放之后会得到新的 correlation map
            for scale in exemplar_scales:
                    h1 = math.ceil(h * scale)
                    w1 = math.ceil(w * scale)
                    if h1 < 1: # use original size if scaled size is too small
                        h1 = h
                    if w1 < 1:
                        w1 = w
                    examples_features_scaled = F.interpolate(examples_features, size=(h1,w1), mode='bilinear')
                    features_scaled = F.conv2d(
                        F.pad(image_features, ((int(w1/2)), int((w1-1)/2), int(h1/2), int((h1-1)/2))),
                        examples_features_scaled
                    )
                    features_scaled = features_scaled.permute([1, 0, 2, 3])
                    # 把所有的correlation map concatenate 在一起
                    combined = torch.cat((combined, features_scaled), dim=1)
            if cnter == 0:
                Combined = 1.0 * combined
            else:
                # 对map4进行上采样，和map3统一宽和高
                if Combined.shape[2] != combined.shape[2] or Combined.shape[3] != combined.shape[3]:
                    combined = F.interpolate(combined, size=(Combined.shape[2], Combined.shape[3]), mode='bilinear')
                Combined = torch.cat((Combined, combined), dim=1)
            cnter += 1
        if ix == 0:
            All_feat = 1.0 * Combined.unsqueeze(0)
        else:
            All_feat = torch.cat((All_feat, Combined.unsqueeze(0)), dim=0)
    return All_feat

2.3 模型预测

  • 预测需要图片标注少量的bbox，如果没有提供bbox文件，则提示用户标注一组边界框
  • 预测过程中的预处理和训练基本一致，不过没有密度图相关的处理。
  • 多尺度特征提取模块提取特征后，可以开启自适应训练(默认100轮)，然后再输出到密度图预测模块
  • 预测所用的loss为MincountLoss和PerturbationLoss的加权和。这两种损失函数已经在1上文中介绍过，不再介绍。
  • 在某些情况下，损失可能会变为零，其中损失是一个值为零的标量而不是张量。因此，仅针对非零情况执行梯度下降。
  • 这里使用官方提供的预训练模型进行检测，可视化的检测结果如下图。
python 复制代码 
# LearningToCountEverything/demo.py
def detect(args):
    if not torch.cuda.is_available() or args.gpu_id < 0:
        use_gpu = False
        print("===> Using CPU mode.")
    else:
        use_gpu = True
        os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
        os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_id)
    # 1、多尺度特征提取模块
    resnet50_conv = Resnet50FPN()
    # 2、密度图预测模块
    # input_channels=6，这是因为Resnet50FPN的map3和map4分别产生三个尺度（0.9、1.0、1.1）的correlation map
    regressor = CountRegressor(input_channels=6, pool='mean')

    # 3、加载训练模型
    if use_gpu:
        resnet50_conv.cuda()
        regressor.cuda()
        regressor.load_state_dict(torch.load(args.model_path))
    else:
        regressor.load_state_dict(torch.load(args.model_path, map_location=torch.device('cpu')))

    resnet50_conv.eval()
    regressor.eval()

    image_name = os.path.basename(args.input_image)
    image_name = os.path.splitext(image_name)[0]

    # 如果没有提供bbox文件，则提示用户输入一组边界框
    # if no bounding box file is given, prompt the user for a set of bounding boxes
    if args.bbox_file is None:
        out_bbox_file = "{}/{}_box.txt".format(args.output_dir, image_name)
        fout = open(out_bbox_file, "w")

        im = cv2.imread(args.input_image)
        cv2.imshow('image', im)
        rects = select_exemplar_rois(im)

        rects1 = list()
        for rect in rects:
            y1, x1, y2, x2 = rect
            rects1.append([y1, x1, y2, x2])
            print(rects1)
            fout.write("{} {} {} {}\n".format(y1, x1, y2, x2))

        fout.close()
        cv2.destroyAllWindows()
        print("selected bounding boxes are saved to {}".format(out_bbox_file))
    else:
        with open(args.bbox_file, "r") as fin:
            lines = fin.readlines()

        rects1 = list()
        for line in lines:
            data = line.split()
            y1 = int(data[0])
            x1 = int(data[1])
            y2 = int(data[2])
            x2 = int(data[3])
            rects1.append([y1, x1, y2, x2])

    print("Bounding boxes: ", end="")
    print(rects1)
    # 3、加载图像、并进行预处理
    image = Image.open(args.input_image)
    image.load()
    sample = {'image': image, 'lines_boxes': rects1}
    sample = Transform(sample)
    image, boxes = sample['image'], sample['boxes']

    if use_gpu:
        image = image.cuda()
        boxes = boxes.cuda()

    # 4、多尺度特征提取模块提取特征
    with torch.no_grad():
        features = extract_features(resnet50_conv, image.unsqueeze(0), boxes.unsqueeze(0), MAPS, Scales)

    # 5、输出预测结果
    if not args.adapt:
        # 不采用自适应，直接输出预测结果
        with torch.no_grad():
            output = regressor(features)
    else:
        # 采用自适应，先进行训练(默认100轮)，再输出
        features.required_grad = True
        adapted_regressor = regressor
        adapted_regressor.train()
        # 仍然只对regressor参数进行微调
        optimizer = optim.Adam(adapted_regressor.parameters(), lr=args.learning_rate)

        pbar = tqdm(range(args.gradient_steps))
        for step in pbar:
            optimizer.zero_grad()
            output = adapted_regressor(features)
            # 此时所用loss为MincountLoss和PerturbationLoss的加权和
            lCount = args.weight_mincount * MincountLoss(output, boxes, use_gpu=use_gpu)
            lPerturbation = args.weight_perturbation * PerturbationLoss(output, boxes, sigma=8, use_gpu=use_gpu)
            Loss = lCount + lPerturbation
            # 在某些情况下，损失可能会变为零，其中损失是一个值为零的标量而不是张量。
            # 因此，仅针对非零情况执行梯度下降。
            if torch.is_tensor(Loss):
                Loss.backward()
                optimizer.step()

            pbar.set_description('Adaptation step: {:<3}, loss: {}, predicted-count: {:6.1f}'.format(step, Loss.item(),
                                                                                                     output.sum().item()))

        features.required_grad = False
        output = adapted_regressor(features)

    print(output.shape)
    print('===> The predicted count is: {:6.2f}'.format(output.sum().item()))
    # 6、可视化预测结果
    rslt_file = "{}/{}_out.png".format(args.output_dir, image_name)
    visualize_output_and_save(image.detach().cpu(), output.detach().cpu(), boxes.cpu(), rslt_file)
    print("===> Visualized output is saved to {}".format(rslt_file))
相关推荐
懒大王爱吃狼40 分钟前
Python教程：python枚举类定义和使用
开发语言·前端·javascript·python·python基础·python编程·python书籍
深度学习实战训练营1 小时前
基于CNN-RNN的影像报告生成
人工智能·深度学习
秃头佛爷2 小时前
Python学习大纲总结及注意事项
开发语言·python·学习
昨日之日20063 小时前
Moonshine - 新型开源ASR（语音识别）模型，体积小，速度快，比OpenAI Whisper快五倍 本地一键整合包下载
人工智能·whisper·语音识别
浮生如梦_3 小时前
Halcon基于laws纹理特征的SVM分类
图像处理·人工智能·算法·支持向量机·计算机视觉·分类·视觉检测
深度学习lover3 小时前
＜项目代码＞YOLOv8 苹果腐烂识别＜目标检测＞
人工智能·python·yolo·目标检测·计算机视觉·苹果腐烂识别
热爱跑步的恒川4 小时前
【论文复现】基于图卷积网络的轻量化推荐模型
网络·人工智能·开源·aigc·ai编程
API快乐传递者4 小时前
淘宝反爬虫机制的主要手段有哪些？
爬虫·python
阡之尘埃6 小时前
Python数据分析案例61——信贷风控评分卡模型(A卡)（scorecardpy 全面解析）
人工智能·python·机器学习·数据分析·智能风控·信贷风控
孙同学要努力8 小时前
全连接神经网络案例——手写数字识别
人工智能·深度学习·神经网络
热门推荐
01【HarmonyOS】HUAWEI DevEco Studio 下载地址汇总02(欧拉)openEuler系统添加网卡文件配置流程、(欧拉)openEuler系统手动配置ipv6地址流程、(欧拉)openEuler系统网络管理说明03组基轨迹建模 GBTM的介绍与实现（Stata 或 R）04【AIGC】重塑未来的科技巨轮05全面解析：构建基于深度学习的安全帽检测系统（UI界面+YOLO代码+数据集）06【经验分享】Ubuntu22.04安装微信(linux官方版)07基于YOLOv10深度学习的CT扫描图像肾结石智能检测系统【python源码+Pyqt5界面+数据集+训练代码】深度学习实战、目标检测08Ubuntu 20.04使用Livox mid 360 测试 FAST_LIO09RAG 实践- Ollama+RagFlow 部署本地知识库10红米手机使用google play