【扒代码】test.py-CSDN博客

本文链接：https://blog.csdn.net/2301_77549977/article/details/140814280
扒代码、debug代码
#!/usr/bin/env python
# coding: utf-8

"""
Test code written by Viresh Ranjan

Last modified by: Minh Hoai Nguyen (minhhoai@cs.stonybrook.edu)
Date: 2021/04/19

① 对于函数：弄明白in&out&deal 需要的参数类型 需要什么参数
② 没有函数
"""

import copy
from model import CountRegressor, Resnet50FPN
from utils import MAPS, Scales, Transform, extract_features
from utils import MincountLoss, PerturbationLoss
from PIL import Image
import os
import torch
import argparse
import json
import numpy as np
from tqdm import tqdm
from os.path import exists
import torch.optim as optim

# 有自定义库

'''
from model import CountRegressor, Resnet50FPN
from utils import MAPS, Scales, Transform, extract_features
from utils import MincountLoss, PerturbationLoss
'''

# 创建命令行接口，接受命令行参数

# 创建 ArgumentParser 对象，用于处理命令行参数
parser = argparse.ArgumentParser(description="Few Shot Counting Evaluation code")
# 添加 --data_path 参数，用于指定FSC147数据集的路径
# parser.add_argument("-dp", "--data_path", type=str, default='/Users/dearr/Downloads/LearningToCountEverything-master/data', help="Path to the FSC147 dataset")
parser.add_argument("-dp", "--data_path", type=str, default='data/', help="Path to the FSC147 dataset")
# 添加 --test_split 参数，用于指定要评估的数据集分割
# 接受的选项包括 'val_PartA', 'val_PartB', 'test_PartA', 'test_PartB', 'test', 'val'
parser.add_argument("-ts", "--test_split", type=str, default='val', choices=["val_PartA","val_PartB","test_PartA","test_PartB","test", "val"], help="what data split to evaluate on")
# 添加 --model_path 参数，用于指定训练模型的路径
parser.add_argument("-m",  "--model_path", type=str, default="./data/pretrainedModels/FamNet_Save1.pth", help="path to trained model")
# 添加 --adapt 参数，如果指定，则执行测试时适应（test time adaptation）
parser.add_argument("-a",  "--adapt", action='store_true', help="If specified, perform test time adaptation")
# 添加 --gradient_steps 参数，用于指定适应过程中的梯度步数
parser.add_argument("-gs", "--gradient_steps", type=int,default=100, help="number of gradient steps for the adaptation")
# 添加 --learning_rate 参数，用于指定适应过程中的学习率
parser.add_argument("-lr", "--learning_rate", type=float,default=1e-7, help="learning rate for adaptation")
# 添加 --weight_mincount 参数，用于指定最小计数损失的权重乘数
parser.add_argument("-wm", "--weight_mincount", type=float,default=1e-9, help="weight multiplier for Mincount Loss")
# 添加 --weight_perturbation 参数，用于指定扰动损失的权重乘数
parser.add_argument("-wp", "--weight_perturbation", type=float,default=1e-4, help="weight multiplier for Perturbation Loss")
# 添加 --gpu-id 参数，用于指定使用的GPU编号，默认为0，使用-1表示使用CPU
parser.add_argument("-g",  "--gpu_id", type=int, default=0, help="GPU id. Default 0 for the first GPU. Use -1 for CPU.")
# 解析命令行参数，将解析后的参数赋值给 args 变量
args = parser.parse_args()

# NOTE 下划线乱用

# 读取数据文件
data_path = args.data_path
# 读取标注信息文件
anno_file = data_path + 'annotation_FSC147_384.json'
# 读取数据分割文件
data_split_file = data_path + 'Train_Test_Val_FSC_147.json'
# 读取图像文件夹
im_dir = data_path + 'images_384_VarV2'

# 检查anno_file和im_dir是否存在
if not exists(anno_file) or not exists(im_dir):
    # 如果anno_file或im_dir不存在，打印提示信息
    print("Make sure you set up the --data_path correctly.")
    # 通知用户当前的数据路径设置，并指出图像目录和注释文件不存在
    print("Current setting is {}, but the image dir and annotation file do not exist.".format(args.data_path))
    # 退出程序
    print("Aborting the evaluation")
    # 退出程序，并返回错误代码-1
    exit(-1)

# 判断GPU是否可用
if not torch.cuda.is_available() or args.gpu_id < 0:
    # 如果GPU不可用或指定使用CPU，设置use_gpu为False
    use_gpu = False
    print("===> Using CPU mode.")
else:
    # 如果GPU可用，设置use_gpu为True
    use_gpu = True
    # 意思就是把gpu的排列顺序变成物理顺序
    # 设置CUDA_DEVICE_ORDER环境变量为PCI_BUS_ID
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    # 设置CUDA_VISIBLE_DEVICES环境变量为args.gpu_id
    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_id)

# 问题： 训练模型评估模型 但是为什么又加载模型？
resnet50_conv = Resnet50FPN()
if use_gpu: resnet50_conv.cuda()
resnet50_conv.eval()

regressor = CountRegressor(6, pool='mean')
regressor.load_state_dict(torch.load(args.model_path))
if use_gpu: regressor.cuda()
regressor.eval()

# 打开标注文件
with open(anno_file) as f:
    annotations = json.load(f)

# 打开数据分割文件
with open(data_split_file) as f:
    data_split = json.load(f)

cnt = 0  #cnt 是一个全局计数器，用于记录已经处理了多少个样本
SAE = 0  # sum of absolute errors 绝对误差
SSE = 0  # sum of square errors 平方误差

# args.test_split = 'val'
print("Evaluation on {} data".format(args.test_split))

# len(data_split[args.test_split]) = 1286
# 从预定义的数据分割中获取当前正在评估的图像ID列表
im_ids = data_split[args.test_split]
# 使用tqdm库创建进度条，传入im_ids作为迭代对象
pbar = tqdm(im_ids)


# 遍历图像ID列表
for im_id in pbar:
    # 根据当前的图像ID从注释字典中获取对应的注释信息
    anno = annotations[im_id]
    # 提取注释中的边界框坐标列表
    bboxes = anno['box_examples_coordinates']
    # 将注释中的点坐标转换为NumPy数组格式
    dots = np.array(anno['points'])

    # 初始化一个空列表，用于存储转换后的边界框坐标
    rects = list()

    # 遍历包含边界框坐标的列表
    # NOTE 确定一个矩形需要两个点即可
    for bbox in bboxes:
        # 从当前边界框bbox中提取左上角和右下角的坐标
        # 假设bbox中的坐标是以[y1, x1, y2, x2]的形式给出的
        # bbox[0]可能是一个包含左上角坐标的元组或列表
        x1, y1 = bbox[0][0], bbox[0][1] # 左上角的x和y坐标
        # bbox[2]可能是一个包含右下角坐标的元组或列表
        x2, y2 = bbox[2][0], bbox[2][1] # 右下角的x和y坐标

        # 将提取的坐标添加到rects列表中，格式为[y1, x1, y2, x2]
        # 注意：这里假设bbox的坐标格式是[y1, x1, y2, x2]，但根据代码似乎是需要转换顺序
        rects.append([y1, x1, y2, x2])

    # 使用Python的PIL库（Pillow）打开图像文件，文件路径由im_dir和im_id指定 
    image = Image.open('{}/{}'.format(im_dir, im_id))
    # 调用load()方法，确保图像被加载到内存中，准备后续处理
    image.load()

    # 创建一个包含当前图像及其边界框信息的字典sample
    sample = {'image': image, 'lines_boxes': rects}
    # 将sample字典传递给Transform函数或类方法，以应用一些变换 
    # TODO Transform函数得好好看看
    sample = Transform(sample)
    # 从Transform函数返回的结果中提取变换后的图像和边界框
    image, boxes = sample['image'], sample['boxes']

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

    # TODO 这个函数extract_features是干嘛的？
    with torch.no_grad(): 
        features = extract_features(resnet50_conv, image.unsqueeze(0), boxes.unsqueeze(0), MAPS, Scales)

    # 检查命令行参数中是否指定了适应（adapt）模式
    if not args.adapt:
        # 如果没有指定适应模式，则使用torch.no_grad()上下文管理器
        with torch.no_grad():
            # 在torch.no_grad()上下文管理器的作用下，不计算梯度
            # 这通常用于推理阶段，以减少内存使用并加快计算速度
            output = regressor(features)
    else:
        '''
        
            首先启用了features张量的梯度计算，这是进行反向传播的前提
            使用copy.deepcopy创建了regressor模型的一个副本adapted_regressor
            这样做是为了避免在适应性训练过程中修改原始模型
            接着，代码进入一个循环，进行指定次数的梯度下降步骤，每次步骤中都会：
                -  清除已有的梯度。
                - 计算模型输出。
                - 计算两种损失（MincountLoss和PerturbationLoss），并将它们相加得到总损失。
                - 如果总损失是一个张量（而不是一个标量），则执行反向传播并更新模型参数
            在完成适应性训练后，代码再次禁用了features张量的梯度计算，
            这通常是为了减少内存消耗，并在模型评估或推理阶段使用
            使用经过适应性训练的模型adapted_regressor获取最终的输出

        '''
        # 将features张量的requires_grad属性设置为True，以便在反向传播时计算梯度
        features.required_grad = True
        # 使用copy.deepcopy创建regressor模型的一个深拷贝，以便在适应性训练中使用
        adapted_regressor = copy.deepcopy(regressor)
        # 将适应性模型设置为训练模式
        adapted_regressor.train()
        # 创建一个Adam优化器，用于后续的梯度下降步骤
        optimizer = optim.Adam(adapted_regressor.parameters(), lr=args.learning_rate)

        # 进行args.gradient_steps次梯度下降步骤
        for step in range(0, args.gradient_steps):
            # 在每次梯度下降之前清除已有的梯度
            optimizer.zero_grad()

            # 通过适应性模型获取输出
            output = adapted_regressor(features)
            # 计算MincountLoss损失，并乘以权重
            lCount = args.weight_mincount * MincountLoss(output, boxes)
            # 计算PerturbationLoss损失，并乘以权重
            lPerturbation = args.weight_perturbation * PerturbationLoss(output, boxes, sigma=8)
            # 将两种损失相加得到总损失
            Loss = lCount + lPerturbation
            # loss can become zero in some cases, where loss is a 0 valued scalar and not a tensor
            # 只有在损失不是零的情况下才执行反向传播和优化器步骤
            # So Perform gradient descent only for non zero cases
            # 这可以防止因为损失是标量而不是张量而导致的错误
            if torch.is_tensor(Loss):
                # 反向传播计算Loss关于模型参数的梯度
                Loss.backward()
                # 更新模型参数
                optimizer.step()

        # 将features张量的requires_grad属性重新设置为False，通常在模型评估或推理阶段这样做
        features.required_grad = False

        # 使用适应性训练后的模型再次获取输出
        output = adapted_regressor(features)


    '''
        模型评估阶段，记录模型性能
        gt_cnt 和 pred_cnt 分别表示每个样本的真实计数和模型预测的计数。
        cnt 是一个全局计数器，用于记录已经处理了多少个样本。
        err 是当前样本的实际计数与预测计数之间的绝对误差。
        SAE（Sum of Absolute Errors）是所有样本绝对误差的累积和。
        SSE（Sum of Squared Errors）是所有样本误差平方的累积和，用于计算RMSE。
        pbar.set_description 更新进度条的描述，显示当前样本的评估信息和当前的MAE与RMSE。
        最后的 print 语句在所有样本处理完毕后，输出整个数据集的MAE和RMSE。
    '''

    # 获取真实计数，即标注点的数量
    gt_cnt = dots.shape[0]

    # 计算预测计数，即模型输出的总和，并将其转换为Python的标量值
    pred_cnt = output.sum().item()

    # 增加计数器，用于记录处理的样本数量
    cnt = cnt + 1

    # 计算当前样本的实际计数与预测计数的绝对误差
    err = abs(gt_cnt - pred_cnt)

    # 将当前样本的误差加到累积的绝对误差SAE中
    SAE += err

    # 将当前样本的误差平方加到累积的平方误差SSE中
    SSE += err**2

    # 使用tqdm库更新进度条的描述信息
    # 描述信息包括图像ID、实际计数、预测计数、误差、当前MAE和RMSE
    pbar.set_description('{:<8}: actual-predicted: {:6d}, {:6.1f}, error: {:6.1f}. Current MAE: {:5.2f}, RMSE: {:5.2f}'.\
                         format(im_id, gt_cnt, pred_cnt, abs(pred_cnt - gt_cnt), SAE/cnt, (SSE/cnt)**0.5))
    
    # 打印空行
    print("")

# 在处理完所有样本后，打印最终的MAE和RMSE结果
# 这些结果是在指定的数据集（由args.test_split指定）上计算得出的
print('On {} data, MAE: {:6.2f}, RMSE: {:6.2f}'.format(args.test_split, SAE/cnt, (SSE/cnt)**0.5))