FSC_test.py

import argparse
import json
import numpy as np
import os
from pathlib import Path
from PIL import Image, ImageDraw
import matplotlib.pyplot as plt
import scipy.ndimage as ndimage
import pandas as pd
import random
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
from torch.utils.data import Dataset
import torchvision
from torchvision import transforms
import torchvision.transforms.functional as TF
import timm
from util.FSC147 import transform_train, transform_val
from tqdm import tqdm
assert "0.4.5" <= timm.__version__ <= "0.4.9"  # version check

import util.misc as misc
import models_mae_cross


def get_args_parser():
    parser = argparse.ArgumentParser('MAE pre-training', add_help=False)

    # Model parameters
    parser.add_argument('--model', default='mae_vit_base_patch16', type=str, metavar='MODEL',
                        help='Name of model to train')
    parser.add_argument('--mask_ratio', default=0.5, type=float,
                        help='Masking ratio (percentage of removed patches).')
    parser.add_argument('--norm_pix_loss', action='store_true',
                        help='Use (per-patch) normalized pixels as targets for computing loss')
    parser.set_defaults(norm_pix_loss=False)

    # Dataset parameters
    parser.add_argument('--data_path', default='./data/FSC147/', type=str,
                        help='dataset path')
    parser.add_argument('--anno_file', default='annotation_FSC147_positive.json', type=str,
                        help='annotation json file')
    parser.add_argument('--anno_file_negative', default='./data/FSC147/annotation_FSC147_neg2.json', type=str,
                        help='annotation json file')
    parser.add_argument('--data_split_file', default='Train_Test_Val_FSC_147.json', type=str,
                        help='data split json file')
    parser.add_argument('--im_dir', default='images_384_VarV2', type=str,
                        help='images directory')
    parser.add_argument('--output_dir', default='./Image',
                        help='path where to save, empty for no saving')
    parser.add_argument('--device', default='cuda',
                        help='device to use for training / testing')
    parser.add_argument('--seed', default=0, type=int)
    parser.add_argument('--resume', default='./output_fim6_dir/checkpoint-0.pth',
                        help='resume from checkpoint')
    parser.add_argument('--external', action='store_true',
                        help='Set this param for using external exemplars')
    parser.add_argument('--box_bound', default=-1, type=int,
                        help='The max number of exemplars to be considered')
    parser.add_argument('--split', default="test", type=str)

    # Training parameters
    parser.add_argument('--num_workers', default=0, type=int)
    parser.add_argument('--pin_mem', action='store_true',
                        help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
    parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem')
    parser.set_defaults(pin_mem=True)
    parser.add_argument('--normalization', default=True, help='Set to False to disable test-time normalization')

    # Distributed training parameters
    parser.add_argument('--world_size', default=1, type=int,
                        help='number of distributed processes')
    parser.add_argument('--local_rank', default=-1, type=int)
    parser.add_argument('--dist_on_itp', action='store_true')
    parser.add_argument('--dist_url', default='env://',
                        help='url used to set up distributed training')

    return parser

os.environ["CUDA_LAUNCH_BLOCKING"] = '5'

class TestData(Dataset):
    def __init__(self, args, split='val', do_aug=True):
        with open(data_path/args.anno_file) as f:
            annotations = json.load(f)
                # Load negative annotations
        with open(args.anno_file_negative) as f:
            neg_annotations = json.load(f)
        with open(data_path/args.data_split_file) as f:
            data_split = json.load(f)

        self.img = data_split[split]
        random.shuffle(self.img)
        self.split = split
        self.img_dir = im_dir
        # self.TransformTrain = transform_train(args, do_aug=do_aug)
        self.TransformVal = transform_val(args)
        self.annotations = annotations
        self.neg_annotations = neg_annotations
        self.im_dir = im_dir

    def __len__(self):
        return len(self.img)

    def __getitem__(self, idx):
        im_id = self.img[idx]
        anno = self.annotations[im_id]
        bboxes = anno['box_examples_coordinates']
        dots = np.array(anno['points'])

        # 加载负样本的框
        neg_anno = self.neg_annotations[im_id]  # 假设每个图像ID在负样本注释中都有对应的条目
        neg_bboxes = neg_anno['box_examples_coordinates']

        rects = list()
        for bbox in bboxes:
            x1 = bbox[0][0]
            y1 = bbox[0][1]
            x2 = bbox[2][0]
            y2 = bbox[2][1]
            if x1 < 0:
                x1 = 0
            if x2 < 0:
                x2 = 0
            if y1 < 0:
                y1 = 0
            if y2 < 0:
                y2 = 0
            
            rects.append([y1, x1, y2, x2])
        neg_rects = list()
        for neg_bbox in neg_bboxes:
            x1 = neg_bbox[0][0]
            y1 = neg_bbox[0][1]
            x2 = neg_bbox[2][0]
            y2 = neg_bbox[2][1]
            if x1 < 0:
                x1 = 0
            if x2 < 0:
                x2 = 0
            if y1 < 0:
                y1 = 0
            if y2 < 0:
                y2 = 0
            
            neg_rects.append([y1, x1, y2, x2])

        image = Image.open('{}/{}'.format(self.im_dir, im_id))
        if image.mode == "RGBA":
            image = image.convert("RGB")
        image.load()
        m_flag = 0

        sample = {'image': image, 'lines_boxes': rects,'neg_lines_boxes': neg_rects, 'dots': dots, 'id': im_id, 'm_flag': m_flag}
        sample = self.TransformTrain(sample) if self.split == "train" else self.TransformVal(sample)
        # if self.split == "train":
        #     sample = self.TransformTrain(sample)
        # # print(sample.keys())
        return sample['image'], sample['gt_density'], len(dots), sample['boxes'], sample['neg_boxes'], sample['pos'],sample['m_flag'], im_id

def batched_rmse(predictions, targets, batch_size=100):
    """
    分批计算RMSE
    :param predictions: 模型预测的值，一个PyTorch张量
    :param targets: 真实的值，一个PyTorch张量，与predictions形状相同
    :param batch_size: 每个批次的大小
    :return: RMSE值
    """
    total_mse = 0.0
    total_count = 0

    # 分批处理
    for i in range(0, len(predictions), batch_size):
        batch_predictions = predictions[i:i+batch_size]
        batch_targets = targets[i:i+batch_size]
        
        # 确保使用float64进行计算以提高精度
        batch_predictions = batch_predictions.double()
        batch_targets = batch_targets.double()

        # 计算批次的MSE
        difference = batch_predictions - batch_targets
        mse = torch.mean(difference ** 2)

        # 累加MSE和计数
        total_mse += mse * len(batch_predictions)
        total_count += len(batch_predictions)

    # 计算平均MSE
    avg_mse = total_mse / total_count

    # 计算RMSE
    rmse_val = torch.sqrt(avg_mse)

    return rmse_val
def batched_mae(predictions, targets, batch_size=100):
    """
    分批计算MAE
    :param predictions: 模型预测的值，一个PyTorch张量
    :param targets: 真实的值，一个PyTorch张量，与predictions形状相同
    :param batch_size: 每个批次的大小
    :return: MAE值
    """
    total_mae = 0.0
    total_count = 0

    # 分批处理
    for i in range(0, len(predictions), batch_size):
        batch_predictions = predictions[i:i+batch_size]
        batch_targets = targets[i:i+batch_size]
        
        # 计算批次的绝对误差
        absolute_errors = torch.abs(batch_predictions - batch_targets)
        
        # 累加绝对误差和计数
        total_mae += torch.sum(absolute_errors)
        total_count += len(batch_predictions)

    # 计算平均绝对误差
    avg_mae = total_mae / total_count

    return avg_mae

def main(args):
    misc.init_distributed_mode(args)

    print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
    print("{}".format(args).replace(', ', ',\n'))

    device = torch.device(args.device)

    # fix the seed for reproducibility
    seed = args.seed + misc.get_rank()
    torch.manual_seed(seed)
    np.random.seed(seed)

    cudnn.benchmark = True

    # dataset_test = TestData(external=args.external, box_bound=args.box_bound, split=args.split)
    dataset_test = TestData(args, split='test')
    num_tasks = misc.get_world_size()
    global_rank = misc.get_rank()
    sampler_test = torch.utils.data.DistributedSampler(
        dataset_test, num_replicas=num_tasks, rank=global_rank, shuffle=True
    )

    data_loader_test = torch.utils.data.DataLoader(
        dataset_test, sampler=sampler_test,
        batch_size=1,
        num_workers=args.num_workers,
        pin_memory=args.pin_mem,
        drop_last=False,
    )

    # define the model
    model = models_mae_cross.__dict__[args.model](norm_pix_loss=args.norm_pix_loss)
    model.to(device)
    model_without_ddp = model

    if args.distributed:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
        model_without_ddp = model.module

    misc.load_model_FSC(args=args, model_without_ddp=model_without_ddp)

    print(f"Start testing.")

    # test
    model.eval()

    # some parameters in training
    train_mae = 0
    train_rmse = 0
    train_nae = 0
    tot_load_time = 0
    tot_infer_time = 0

    loss_array = []
    gt_array = []
    pred_arr = []
    name_arr = []
    empties = []

    total_mae = 0.0
    total_mse = 0.0
    total_nae = 0.0
    total_count = 0
    sub_batch_size = 50
    for val_samples, val_gt_density, val_n_ppl, val_boxes,neg_val_boxes, val_pos, _, val_im_names in tqdm(data_loader_test, total=len(data_loader_test), desc="Validation"):
        val_samples = val_samples.to(device, non_blocking=True, dtype=torch.float)  # 使用更高精度
        val_gt_density = val_gt_density.to(device, non_blocking=True, dtype=torch.float)
        val_boxes = val_boxes.to(device, non_blocking=True, dtype=torch.float)
        neg_val_boxes = neg_val_boxes.to(device, non_blocking=True, dtype=torch.float)
        num_samples = val_samples.size(0)
        total_count += num_samples

        for i in range(0, num_samples, sub_batch_size):
            sub_val_samples = val_samples[i:i+sub_batch_size]
            sub_val_gt_density = val_gt_density[i:i+sub_batch_size]

            with torch.no_grad():
                with torch.cuda.amp.autocast():
                    sub_val_output = model(sub_val_samples, val_boxes[i:i+sub_batch_size], 3)
            with torch.no_grad():
                with torch.cuda.amp.autocast():
                    neg_sub_val_output = model(sub_val_samples, neg_val_boxes[i:i+sub_batch_size], 3)
                # output = torch.clamp((sub_val_output-neg_sub_val_output),min=0)
                sub_val_pred_cnt = torch.abs(sub_val_output.sum()) / 60
                # sub_val_pred_cnt = torch.abs(output.sum()) / 60
                # neg_sub_val_pred_cnt = torch.abs(neg_sub_val_output.sum()) / 60
                sub_val_gt_cnt = sub_val_gt_density.sum() / 60
                
                sub_val_cnt_err = torch.abs(sub_val_pred_cnt - sub_val_gt_cnt)

                # 逐项添加并检查
                if not torch.isinf(sub_val_cnt_err) and not torch.isnan(sub_val_cnt_err):
                    batch_mae = sub_val_cnt_err.item()
                    batch_mse = sub_val_cnt_err.item() ** 2
                    batch_nae = sub_val_cnt_err.item() / sub_val_gt_cnt.item() if sub_val_gt_cnt.item() != 0 else 0

                    total_mae += batch_mae * sub_val_samples.size(0)
                    total_mse += batch_mse * sub_val_samples.size(0)
                    total_nae += batch_nae * sub_val_samples.size(0)
                sub_val_pred_cnt = (sub_val_pred_cnt).int()
    final_mae = total_mae / total_count
    final_rmse = (total_mse / total_count) ** 0.5
    final_nae = total_nae / total_count

    print(f'MAE: {final_mae}, RMSE: {final_rmse}, NAE: {final_nae}')



if __name__ == '__main__':
    args = get_args_parser()
    args = args.parse_args()

    # load data
    data_path = Path(args.data_path)
    anno_file = data_path / args.anno_file
    data_split_file = data_path / args.data_split_file
    im_dir = data_path / args.im_dir

    with open(anno_file) as f:
        annotations = json.load(f)

    with open(data_split_file) as f:
        data_split = json.load(f)

    if args.output_dir:
        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
    main(args)