test.py

# Copyright (c) 2025 Robert Bosch GmbH
# SPDX-License-Identifier: AGPL-3.0

import os
import cv2
import json
import torch
import random
import logging
import argparse
import numpy as np
from PIL import Image
from skimage import measure
from tabulate import tabulate
import torch.nn.functional as F
import torchvision.transforms as transforms
from sklearn.metrics import auc, roc_auc_score, average_precision_score, f1_score, precision_recall_curve, pairwise
from torchvision.transforms import InterpolationMode
from concurrent.futures import ProcessPoolExecutor, as_completed

import open_clip
from domain_adaption import memory as memory_da
from few_shot import memory as memory_fs
from model import LinearLayer
from dataset import VisaDataset, VisaDatasetTest, MVTecDataset, MPDDDataset, MADDataset, RealIADDataset_v2
from prompts.prompt_ensemble_visa_19cls_test import encode_text_with_prompt_ensemble as encode_text_with_prompt_ensemble_visa
from prompts.prompt_ensemble_mvtec_20cls import encode_text_with_prompt_ensemble as encode_text_with_prompt_ensemble_mvtec
from prompts.new_prompt_ensemble_mpdd import encode_text_with_prompt_ensemble as encode_text_with_prompt_ensemble_mpdd
from prompts.prompt_ensemble_mad_sim import encode_text_with_prompt_ensemble as encode_text_with_prompt_ensemble_mad_sim
from prompts.prompt_ensemble_mad_real import encode_text_with_prompt_ensemble as encode_text_with_prompt_ensemble_mad_real
from prompts.prompt_ensemble_real_IAD_simple import encode_text_with_prompt_ensemble as encode_text_with_prompt_ensemble_real_iad
import re
from tqdm import tqdm

import pdb


def setup_seed(seed):
	torch.manual_seed(seed)
	torch.cuda.manual_seed_all(seed)
	np.random.seed(seed)
	random.seed(seed)
	torch.backends.cudnn.deterministic = True
	torch.backends.cudnn.benchmark = False


def normalize(pred, max_value=None, min_value=None):
	if max_value is None or min_value is None:
		den = pred.max() - pred.min()
		if den < 1e-12:
			return np.zeros_like(pred)
		return (pred - pred.min()) / den
	else:
		den = max_value - min_value
		if den < 1e-12:
			return np.zeros_like(pred)
		return (pred - min_value) / den


def apply_ad_scoremap(image, scoremap, alpha=0.5):
	np_image = np.asarray(image, dtype=float)
	scoremap = (scoremap * 255).astype(np.uint8)
	scoremap = cv2.applyColorMap(scoremap, cv2.COLORMAP_JET)
	scoremap = cv2.cvtColor(scoremap, cv2.COLOR_BGR2RGB)
	return (alpha * np_image + (1 - alpha) * scoremap).astype(np.uint8)


def cal_pro_score(masks, amaps, max_step=200, expect_fpr=0.3):
	# ref: https://github.com/gudovskiy/cflow-ad/blob/master/train.py
	binary_amaps = np.zeros_like(amaps, dtype=bool)
	min_th, max_th = amaps.min(), amaps.max()
	if abs(max_th - min_th) < 1e-12:
		return 0.0

	delta = (max_th - min_th) / max_step
	pros, fprs, ths = [], [], []
	for th in np.arange(min_th, max_th, delta):
		binary_amaps[amaps <= th], binary_amaps[amaps > th] = 0, 1
		pro = []
		for binary_amap, mask in zip(binary_amaps, masks):
			for region in measure.regionprops(measure.label(mask)):
				tp_pixels = binary_amap[region.coords[:, 0], region.coords[:, 1]].sum()
				pro.append(tp_pixels / region.area)
		inverse_masks = 1 - masks
		fp_pixels = np.logical_and(inverse_masks, binary_amaps).sum()
		den = inverse_masks.sum()
		fpr = fp_pixels / den if den > 0 else 0.0
		pros.append(np.array(pro).mean() if len(pro) > 0 else 0.0)
		fprs.append(fpr)
		ths.append(th)

	pros, fprs, ths = np.array(pros), np.array(fprs), np.array(ths)
	idxes = fprs < expect_fpr
	if idxes.sum() < 2:
		return 0.0

	fprs = fprs[idxes]
	pros = pros[idxes]

	den = fprs.max() - fprs.min()
	if den < 1e-12:
		return 0.0

	fprs = (fprs - fprs.min()) / den
	pro_auc = auc(fprs, pros)
	return pro_auc


def evaluate_one_object(obj, results):
	table = []
	gt_px = []
	pr_px = []
	gt_sp = []
	pr_sp = []
	pr_sp_tmp = []
	table.append(obj)

	for idxes in range(len(results['cls_names'])):
		if results['cls_names'][idxes] == obj:
			gt_px.append(results['imgs_masks'][idxes].squeeze(1).numpy())
			pr_px.append(results['anomaly_maps'][idxes])
			pr_sp_tmp.append(np.mean(np.partition(results['anomaly_maps'][idxes].reshape(-1), -3)[-3:]))
			gt_sp.append(results['gt_sp'][idxes])
			pr_sp.append(results['pr_sp'][idxes])

	gt_px = np.array(gt_px)
	gt_sp = np.array(gt_sp)
	pr_px = np.array(pr_px)
	pr_sp = np.array(pr_sp)

	pr_sp_tmp = np.array(pr_sp_tmp)
	if len(pr_sp_tmp) > 0:
		den = pr_sp_tmp.max() - pr_sp_tmp.min()
		pr_sp_tmp = (pr_sp_tmp - pr_sp_tmp.min()) / den if den > 1e-12 else np.zeros_like(pr_sp_tmp)
	else:
		pr_sp_tmp = np.zeros_like(pr_sp)

	# keep your current logic
	pr_sp = pr_sp_tmp

	auroc_px = roc_auc_score(gt_px.ravel(), pr_px.ravel())
	auroc_sp = roc_auc_score(gt_sp, pr_sp)
	ap_sp = average_precision_score(gt_sp, pr_sp)
	ap_px = average_precision_score(gt_px.ravel(), pr_px.ravel())

	precisions, recalls, thresholds = precision_recall_curve(gt_sp, pr_sp)
	f1_scores = (2 * precisions * recalls) / (precisions + recalls + 1e-12)
	f1_sp = np.max(f1_scores[np.isfinite(f1_scores)])

	precisions, recalls, thresholds = precision_recall_curve(gt_px.ravel(), pr_px.ravel())
	f1_scores = (2 * precisions * recalls) / (precisions + recalls + 1e-12)
	f1_px = np.max(f1_scores[np.isfinite(f1_scores)])

	if len(gt_px.shape) == 4:
		gt_px = gt_px.squeeze(1)
	if len(pr_px.shape) == 4:
		pr_px = pr_px.squeeze(1)
	aupro = cal_pro_score(gt_px, pr_px)

	table.append(str(np.round(auroc_px * 100, decimals=1)))
	table.append(str(np.round(f1_px * 100, decimals=1)))
	table.append(str(np.round(ap_px * 100, decimals=1)))
	table.append(str(np.round(aupro * 100, decimals=1)))
	table.append(str(np.round(auroc_sp * 100, decimals=1)))
	table.append(str(np.round(f1_sp * 100, decimals=1)))
	table.append(str(np.round(ap_sp * 100, decimals=1)))

	return {
		"obj": obj,
		"table": table,
		"auroc_sp": auroc_sp,
		"auroc_px": auroc_px,
		"f1_sp": f1_sp,
		"f1_px": f1_px,
		"aupro": aupro,
		"ap_sp": ap_sp,
		"ap_px": ap_px,
	}


def test(args):
	img_size = args.image_size
	features_list = args.features_list
	few_shot_features = args.few_shot_features
	dataset_dir = args.data_path
	save_path = args.save_path
	dataset_name = args.dataset
	if not os.path.exists(save_path):
		os.makedirs(save_path)
	device = "cuda" if torch.cuda.is_available() else "cpu"
	txt_path = os.path.join(save_path, 'log.txt')

	# clip
	model, _, preprocess = open_clip.create_model_and_transforms(args.model, img_size, pretrained=args.pretrained)
	model.to(device)
	model.eval()
	tokenizer = open_clip.get_tokenizer(args.model)

	# logger
	root_logger = logging.getLogger()
	for handler in root_logger.handlers[:]:
		root_logger.removeHandler(handler)
	root_logger.setLevel(logging.WARNING)
	logger = logging.getLogger('test')
	logger.handlers.clear()
	formatter = logging.Formatter('%(asctime)s.%(msecs)03d - %(levelname)s: %(message)s',
								  datefmt='%y-%m-%d %H:%M:%S')
	logger.setLevel(logging.INFO)
	logger.propagate = False
	file_handler = logging.FileHandler(txt_path, mode='a')
	file_handler.setFormatter(formatter)
	logger.addHandler(file_handler)
	console_handler = logging.StreamHandler()
	console_handler.setFormatter(formatter)
	logger.addHandler(console_handler)

	# record parameters
	for arg in vars(args):
		if args.mode == 'zero_shot' and (arg == 'k_shot' or arg == 'few_shot_features'):
			continue
		logger.info(f'{arg}: {getattr(args, arg)}')

	# seg
	with open(args.config_path, 'r') as f:
		model_configs = json.load(f)
	linearlayer = LinearLayer(model_configs['vision_cfg']['width'], model_configs['embed_dim'],
							  len(features_list), args.model).to(device)
	checkpoint = torch.load(args.checkpoint_path, map_location=device)
	linearlayer.load_state_dict(checkpoint["trainable_linearlayer"])
	linearlayer.eval()

	# dataset
	target_transform_b = transforms.Compose([
		transforms.Resize((img_size, img_size)),
		transforms.CenterCrop(img_size),
		transforms.ToTensor()
	])
	target_transform_type = transforms.Compose([
		transforms.Resize((img_size, img_size), interpolation=InterpolationMode.NEAREST),
		transforms.CenterCrop(img_size),
		transforms.PILToTensor(),
		transforms.Lambda(lambda x: x.squeeze(0).long()),
	])

	if args.dataset == 'mvtec':
		test_data = MVTecDataset(root=dataset_dir, transform=preprocess, target_transform=target_transform_b, target_transform_type=target_transform_type, aug_rate=-1, mode='test')
	elif args.dataset == 'visa':
		test_data = VisaDatasetTest(
			root=args.data_path,
			transform=preprocess,
			target_transform=target_transform_b,
			mode="test"
		)
	elif args.dataset == 'mpdd':
		test_data = MPDDDataset(root=dataset_dir, transform=preprocess, target_transform=target_transform_b, target_transform_type=target_transform_type, mode='test')
	elif args.dataset == 'mad_sim':
		test_data = MADDataset(root=dataset_dir, transform=preprocess, target_transform=target_transform_b, target_transform_type=target_transform_type, mode='test', datatype='sim')
	elif args.dataset == 'mad_real':
		test_data = MADDataset(root=dataset_dir, transform=preprocess, target_transform=target_transform_b, target_transform_type=target_transform_type, mode='test', datatype='real')
	elif args.dataset == 'real_iad':
		test_data = RealIADDataset_v2(root=dataset_dir, transform=preprocess, target_transform=target_transform_b, target_transform_type=target_transform_type, mode='test')
	else:
		raise ValueError(f"Unsupported dataset: {args.dataset}")

	test_dataloader = torch.utils.data.DataLoader(test_data, batch_size=1, shuffle=False)
	obj_list = test_data.get_cls_names()

	# few shot
	if args.mode == 'few_shot':
		mem_features = memory_fs(args.model, model, obj_list, dataset_dir, save_path, preprocess, target_transform_b, target_transform_type,
							  args.k_shot, few_shot_features, dataset_name, device)

	if args.mode == 'domain_adaption':
		mem_features = memory_da(args.model, model, obj_list, dataset_dir, save_path, preprocess, target_transform_b, target_transform_type,
							 few_shot_features, dataset_name, device)

	# text prompt
	with torch.amp.autocast('cuda', enabled=torch.cuda.is_available()), torch.no_grad():
		if args.dataset == 'mvtec':
			text_prompts = encode_text_with_prompt_ensemble_mvtec(model, obj_list, tokenizer, device)
		elif args.dataset == 'visa':
			text_prompts = encode_text_with_prompt_ensemble_visa(model, obj_list, tokenizer, device)
		elif args.dataset == 'mpdd':
			text_prompts = encode_text_with_prompt_ensemble_mpdd(model, obj_list, tokenizer, device)
		elif args.dataset == 'mad_sim':
			text_prompts = encode_text_with_prompt_ensemble_mad_sim(model, obj_list, tokenizer, device)
		elif args.dataset == 'mad_real':
			text_prompts = encode_text_with_prompt_ensemble_mad_real(model, obj_list, tokenizer, device)
		elif args.dataset == 'real_iad':
			text_prompts = encode_text_with_prompt_ensemble_real_iad(model, obj_list, tokenizer, device)
		else:
			raise ValueError(f"Unsupported dataset: {args.dataset}")

	results = {}
	results['cls_names'] = []
	results['imgs_masks'] = []
	results['anomaly_maps'] = []
	results['gt_sp'] = []
	results['pr_sp'] = []

	for items in tqdm(test_dataloader, desc="Testing", leave=True):
		image = items['img'].to(device)
		cls_name = items['cls_name']
		results['cls_names'].append(cls_name[0])

		gt_mask = items['img_mask_b'].clone()

		for i in range(gt_mask.size(0)):
			gt_mask[i][gt_mask[i] > 0.5], gt_mask[i][gt_mask[i] <= 0.5] = 1, 0

		results['imgs_masks'].append(gt_mask)
		results['gt_sp'].append(items['anomaly'].item())

		with torch.no_grad(), torch.amp.autocast('cuda', enabled=torch.cuda.is_available()):
			image_features, patch_tokens = model.encode_image(image, features_list)
			image_features /= image_features.norm(dim=-1, keepdim=True)
			text_features = []

			for cls in cls_name:
				text_features.append(text_prompts[cls])
			text_features = torch.stack(text_features, dim=0)

			# sample
			text_probs = ((image_features @ text_features[0]) * 100.).softmax(dim=-1)
			results['pr_sp'].append(sum(text_probs[0][1:]).cpu().item())

			# pixel
			patch_tokens = linearlayer(patch_tokens)
			anomaly_maps = []
			for layer in range(len(patch_tokens)):
				patch_tokens[layer] /= patch_tokens[layer].norm(dim=-1, keepdim=True)
				anomaly_map = ((patch_tokens[layer] @ text_features) * 100.)
				B, L, C = anomaly_map.shape
				H = int(np.sqrt(L))
				anomaly_map = F.interpolate(
					anomaly_map.permute(0, 2, 1).view(B, C, H, H),
					size=img_size, mode='bilinear', align_corners=True
				)
				anomaly_map = torch.sum(torch.softmax(anomaly_map, dim=1)[:, 1:, :, :], dim=1)
				anomaly_maps.append(anomaly_map.cpu().numpy())
			anomaly_map = np.sum(anomaly_maps, axis=0)

			# few shot
			if args.mode == 'few_shot':
				image_features, patch_tokens = model.encode_image(image, few_shot_features)
				anomaly_maps_few_shot = []
				for idx, p in enumerate(patch_tokens):
					if 'ViT' in args.model:
						p = p[0, 1:, :]
					else:
						p = p[0].view(p.shape[1], -1).permute(1, 0).contiguous()
					cos = pairwise.cosine_similarity(mem_features[cls_name[0]][idx].cpu(), p.cpu())
					height = int(np.sqrt(cos.shape[1]))
					anomaly_map_few_shot = np.min((1 - cos), 0).reshape(1, 1, height, height)
					anomaly_map_few_shot = F.interpolate(torch.tensor(anomaly_map_few_shot),
														 size=img_size, mode='bilinear', align_corners=True)
					anomaly_maps_few_shot.append(anomaly_map_few_shot[0].cpu().numpy())
				anomaly_map_few_shot = np.sum(anomaly_maps_few_shot, axis=0)
				anomaly_map = anomaly_map + anomaly_map_few_shot

			# domain_adaption
			if args.mode == 'domain_adaption':
				image_features, patch_tokens = model.encode_image(image, few_shot_features)
				anomaly_maps_few_shot = []
				for idx, p in enumerate(patch_tokens):
					if 'ViT' in args.model:
						p = p[0, 1:, :]
					else:
						p = p[0].view(p.shape[1], -1).permute(1, 0).contiguous()
					cos = pairwise.cosine_similarity(mem_features[cls_name[0]][idx].cpu(), p.cpu())
					M, _ = cos.shape
					height = int(np.sqrt(cos.shape[1]))
					distances = 1.0 - cos

					k = max(1, int(np.ceil(M * args.quantile)))
					smallest_kplus1 = np.partition(distances, k, axis=0)[:k+1, :]
					sorted_smallest_kplus1 = np.sort(smallest_kplus1, axis=0)
					smallest_k = sorted_smallest_kplus1[1:k+1, :]
					avg_smallest = smallest_k.mean(axis=0)
					anomaly_map_few_shot = avg_smallest.reshape(1, 1, height, height)
					anomaly_map_few_shot = F.interpolate(torch.tensor(anomaly_map_few_shot),
														 size=img_size, mode='bilinear', align_corners=True)
					anomaly_maps_few_shot.append(anomaly_map_few_shot[0].cpu().numpy())
				anomaly_map_few_shot = np.sum(anomaly_maps_few_shot, axis=0)
				anomaly_map = anomaly_map + anomaly_map_few_shot

			results['anomaly_maps'].append(anomaly_map)

			# visualization
			if args.visualization:
				path = items['img_path']
				cls = path[0].split('/')[-2]
				filename = path[0].split('/')[-1]
				vis = cv2.cvtColor(cv2.resize(cv2.imread(path[0]), (img_size, img_size)), cv2.COLOR_BGR2RGB)
				mask = normalize(anomaly_map[0])
				vis = apply_ad_scoremap(vis, mask)
				vis = cv2.cvtColor(vis, cv2.COLOR_RGB2BGR)
				save_vis = os.path.join(save_path, 'imgs', cls_name[0], cls)
				if not os.path.exists(save_vis):
					os.makedirs(save_vis)
				cv2.imwrite(os.path.join(save_vis, filename), vis)

	# metrics (parallel)
	table_ls = []
	auroc_sp_ls = []
	auroc_px_ls = []
	f1_sp_ls = []
	f1_px_ls = []
	aupro_ls = []
	ap_sp_ls = []
	ap_px_ls = []

	num_workers = min(len(obj_list), os.cpu_count() or 1)

	with ProcessPoolExecutor(max_workers=num_workers) as executor:
		futures = {
			executor.submit(evaluate_one_object, obj, results): obj
			for obj in obj_list
		}

		for future in tqdm(as_completed(futures), total=len(futures), desc="Evaluating objects", leave=True):
			ret = future.result()
			table_ls.append(ret["table"])
			auroc_sp_ls.append(ret["auroc_sp"])
			auroc_px_ls.append(ret["auroc_px"])
			f1_sp_ls.append(ret["f1_sp"])
			f1_px_ls.append(ret["f1_px"])
			aupro_ls.append(ret["aupro"])
			ap_sp_ls.append(ret["ap_sp"])
			ap_px_ls.append(ret["ap_px"])

	# keep order same as obj_list
	table_ls = sorted(table_ls, key=lambda x: obj_list.index(x[0]))

	# logger
	table_ls.append([
		'mean',
		str(np.round(np.mean(auroc_px_ls) * 100, decimals=1)),
		str(np.round(np.mean(f1_px_ls) * 100, decimals=1)),
		str(np.round(np.mean(ap_px_ls) * 100, decimals=1)),
		str(np.round(np.mean(aupro_ls) * 100, decimals=1)),
		str(np.round(np.mean(auroc_sp_ls) * 100, decimals=1)),
		str(np.round(np.mean(f1_sp_ls) * 100, decimals=1)),
		str(np.round(np.mean(ap_sp_ls) * 100, decimals=1))
	])
	results_table = tabulate(
		table_ls,
		headers=['objects', 'auroc_px', 'f1_px', 'ap_px', 'aupro', 'auroc_sp', 'f1_sp', 'ap_sp'],
		tablefmt="pipe"
	)
	logger.info("\n%s", results_table)


if __name__ == '__main__':
	parser = argparse.ArgumentParser("MultiADS", add_help=True)
	# paths
	parser.add_argument("--data_path", type=str, default="./data/visa", help="path to test dataset")
	parser.add_argument("--save_path", type=str, default='./results/visa/zero_shot/', help='path to save results')
	parser.add_argument("--checkpoint_path", type=str, default='./exps/mvtec/epoch_1.pth', help='path to save results')
	parser.add_argument("--config_path", type=str, default='./open_clip/model_configs/ViT-L-14-336.json', help="model configs")
	# model
	parser.add_argument("--dataset", type=str, default='visa', help="test dataset")
	parser.add_argument("--model", type=str, default="ViT-L-14-336", help="model used")
	parser.add_argument("--pretrained", type=str, default="openai", help="pretrained weight used")
	parser.add_argument("--features_list", type=int, nargs="+", default=[6, 12, 18, 24], help="features used")
	parser.add_argument("--few_shot_features", type=int, nargs="+", default=[6, 12, 18, 24], help="features used for few shot")
	parser.add_argument("--image_size", type=int, default=518, help="image size")
	parser.add_argument("--mode", type=str, default="zero_shot", help="zero shot or few shot or domain adaption")
	# few shot
	parser.add_argument("--k_shot", type=int, default=10, help="e.g., 10-shot, 5-shot, 1-shot")
	# domain adaption
	parser.add_argument("--quantile", type=float, default=0.0001, help="percent of the qunatile of nearest neighbour")

	parser.add_argument("--seed", type=int, default=42, help="random seed")
	parser.add_argument('--visualization', action='store_true')
	args = parser.parse_args()

	setup_seed(args.seed)
	test(args)