人群计数[CAN](Context-Aware Crowd Counting) 代码解读

之前已经完成了代码复现，并且自己也将ShanghaiTech中part_A_final和part_B_final数据集结果跑出来了。现在对代码部分做一个详细解读，加深自己的理解，如果有不对的地方请大家多多指教！！
（本文默认已经阅读过此篇论文并且对python的语法有一定的了解）

代码分析

可以看见里面有create_json.py、dataset.py、image.py、make_dataset.py、model.py、test.py、train.py和utils.py共8个py代码部分。我们按照代码的执行顺序来依次理解。

1、images.py

image.py主要实现了对图片的处理和变换。论文中，提出在训练过程中，在不同的位置随机剪裁原始图像的1/4的图像块。本文件中只定义了一个函数load_data（），主要作用就是生成h5py文件、读取图像和密度图以及剪裁图像。

import random
import os
from PIL import Image
import numpy as np
import h5py
import cv2


#主要实现了对图片的处理和变换。
def load_data(img_path, train=True):
    # 生成一个h5py文件的文件名：replace函数可以把字符串里面的old字符串替换成new字符串
    gt_path = img_path.replace('.jpg', '.h5').replace('images', 'ground_truth')
    #读出图像，后面加上.convert('RGB')是将图像转换为RGB
    #如果不使用.convert('RGB')进行转换的话，读出来的图像是RGBA四通道的，A通道为透明通道，该对深度学习模型训练来说暂时用不到。
    img = Image.open(img_path).convert('RGB')
    #生成h5py文件：为图片建立了h5py文件，并且在程序中打开为gt_file
    gt_file = h5py.File(gt_path, 'r')
    #将文件里图片对应的密度图读取，并转化为numpy格式
    target = np.asarray(gt_file['density'])
    #在不同的位置随机切割图像
    if train:
        ratio = 0.5
        crop_size = (int(img.size[0]*ratio), int(img.size[1]*ratio))
        #生产随机数
        rdn_value = random.random()
        if rdn_value < 0.25:
            dx = 0
            dy = 0
        elif rdn_value < 0.5:
            dx = int(img.size[0]*ratio)
            dy = 0
        elif rdn_value < 0.75:
            dx = 0
            dy = int(img.size[1]*ratio)
        else:
            dx = int(img.size[0]*ratio)
            dy = int(img.size[1]*ratio)

        img = img.crop((dx, dy, crop_size[0]+dx, crop_size[1]+dy))
        target = target[dy:(crop_size[1]+dy), dx:(crop_size[0]+dx)]
        if random.random() > 0.8:
            #作用是将数组在左右方向上翻转
            target = np.fliplr(target)
            #矩阵转置
            img = img.transpose(Image.FLIP_LEFT_RIGHT)

    #对target采取了宽高取1/8的操作因为CAN的输出结果就是原图大小的1/8，采用了INTER_CUBIC变换法，并在最后乘了64以保证图片像素之和不变。
    target = cv2.resize(target, (target.shape[1]//8, target.shape[0]//8), interpolation=cv2.INTER_CUBIC)*64

    return img, target

2、dataset.py

dataset文件主要实现了数据集的创建。在__getitem__（）函数中调用了image.py中的load_data（）函数，来得到处理后的图像和密度图。

import os
import random
import torch
import numpy as np
from torch.utils.data import Dataset
from PIL import Image
from image import *
import torchvision.transforms.functional as F

#dataset文件主要实现了数据集的创建
class listDataset(Dataset):
    #root是指用于训练的图片的路径，是列表类型；shuffle表示是否需要将root的顺序打乱；train表示这个数据集是否是用于训练的；batch_size表示了训练批数据的大小
    def __init__(self, root, shape=None, shuffle=True, transform=None,  train=False, seen=0, batch_size=1, num_workers=4):
        #打乱路径
        random.shuffle(root)
        #对变量成员的初始化
        self.nSamples = len(root)  #计算图片的数量
        self.lines = root
        self.transform = transform
        self.train = train
        self.shape = shape
        self.seen = seen
        self.batch_size = batch_size
        self.num_workers = num_workers
        
    #在构造函数中，定义了self.nSamples为root列表的长度，所以在只需要return self.nSample即可。
    def __len__(self):
        return self.nSamples

    #得到图片和密度图
    def __getitem__(self, index):
        # getitem()里，assert语句用来声明某个条件是真的，后面一般会跟一个逻辑式，式子值为True的时候不发生任何事情，式子值为False时会引发异常。这里显然是判断下标是否溢出。
        assert index <= len(self), 'index range error'
        #img_path从成员函数self.lines处读取了图片的地址
        img_path = self.lines[index]
        #load_data()函数。这个函数在image.py文件里被定义。这个函数读取了图片地址，返回用于训练的图片和图片对应对的密度图的numpy变量
        img, target = load_data(img_path, self.train)

        #下面判断图片是否需要做变换，最后return图片和图片对应的密度图。
        #由transform为None知道，此句子不用执行
        if self.transform is not None:
            img = self.transform(img)
        return img, target

3、make_dataset.py

首先运行make_dataset.py文件，主要是为了让image和mat文件经过高斯核的计算，产生hdf5文件（target文件），运行make_dataset.py脚本后，hdf5文件（后缀.h5）会产生在数据集ground_truth文件夹中。

import h5py
import scipy.io as io
import PIL.Image as Image
import numpy as np
import os
import glob
from matplotlib import pyplot as plt
from scipy.ndimage.filters import gaussian_filter
from matplotlib import cm as CM
from image import *

#通过make_dataset.py创建h5py文件（target文件）
#让image和mat文件经过高斯核的计算，产生target文件
#运行成功后，h5py文件（后缀.h5）会产生在数据集的ground_truth文件夹中

# ShanghaiTech数据集的路径
root = './data/'

#直接将地址赋值和使用os.path.join()函数等同
#part_B_final_train = './part_B_final/train_data/images'
#os.path.join()函数：连接两个或更多的路径名组件
part_B_final_train = os.path.join(root, 'part_B_final/train_data', 'images')
part_B_final_test = os.path.join(root, 'part_B_final/test_data', 'images')
path_sets = [part_B_final_train, part_B_final_test]


img_paths = []
#glob.glob()函数：返回所有匹配的文件路径列表
#append()函数：用于在列表末尾添加新的对象。
#就是将所有的路径全部依次放在img_paths中
for path in path_sets:
    for img_path in glob.glob(os.path.join(path, '*.jpg')):
        img_paths.append(img_path)

#遍历处理图像，生成.h5文件
for img_path in img_paths:
    print(img_path)
    #使用模块scipy.io的函数loadmat和savemat可以实现Python对mat数据的读写
    #replace函数可以把字符串里面的old字符串替换成new字符串，max参数指替换不超过max次;replace(old,new,max)
    mat = io.loadmat(img_path.replace('.jpg', '.mat').replace('images', 'ground_truth').replace('IMG_', 'GT_IMG_'))
    #matplotlib库的pyplot模块中的imread()函数用于将文件中的图像读取到数组中。
    img = plt.imread(img_path)
    #shape函数是numpy.core.fromnumeric中的函数，它的功能是读取矩阵的长度，比如shape[0]就是读取矩阵第一维度的长度即矩阵的行数；shape[0]是读取列数
    #创建一个与img同型的全0矩阵
    k = np.zeros((img.shape[0], img.shape[1]))
    #image_info获取图片或者是目录下所有图片
    gt = mat["image_info"][0, 0][0, 0][0]
    #range(start,end),不包含end

    # 让image和mat文件经过高斯核的计算，产生target文件，就是密度图
    for i in range(0, len(gt)):
        if int(gt[i][1]) < img.shape[0] and int(gt[i][0]) < img.shape[1]:
            k[int(gt[i][1]), int(gt[i][0])] = 1
    #高斯滤波gaussian_filter(),高斯滤波是一种线性平滑滤波，可以去除高斯噪声，其效果是降低图像灰度的尖锐变化，也就是图像模糊了。
    k = gaussian_filter(k, 15)
    #将生成的h5py文件（后缀.h5）产生在数据集的ground_truth文件夹中
    with h5py.File(img_path.replace('.jpg', '.h5').replace('images', 'ground_truth'), 'w') as hf:
            hf['density'] = k
            

4、create_json.py

然后执行create_json.py文件来产生json文件，train.json和val.json文件分别是训练和测试文件的路径。比较简单，不多说了。

import json
from os.path import join
import glob

#产生json文件，train.json和val.json文件分别是训练和测试文件的路径，里面包含所有图片的路径
if __name__ == '__main__':
    # 包含图像的文件夹路径
    #img_folder = './data/part_B_final/train_data/images'
    img_folder = './data/part_B_final/test_data/images'

    # 最终json文件的路径
    #output_json = './train.json'
    output_json = './val.json'

    img_list = []

    # glob.glob()函数：返回所有匹配的文件路径列表
    # append()函数：用于在列表末尾添加新的对象。
    for img_path in glob.glob(join(img_folder, '*.jpg')):
        img_list.append(img_path)

    #将图片路径写入json文件
    with open(output_json, 'w') as f:
        json.dump(img_list, f)

5、utils.py

里面定义了三个函数，前两个没有用到，所以就只讲解第三个函数了。（主要是我也没看明白前两个函数是干什么的，如果有大神知道，可以告诉我一下）

import h5py
import torch
import shutil
import numpy as np


def save_net(fname, net):
    with h5py.File(fname, 'w') as h5f:
        for k, v in net.state_dict().items():
            h5f.create_dataset(k, data=v.cpu().numpy())

def load_net(fname, net):
    with h5py.File(fname, 'r') as h5f:
        for k, v in net.state_dict().items():        
            param = torch.from_numpy(np.asarray(h5f[k]))         
            v.copy_(param)

#前两个函数并没有用到，重点是这个函数
#每次训练完之后会生成一个模型，为checkpoint.pth.tar文件，这个函数主要就是为了判断这个模型是否是最好效果的模型，如果是，将他复制到文件model_best.pth.tar
def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
    # torch.save()用法：保存模型参数
    torch.save(state, filename)
    if is_best:
        #shutil.copyfile（）复制文件
        shutil.copyfile(filename, 'model_best.pth.tar')            

6、model.py

model.py文件主要实现了网络模型的定义（是里面最复杂的文件了），我们先了解了解基础。
a、pytorch里面一切自定义操作基本上都是继承nn.Module类来实现的
b、一般把网络中具有可学习参数的层（如全连接层、卷积层等）放在构造函数__init__()中，当然也可以把不具有参数的层也放在里面；
c、一般把不具有可学习参数的层(如ReLU、dropout、BatchNormanation层)可放在构造函数中，也可不放在构造函数中，如果不放在构造函数__init__里面，则在forward方法里面可以使用nn.functional来代替
在这个文件中，主要实现了提取多尺度的上下文信息的模型和本文自己的网络模型。（多结合论文看看，很多计算，论文有提到，并且说明了原因）

import torch.nn as nn
import torch
from torch.nn import functional as F
from torchvision import models

#model.py主要实现了对网络模型的定义

#提取多尺度的上下文信息的模型定义
class ContextualModule(nn.Module):
    #我们使用S = 4个不同的尺度，对应的块大小k(j)∈{1,2,3,6}，因为它比其他设置表现出更好的性能。
    #features是指每个输入样本的大小；out_features是指每个输出样本的大小；sizes是指尺度大小
    def __init__(self, features, out_features=512, sizes=(1, 2, 3, 6)):
        super(ContextualModule, self).__init__()
        self.scales = []
        #利用特征金字塔来计算尺度特征？
        self.scales = nn.ModuleList([self._make_scale(features, size) for size in sizes])
        #bottleneck 的作用还是用更合理的方法减少了参数的数量，同时在尽可能不删除关键的特征的情况下
        #nn.Conv2d(in_channels, out_channels, kernel_size)作为二维卷积的实现；features*2是指输入张量的通道数;kernel_size是指卷积核的大小为1*1
        self.bottleneck = nn.Conv2d(features * 2, out_features, kernel_size=1)
        #nn.ReLU()为激活函数，也是卷积层
        self.relu = nn.ReLU()
        #每个这样的网络输出一个特定比例的表单权重图
        self.weight_net = nn.Conv2d(features, features, kernel_size=1)

    #计算预测权值w的函数
    def __make_weight(self, feature, scale_feature):
        #原文中给出的对比特征的计算公式：C_j = S_j - f_j，就是计算出对比特征
        weight_feature = feature - scale_feature
        #Sgimoid函数即形似S的函数，也成为S函数。在机器学习中经常用作分类
        #sigmoid函数来避免除0
        #使用self.weight_net()来进一步学习尺度感知特征的权值
        return F.sigmoid(self.weight_net(weight_feature))

    #计算尺度
    def _make_scale(self, features, size):
        #nn.AdaptiveAvgPool2d()自适应平均池化函数
        prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
        #二维卷积
        conv = nn.Conv2d(features, features, kernel_size=1, bias=False)
        #nn.Sequential()一个有序的容器，神经网络模块将按照传入构造器的顺序依次被添加到计算图中执行，
        return nn.Sequential(prior, conv)

    #反馈给后端网络
    def forward(self, feats):
        h, w = feats.size(2), feats.size(3)
        #自适应处理，学习计算尺度
        multi_scales = [F.upsample(input=stage(feats), size=(h, w), mode='bilinear') for stage in self.scales]
        #根据自适应处理得到得到尺度来计算权值
        weights = [self.__make_weight(feats, scale_feature) for scale_feature in multi_scales]
        #利用权值来计算上下文特征（原文的公式（5））
        overall_features = [(multi_scales[0]*weights[0]+multi_scales[1]*weights[1]+multi_scales[2]*weights[2]+multi_scales[3]*weights[3])/(weights[0]+weights[1]+weights[2]+weights[3])]+ [feats]
        #进行卷积处理
        bottle = self.bottleneck(torch.cat(overall_features, 1))
        return self.relu(bottle)

#本文自己网络的模型
class CANNet(nn.Module):
    def __init__(self, load_weights=False):
        super(CANNet, self).__init__()
        self.seen = 0
        #使用前面已经建立的提取多尺度的上下文信息的模型
        self.context = ContextualModule(512, 512)
        #网络前端，‘M’表示这一层是一个MaxPooling池化层，frontend_feat表示的是VGG-16网络的前10层
        self.frontend_feat = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512]
        #网络后端，back_feat表示的是采用空洞卷积的后端
        self.backend_feat = [512, 512, 512, 256, 128, 64]
        #前端和后端都使用了make_layers()函数自定义实现
        self.frontend = make_layers(self.frontend_feat)
        self.backend = make_layers(self.backend_feat, in_channels=512, batch_norm=True, dilation=True)
        #最后的output层采用了1*1卷积核实现了特征平面数量向1的转变。
        self.output_layer = nn.Conv2d(64, 1, kernel_size=1)
        #根据前面load_weights为False，所以这个if一定会执行
        if not load_weights:
            #保存VGG-16 前10层的网络结构
            mod = models.vgg16(pretrained=True)
            #先调用self._initialize_weights()方法进行手动初始化
            self._initialize_weights()
            #net.state_dict()是网络全部参数的字典，里面的键是网络各层参数的名字，值是封装好参数的Tensor
            #dict.items()语句会返回一个可遍历的（键，值）元组，通过遍历这个元组的方式，即通过逐个的i访问mod.state_dict().items()[i][1].data[:]，即可遍历完所有的参数，完成拷贝。
            for i in range(len(self.frontend.state_dict().items())):
                list(self.frontend.state_dict().items())[i][1].data[:] = list(mod.state_dict().items())[i][1].data[:]

    #对网络forward的定义和初始化的方法
    def forward(self, x):
        x = self.frontend(x)
        x = self.context(x)
        x = self.backend(x)
        x = self.output_layer(x)
        return x

    #net.modules()是网络各层的列表，我们使用m来遍历列表中的元素，判断m的层类型，然后分别使用init下的函数来完成初始化。
    def _initialize_weights(self):
        for m in self.modules():
            #正态分部：nn.init.normal_(tensor, mean=0, std=1)
            if isinstance(m, nn.Conv2d):
                nn.init.normal_(m.weight, std=0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

#网络结构自定义的函数，是写在类的定义之外的
def make_layers(cfg, in_channels = 3, batch_norm=False, dilation = False):
    #通过dilation与否判断空洞率的大小，若是False则是网络前端，空洞率为1，若是True则是网络后端，空洞率为2。
    if dilation:
        d_rate = 2
    else:
        d_rate = 1
    layers = []
    #接下来遍历传入的frontend_feat和backend_feat，以此来确定网络层的类型，若是’M’则是MaxPooling层；
    #若是数字，则是卷积层，一套卷积层包括了Conv2d层，BatchNorm层和ReLU层，遍历完毕后，即可完成对网络前端和后端的创建。
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=d_rate, dilation=d_rate)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                #由于batch_norm=False,直接执行else语句
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
    return nn.Sequential(*layers)

7、train.py

就是训练模型，使其输入一张图像，能输出对应的较为准确的密度图。

import sys
import os

import warnings

from model import CANNet

from utils import save_checkpoint

import torch
import torch.nn as nn
#torch.autograd提供了类和函数用来对任意标量函数进行求导。
from torch.autograd import Variable
from torchvision import datasets, transforms

import numpy as np
import argparse
import json
import cv2
import dataset
import time

#argparse模块可以让python命令行启动的时候接收参数
#通过argparse.ArgumentParser()来创建一个解析对象
parser = argparse.ArgumentParser(description='PyTorch CANNet')

#通过parser.add_argument()函数来增加命令行参数，metavar能改变显示出来的名字，help参数会在命令行打出-h或-help的时候出来
parser.add_argument('train_json', metavar='TRAIN',
                    help='path to train json')
parser.add_argument('val_json', metavar='VAL',
                    help='path to val json')

def main():

    global args, best_prec1

    best_prec1 = 1e6

    args = parser.parse_args()
    args.lr = 1e-4
    args.batch_size = 1

    args.decay = 5*1e-4
    args.start_epoch = 0
    args.epochs = 1000
    args.workers = 4
    args.seed = int(time.time())
    args.print_freq = 4
    #json.load()可以从文件中读取json字符
    #就是读取图片路径
    with open(args.train_json, 'r') as outfile:
        train_list = json.load(outfile)
    with open(args.val_json, 'r') as outfile:
        val_list = json.load(outfile)

    #为当前GPU设置随机种子，多GPU的时候应该使用
    torch.cuda.manual_seed(args.seed)

    #定义网络对象model
    model = CANNet()

    #将model转移到GPU上
    model = model.cuda()

    #定义误差函数为MSE(均方误差)
    criterion = nn.MSELoss(size_average=False).cuda()

    #定义优化器optimizer，采用随机梯度下降法，提供了学习率、动量以及衰退率
    optimizer = torch.optim.Adam(model.parameters(), args.lr,
                                    weight_decay=args.decay)

    #循环，一步一步训练网络
    for epoch in range(args.start_epoch, args.epochs):
        #网络向前传播和误差逆传播
        train(train_list, model, criterion, optimizer, epoch)
        #准确率检测功能
        prec1 = validate(val_list, model, criterion)

        #判断validate返回的MAE是否最优
        is_best = prec1 < best_prec1
        #保存最优的MEA，并且输出到屏幕上面
        best_prec1 = min(prec1, best_prec1)
        print(' * best MAE {mae:.3f} '
              .format(mae=best_prec1))
        #state_dict是一个列表，包含了目前训练到的epoch、网络和优化器的所有参数字典、最优的MAE
        save_checkpoint({
            'state_dict': model.state_dict(),
        }, is_best)

#主要是对train函数的实现，也是训练的核心部分，主要是完成了训练批数据的定义、网络向前传播和误差逆传播
def train(train_list, model, criterion, optimizer, epoch):

    losses = AverageMeter()
    batch_time = AverageMeter()
    data_time = AverageMeter()

    #使用torch.utils.data.DataLoader()方法进行了训练批数据的创建，第一个参数为数据集，数据集已经在dataset.py文件定义好了
    #train_list为训练图片的地址，打乱shuffle为True，对图片进行了transforms的Normalize
    train_loader = torch.utils.data.DataLoader(
        dataset.listDataset(train_list,
                       shuffle=True,
                       transform=transforms.Compose([
                       transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225]),
                   ]),
                       train=True,
                       seen=model.seen,
                       batch_size=args.batch_size,
                       num_workers=args.workers),
        batch_size=args.batch_size)
    print('epoch %d, processed %d samples, lr %.10f' % (epoch, epoch * len(train_loader.dataset), args.lr))

    #我们的网络包含了BN，所以在训练之前我们需要声明model.train()
    model.train()
    #time.time() 返回当前时间的时间戳（1970纪元后经过的浮点秒数）。
    end = time.time()

    #使用enumerate来连下标一起读取训练批数据，img是图片，target是真实的密度图（端到端）
    for i, (img, target) in enumerate(train_loader):
        data_time.update(time.time() - end)

        #前向和逆向传播的过程
        #将img转移到GPU上面
        img = img.cuda()
        #将img声明为Variable变量，要想使用自动求导，将所有的tensor包含进Variable对象中即可。
        img = Variable(img)
        #将img传入网络，得出预测的密度图
        output = model(img)[:, 0, :, :]

        #真实密度图
        target = target.type(torch.FloatTensor).cuda()
        target = Variable(target)

        #比较output与target，得出loss（损失函数）
        loss = criterion(output, target)

        #每次更新loss
        losses.update(loss.item(), img.size(0))
        #梯度清理
        optimizer.zero_grad()
        #反向传播，计算当前梯度
        loss.backward()
        #根据梯度更新网络参数
        optimizer.step()
        #计算此次训练的时间
        batch_time.update(time.time() - end)
        #记录结束时间
        end = time.time()

        if i % args.print_freq == 0:
            print('Epoch: [{0}][{1}/{2}]\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  .format(
                   epoch, i, len(train_loader), batch_time=batch_time,
                   data_time=data_time, loss=losses))

#完成批数据测试的建立
def validate(val_list, model, criterion):
    print('begin val')
    val_loader = torch.utils.data.DataLoader(
    dataset.listDataset(val_list,
                   shuffle=False,
                   transform=transforms.Compose([
                       transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225]),
                   ]),  train=False),
    batch_size=1)

    #声明model.eval()
    model.eval()
    # 初始化mae
    mae = 0

    #训练网络
    for i, (img, target) in enumerate(val_loader):
        #img.shape[]获取图片的高和宽
        h, w = img.shape[2:4]
        #将图像切割成四份
        h_d = h//2
        w_d = w//2
        ##将img声明为Variable变量，要想使用自动求导，将所有的tensor包含进Variable对象中即可。
        img_1 = Variable(img[:, :, :h_d, :w_d].cuda())
        img_2 = Variable(img[:, :, :h_d, w_d:].cuda())
        img_3 = Variable(img[:, :, h_d:, :w_d].cuda())
        img_4 = Variable(img[:, :, h_d:, w_d:].cuda())
        #求出每块的密度图
        density_1 = model(img_1).data.cpu().numpy()
        density_2 = model(img_2).data.cpu().numpy()
        density_3 = model(img_3).data.cpu().numpy()
        density_4 = model(img_4).data.cpu().numpy()

        #相加得出预测密度图
        pred_sum = density_1.sum()+density_2.sum()+density_3.sum()+density_4.sum()

        #计算损失函数
        mae += abs(pred_sum-target.sum())

    mae = mae/len(val_loader)
    print(' * MAE {mae:.3f} '
              .format(mae=mae))

    return mae

#用来封装统计量，并对他们进行计算
class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        #方差
        self.val = 0
        #平均数
        self.avg = 0
        #和
        self.sum = 0
        #数量
        self.count = 0

    #在update()中提供了计算这些量的方法
    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

if __name__ == '__main__':
    main()

8、test.py

终于到了最后一个文件，这个文件就是将训练好的模型，把测试图片生成密度图，然后计算出人群数量。

import h5py
import PIL.Image as Image
import numpy as np
import os
import glob
import scipy
from image import *
from model import CANNet
import torch
from torch.autograd import Variable

from sklearn.metrics import mean_squared_error,mean_absolute_error

from torchvision import transforms

#torchvision.transforms.Compose()类，这个类的主要作用是串联多个图片变换的操作。
transform = transforms.Compose([
                       transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225]),
                   ])

# the folder contains all the test images
img_folder = './data/part_B_final/test_data/images'
img_paths = []

#glob.glob()函数：返回所有匹配的文件路径列表
#append()函数：用于在列表末尾添加新的对象。
for img_path in glob.glob(os.path.join(img_folder, '*.jpg')):
    img_paths.append(img_path)

model = CANNet()

model = model.cuda()

#将最好的效果的模型文件导入
checkpoint = torch.load('model_best.pth.tar')

#model.state_dict()是浅拷贝，拷贝最外层的数值和指针，不拷贝更深层次的对象，即只拷贝了父对象
#model.load_state_dict() 是深拷贝，拷贝数值、指针和指针指向的深层次内存空间，拷贝了父对象及其子对象
model.load_state_dict(checkpoint['state_dict'])

#声明model.eval()
model.eval()

pred= []
gt = []

#遍历测试图像
for i in range(len(img_paths)):
    #转化为GPU可用
    img = transform(Image.open(img_paths[i]).convert('RGB')).cuda()
    #增加一个维度
    #经常用于CNN，因为conv2d的输入必须是四维的（batch,channel,height,width）,如果输入的是文本的话通常只是三维的（batch,length,dim）
    # 因此需要unsqueeze(1)，增加一维channel，才能做卷积操作（这里不明白为什么代码中是0）
    img = img.unsqueeze(0)
    # img.shape[]获取图片的高和宽
    h, w = img.shape[2:4]
    h_d = h//2
    w_d = w//2
    img_1 = Variable(img[:, :, :h_d, :w_d].cuda())
    img_2 = Variable(img[:, :, :h_d, w_d:].cuda())
    img_3 = Variable(img[:, :, h_d:, :w_d].cuda())
    img_4 = Variable(img[:, :, h_d:, w_d:].cuda())
    density_1 = model(img_1).data.cpu().numpy()
    density_2 = model(img_2).data.cpu().numpy()
    density_3 = model(img_3).data.cpu().numpy()
    density_4 = model(img_4).data.cpu().numpy()

    #os.path.splitext（）分割路径，返回路径名和文件扩展名的元组
    pure_name = os.path.splitext(os.path.basename(img_paths[i]))[0]
    # 生成h5py文件：为图片建立了h5py文件，并且在程序中打开为gt_file
    gt_file = h5py.File(img_paths[i].replace('.jpg', '.h5').replace('images', 'ground_truth'), 'r')
    #将文件里图片对应的密度图读取，并转化为numpy格式
    groundtruth = np.asarray(gt_file['density'])
    #预测人数（密度图就是一个点代表一个人，在矩阵中为1，sum（）就是将矩阵的所有元素相加求和，所有计算出来的是人群数量）
    pred_sum = density_1.sum()+density_2.sum()+density_3.sum()+density_4.sum()
    pred.append(pred_sum)
    #真实人数直接求和
    gt.append(np.sum(groundtruth))

#计算平均绝对误差
mae = mean_absolute_error(pred, gt)
#计算根均方误差
rmse = np.sqrt(mean_squared_error(pred, gt))

print('pred:', pred)
print('gt:', gt)
print('MAE: ', mae)
print('RMSE: ', rmse)

后记

感觉model.py基本代码都看懂了，还是不太明白是怎么回事，主要还是不明白卷积操作。