codebase/inference/inference_tool.py

import logging
import pdb
import tqdm
import numpy as np
import open_clip
import torch
import torch.nn.functional as F
import os
from classname_and_prompt import *
from torchrs.datasets import AID, RESISC45, EuroSATRGB
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import pandas as pd
from clip_benchmark.datasets.builder import get_dataset_collate_fn
from clip_benchmark.metrics.zeroshot_retrieval import recall_at_k, batchify, dataloader_with_indices
from functools import reduce
import cv2
from scipy.ndimage import maximum_filter
from skimage import measure
import json
from datetime import datetime
from torchvision import transforms


def _convert_to_rgb(image):
    return image.convert('RGB')


def get_preprocess(image_resolution=224, is_train=False, subset_name="clip", aug=None):

    if subset_name == "clip":
        normalize = transforms.Normalize(
            mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711]
        )
    elif subset_name == "imagenet":
        normalize = transforms.Normalize(
            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
        )

    elif subset_name == "rs5m":
        normalize = transforms.Normalize(
            mean=[0.406, 0.423, 0.390], std=[0.188, 0.175, 0.185]
        )

    elif subset_name == "pub11":
        normalize = transforms.Normalize(
            mean=[0.445, 0.469, 0.441], std=[0.208, 0.193, 0.213]
        )

    elif subset_name == "rs3":
        normalize = transforms.Normalize(
            mean=[0.350, 0.356, 0.316], std=[0.158, 0.147, 0.143]
        )

    elif subset_name == "geometa":
        normalize = transforms.Normalize(
            mean=[0.320, 0.322, 0.285], std=[0.179, 0.168, 0.166]
        )

    if is_train:
        preprocess_train = transforms.Compose([
            transforms.RandomResizedCrop(
                image_resolution,
                interpolation=transforms.InterpolationMode.BICUBIC,
                scale=(0.9, 1.0)
            ),
            _convert_to_rgb,
            transforms.RandomHorizontalFlip(),
            transforms.RandomRotation(degrees=(0, 360)),
            transforms.ToTensor(),
            normalize,
        ])
        return preprocess_train
    else:
        preprocess_val = transforms.Compose([
            transforms.Resize(
                size=image_resolution,
                interpolation=transforms.InterpolationMode.BICUBIC,
            ),
            transforms.CenterCrop(image_resolution),
            _convert_to_rgb,
            transforms.ToTensor(),
            normalize,
        ])
        return preprocess_val


def zeroshot_get_dataset(dataset_name, root, split, transform=None):

    if dataset_name == "EuroSAT":
        EuroSAT_root = os.path.join(root, "eurosat-rgb")
        os.makedirs(EuroSAT_root, exist_ok=True)
        dataset = EuroSATRGB(
            root=EuroSAT_root,
            transform=transform
        )
        dataset.classes = dataset.classes
        dataset.templates = RSEuroSAT.templates

    elif dataset_name == "AID":
        AID_root = os.path.join(root, "AID")
        os.makedirs(AID_root, exist_ok=True)
        dataset = AID(
            root=AID_root,
            transform=transform
        )
        dataset.classes = dataset.classes
        dataset.templates = RSAID.templates

    elif dataset_name == "RESISC45":
        RESISC45_root = os.path.join(root, "RESISC45")
        os.makedirs(RESISC45_root, exist_ok=True)
        dataset = RESISC45(
            root=RESISC45_root,
            transform=transform
        )
        dataset.classes = dataset.classes
        dataset.templates = RSRESISC45.templates

    dataset.classes = [dataset.classes[i].replace('_', ' ') for i in range(len(dataset.classes))]
    dataset.classes = [dataset.classes[i].replace('/', ' ') for i in range(len(dataset.classes))]
    dataset.classes = [dataset.classes[i].lower() for i in range(len(dataset.classes))]

    return dataset


def zeroshot_classifier(model, classnames, templates, args):
    tokenizer = open_clip.tokenize
    with torch.no_grad():
        zeroshot_weights = []
        for classname in classnames:
            texts = [template.replace('{}', classname) for template in templates]
            context_length = 77
            texts = tokenizer(texts, context_length=context_length).to(args.device)

            class_embeddings = model.encode_text(texts)
            class_embeddings = class_embeddings.mean(dim=0)
            class_embedding = F.normalize(class_embeddings, dim=-1)
            class_embedding /= class_embedding.norm()
            zeroshot_weights.append(class_embedding.cpu())
        zeroshot_weights = torch.stack(zeroshot_weights, dim=1)
    return zeroshot_weights


def zeroshot_evaluation(model, zeroshot_dataset, preprocess, args):

    dataset = zeroshot_get_dataset(dataset_name=zeroshot_dataset, split='test', root=args.test_dataset_dir, transform=preprocess)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, num_workers=args.workers)

    logging.info(f'Calculating classifier for {zeroshot_dataset}')
    classnames, prompt_templates = dataset.classes, dataset.templates
    import copy
    classnames = copy.deepcopy(classnames)
    classifier = zeroshot_classifier(model, classnames, prompt_templates, args)

    logging.info(f'Calculating image features for {zeroshot_dataset}')
    results = {}
    acc, features, labels = zeroshot_run(model, classifier, dataloader, args)
    logging.info(f'{zeroshot_dataset} zero-shot accuracy: {acc}%')
    results[f'{zeroshot_dataset}-zeroshot-acc'] = acc

    for key, item in results.items():
        results[key] = float(item)

    return results


def zeroshot_accuracy(output, target, topk=(1,)):
    pred = output.topk(max(topk), 1, True, True)[1].t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    return float(correct[0].reshape(-1).float().sum(0, keepdim=True).cpu().numpy()) * 100 / len(target)


def zeroshot_run(model, classifier, dataloader, args):
    with torch.no_grad():
        all_image_features = []
        all_labels = []
        all_logits = []
        for images, target in tqdm.tqdm(dataloader, unit_scale=args.batch_size):
            images = images.to(args.device)
            image_features = model.encode_image(images)
            image_features = F.normalize(image_features, dim=-1).detach().cpu()
            logits = 100. * image_features @ classifier
            all_image_features.append(image_features)
            all_labels.append(target)
            all_logits.append(logits)

    all_image_features = torch.cat(all_image_features)
    all_labels = torch.cat(all_labels)
    all_logits = torch.cat(all_logits)

    acc = zeroshot_accuracy(all_logits, all_labels, topk=(1,))
    return round(acc, 2), all_image_features, all_labels


class CsvDataset(Dataset):
    def __init__(self, input_filename, transforms, img_key, caption_key, sep="\t", nori_dataset=False,
                 images_dir=''):
        logging.debug(f'Loading csv data from {input_filename}.')
        if 'rsicd' in input_filename:
            df = pd.read_csv(input_filename, sep=sep, encoding='gb18030')
        else:
            df = pd.read_csv(input_filename, sep=sep)

        self.nori_dataset = nori_dataset
        self.f = None
        self.images_dir = images_dir

        self.images = df[img_key].tolist()
        self.captions = df[caption_key].tolist()

        self.transforms = transforms

        self.duplicate()

        logging.debug('Done loading data.')

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

    def __getitem__(self, index):
        texts = self.captions[index]
        image = Image.open(os.path.join(self.images_dir, str(self.images[index])))
        image = self.transforms(image)

        return image, texts

    def duplicate(self):
        unique_images, indexs = np.unique(self.images, return_index=True)
        if len(unique_images) != len(self.images):
            logging.debug(
                f'Amoung all {len(self.images)} images, there are only {len(unique_images)} unique images. Dupication will be performed to enable one-image-to-multiple-text retrieval.')
            self.duplicated_images = []
            self.duplicated_captions = []
            for index in indexs:
                self.duplicated_images.append(self.images[index])
                same_indexs = [i for i, x in enumerate(self.images) if x == self.images[index]]
                captions = []
                for same_index in same_indexs:
                    captions.append(self.captions[same_index])
                self.duplicated_captions.append(captions)

            self.images = self.duplicated_images
            self.captions = self.duplicated_captions


def retrieval_evaluation(model, preprocess, args, recall_k_list=[1, 5, 10], dataset_name=None):
    """
    Modified from https://github.com/LAION-AI/CLIP_benchmark/blob/main/clip_benchmark/metrics/zeroshot_retrieval.py
    Evaluate the model on the given dataset

    Parameters
    ----------

    model: torch.nn,Module
        CLIP-like model with `encode_image` and `encode_text`

    dataloader: torch.utils.data.Dataloader
        dataloader to use for evaluation

    tokenizer:
        text tokenizer, i.e. convert list of strings to torch.Tensor of integers

    device: cpu/cuda
    recall_k_list: list of int
        recall@k k's to use

    Returns
    -------

    dict of retrieval metrics
    """

    if dataset_name == "rsitmd":
        dataset = CsvDataset(
            input_filename=os.path.join(args.test_dataset_dir, "rsitmd", "rsitmd_test.csv"),
            transforms=preprocess,
            img_key="filename",
            caption_key="title",
            sep=",",
            images_dir=os.path.join(args.test_dataset_dir, "rsitmd", "images")
        )
    elif dataset_name == "rsicd":
        dataset = CsvDataset(
            input_filename=os.path.join(args.test_dataset_dir, "rsicd", "rsicd_test.csv"),
            transforms=preprocess,
            img_key="filename",
            caption_key="title",
            sep=",",
            images_dir=os.path.join(args.test_dataset_dir, "rsicd", "RSICD_images")
        )

    dataloader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        num_workers=args.workers,
        collate_fn=get_dataset_collate_fn('mscoco_captions')
    )
    n_batches = len(dataloader)
    tokenizer = open_clip.tokenize
    # list of batch of images embedding
    batch_images_emb_list = []
    # list of batch of text embedding
    batch_texts_emb_list = []
    # for each text, we collect the corresponding image index, as each image can have multiple corresponding texts
    texts_image_index = []
    dataloader = dataloader_with_indices(dataloader)

    for batch_images, batch_texts, inds in tqdm.tqdm(dataloader, total=n_batches):
        batch_images = batch_images.to(args.device)
        # store the index of image for each text
        batch_texts_image_index = [ind for ind, texts in zip(inds, batch_texts) for text in texts]
        # tokenize all texts in the batch
        batch_texts = tokenizer([text for i, texts in enumerate(batch_texts) for text in texts]).to(args.device)

        # compute the embedding of images and texts
        with torch.no_grad():
            batch_image_features = model.encode_image(batch_images)
            batch_text_features = model.encode_text(batch_texts)
            batch_images_emb = F.normalize(batch_image_features, dim=-1)
            batch_texts_emb = F.normalize(batch_text_features, dim=-1)

        batch_images_emb_list.append(batch_images_emb.cpu())
        batch_texts_emb_list.append(batch_texts_emb.cpu())
        texts_image_index.extend(batch_texts_image_index)

    batch_size = len(batch_images_emb_list[0])

    # concatenate all embeddings
    images_emb = torch.cat(batch_images_emb_list)
    texts_emb = torch.cat(batch_texts_emb_list)

    # get the score for each text and image pair
    scores = texts_emb @ images_emb.t()

    # construct a the positive pair matrix, which tells whether each text-image pair is a positive or not
    positive_pairs = torch.zeros_like(scores, dtype=bool)
    positive_pairs[torch.arange(len(scores)), texts_image_index] = True
    metrics = {}
    for recall_k in recall_k_list:
        '''
        Note that recall_at_k computes **actual** recall i.e. nb_true_positive/nb_positives, where the number
        of true positives, e.g. for text retrieval, is, for each image,  the number of retrieved texts matching that image among the top-k.
        Also, the number of positives are the total number of texts matching the image in the dataset, as we have a set of captions
        for each image, that number will be greater than 1 for text retrieval.
        However, image/text retrieval recall@k, the way it is done in CLIP-like papers, is a bit different.
        recall@k, in CLIP-like papers, is, for each image, either 1 or 0. It is 1 if atleast one text matches the image among the top-k.
        so we can easily compute that using the actual recall, by checking whether there is at least one true positive,
        which would be the case if the recall is greater than 0. One we compute the recal for each image (or text), we average
        it over the dataset.
        '''
        metrics[f"retrieval-image2text-R@{recall_k}-{dataset_name}"] = (batchify(recall_at_k, scores.T,
                                                                                 positive_pairs.T, batch_size,
                                                                                 args.device,
                                                                                 k=recall_k) > 0).float().mean().item() * 100

    for recall_k in recall_k_list:
        metrics[f"retrieval-text2image-R@{recall_k}-{dataset_name}"] = (batchify(recall_at_k, scores, positive_pairs,
                                                                                 batch_size, args.device,
                                                                                 k=recall_k) > 0).float().mean().item() * 100

    metrics[f"retrieval-mean-recall-{dataset_name}"] = np.mean(list(metrics.values()))

    for key, item in metrics.items():
        metrics[key] = round(float(item), 2)
    logging.info(f'{dataset_name} retrieval recall: {metrics}%')

    return metrics


class SLM(object):

    # **
    # * Copyright @2022 AI, AIRCAS. (mails.ucas.ac.cn)
    #
    # @author yuanzhiqiang <yuanzhiqiang19@mails.ucas.ac.cn>
    #         2022/03/08

    def __init__(self):
        # logging
        logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        self.logger = logging.getLogger()

        # parameters
        self.rsu_beta = 0.707
        self.rsu_eps = 1e-7

        self.ras_expand_factor = 1.5
        self.ras_filter_times = 5
        self.ras_scala_beta = 3

        self.rda_eta = 0.5

        self.rmi_wsu = 0.4
        self.rmi_was = 0.35
        self.rmi_wda = 0.25

        # visual settings
        self.visual_ras = False
        self.src_addmap_path = None

        # sum indicator
        self.all_metrics = self._format_output_dict()

    def _format_output_dict(self, *params):
        """
        format output dict
        :param params: keys
        :return: format dict
        """
        len_params = len(params)
        if len_params == 0: init_param = [[] for i in range(4)]
        elif len_params == 4: init_param = params
        else: raise NotImplementedError

        return {
            "↑ Rsu [0 ~ 1]": init_param[0],
            "↑ Rda [0 ~ 1]": init_param[1],
            "↓ Ras [0 ~ 1]": init_param[2],
            "↑ Rmi [0 ~ 1]": init_param[3]
        }

    def logging_acc(self, metrics_dict, prob_path = None, ave = False):
        """
        logging the metrics
        :param metrics_dict: dict of metrics
        :param prob_path: path
        :return: 0
        """

        if not ave:
            self.logger.info("Eval {}".format(prob_path))
        else:
            self.logger.info("+++++++++++++++Average++++++++++++++")

        self.logger.info("+++++++ Calc the SLM METRICS +++++++")
        for metric, value in metrics_dict.items():
            self.logger.info("++++     {}:{:.4f}   ++++".format(metric, value))
        self.logger.info("++++++++++++++++++++++++++++++++++++\n")

    def set_visual_options(self, visual_ras, src_addmap_path):
        """
        set visual options
        :param visual_ras: flag
        :param src_addmap_path: set src addmap path
        """
        self.visual_ras = visual_ras
        self.src_addmap_path = src_addmap_path
        return True

    def read_gray_to_prob(self, probmap_path):
        """
        Read the prob maps, and trans to probility
        :param probmap_path: probmap routh
        :return: probability
        """
        gray_image = cv2.imread(probmap_path, cv2.IMREAD_GRAYSCALE)
        prob = gray_image / 255.0
        return prob

    def generate_mask_by_points(self, prob, points_list):
        """
        Generate mask by regions
        :param prob: probability
        :param points_list: regions
        :return: mask
        """
        H, W = prob.shape

        mask = np.zeros((H, W))
        points_list = [np.array(i, np.int32) for i in points_list]
        # fill
        cv2.fillPoly(mask, points_list, 1)
        return mask

    def _get_region_center_radius(self, region_point):
        """
        get the region center and radius
        :param region_point: regions
        :return: mid_x, mid_y, radius
        """
        mid_x = int(reduce(lambda x, y: x+y, np.array(region_point)[:, 0]) / len(region_point))
        mid_y = int(reduce(lambda x, y: x+y, np.array(region_point)[:, 1]) / len(region_point))
        radius = int(np.mean([np.linalg.norm(np.array(point) - np.array([mid_x, mid_y])) for point in region_point]) * self.ras_expand_factor)
        return mid_x, mid_y, radius

    def _get_prob_center_in_gray(self, prob):
        """
        get the top point with the highest probability from the probability map
        :param prob: probability
        :return: centers
        """

        # recover the prob
        gray_img = np.asarray(prob * 255.0, dtype=np.uint8)
        # cv2.imwrite("./gray_img.jpg", gray_img)
        # construct continuous area
        continuous_area = np.asarray(gray_img > 150, np.uint8) * 255
        # cv2.imwrite("./continuous_area_img_0.jpg", continuous_area)
        continuous_area = np.uint8(measure.label(continuous_area, connectivity=2))
        # cv2.imwrite("./continuous_area_img_1.jpg", continuous_area)

        # soften
        for i in range(self.ras_filter_times):
            gray_img = cv2.boxFilter(gray_img, ddepth=-1, ksize=(50, 50))

        # get probability binary map
        mx = maximum_filter(gray_img, size=1000)
        gray_img = np.where(mx == gray_img, gray_img, 0)
        # cv2.imwrite("./local_maxima_before_filter.jpg", gray_img)
        gray_img = np.asarray(gray_img > 0, np.uint8) * 255
        # cv2.imwrite("./local_maxima_after_filter.jpg", gray_img)

        # get probability area information
        labels = measure.label(gray_img, connectivity=2)
        all_region_infos = measure.regionprops(labels)
        centers = [[int(i) for i in prop.centroid][::-1] for prop in all_region_infos]

        # construct v-center list and sort
        v_center = [[c[0], c[1], prob[c[1]][c[0]]] for c in centers]
        v_center.sort(key= lambda x: x[2], reverse=True)
        centers = list(map(lambda x: x[:2], v_center))

        # filter centers
        centers = [i for i in centers if prob[i[1]][i[0]] >= 0.5]

        return centers, continuous_area

    def _get_offset_between_real_and_synthetic(self, real_center_radius, prob_centers, bina_img):
        """
        calculate true center offset from result center
        :param real_center_radius: real_center_radius
        :param prob_centers: prob_centers
        :return: offsets
        """

        # check prob_centers is not None
        if len(prob_centers) == 0 : return [real_center_radius[0][2]]

        offsets = []
        for center_radius in real_center_radius:
            x, y, r = center_radius

            # calc the l2 dis
            dises = list(map(lambda p: np.linalg.norm(np.array([x, y] - np.array(p))), prob_centers))

            # filter the dis in cicle
            dises = list(filter(lambda d: d <= r, dises))

            # if no prob center set it to radius
            offsets.append(np.mean(dises) if len(dises) != 0 else r)

        return offsets

    def _trans_ras_offset_to_scalable_ras(self, offsets, centers_and_radius):
        """
        convert distance offset to ras value
        :param offsets: offsets
        :return: centers_and_radius
        """

        # granular transformation
        granular_offet = np.mean([off/v[2] for off, v in zip(offsets, centers_and_radius)])

        # scala transformation
        granular_offet = (np.exp(self.ras_scala_beta * granular_offet) - 1) / (np.exp(self.ras_scala_beta) - 1)

        return granular_offet

    def ras(self, region_lists, prob, visual=True, src_img=None):
        """
        calc the matric of ras: makes attention center close to annotation center
        :param region_lists: regions
        :param prob: probability
        :return: ras
        """

        # get the annotation center and radius
        centers_and_radius = [self._get_region_center_radius(i) for i in region_lists]

        # get the point with the highest probability from the probability map
        prob_centers, bina_img = self._get_prob_center_in_gray(prob)

        # calculate true center offset from result center
        offsets = self._get_offset_between_real_and_synthetic(centers_and_radius, prob_centers, bina_img)

        # convert distance offset to rcs value
        ras = self._trans_ras_offset_to_scalable_ras(offsets, centers_and_radius)

        # visual
        if visual and (src_img != None):
            src_img = cv2.imread(src_img)

            # logging something
            # print("centers_and_radius: ", centers_and_radius)
            # print("prob_centers: ", prob_centers)
            # print("offsets: ", offsets)

            # backup area
            for c_r in centers_and_radius:
                cv2.circle(src_img, (c_r[0], c_r[1]), c_r[2], 2, 3)

            # candidate points
            for idx, point in enumerate(prob_centers):
                cv2.circle(src_img, tuple(point), 6*(idx+1), 1, 4)
                cv2.putText(src_img, str(idx+1), tuple(point), cv2.FONT_HERSHEY_COMPLEX, 6, (0, 0, 0), 25)

            cv2.imwrite("./img_circle.jpg", src_img)

            # print(prob_centers)

        return ras

    def rsu(self, prob, mask):
        """
        calc the salient area proportion
        :param prob: probability
        :param mask: mask
        :return: rsu
        """

        all_mask_value = np.sum(np.multiply(prob, mask))
        all_value = np.sum(prob)
        H, W = np.shape(mask)
        all_mask = np.sum(mask)

        left_frac = all_mask_value / (all_value - all_mask_value + self.rsu_eps)

        right_frac = (H * W - all_mask) / all_mask

        rsu = -np.exp(-1 * self.rsu_beta * left_frac * right_frac) + 1

        return rsu

    def rda(self, region_lists, prob):
        """
        calc the matric of rda: makes attention center focus on one point
        :param region_lists: regions
        :param prob: probability
        :return: rda
        """

        # get the annotation center and radius
        centers_and_radius = [self._get_region_center_radius(i) for i in region_lists]

        # get the point with the highest probability from the probability map
        prob_centers, bina_img = self._get_prob_center_in_gray(prob)

        # set value
        rda = []
        for c_r in centers_and_radius:
            x, y, r = c_r

            # calc the backup points
            backup_points = list(filter(lambda p: np.linalg.norm(np.array([x, y] - np.array(p))) <= r, prob_centers))

            # margin condition
            len_backup_points = len(backup_points)
            if len_backup_points <= 1 :
                rda.append(float(len_backup_points))
                continue

            # if len_backup_points >= 2, calc the attention discrete
            centers_attention = np.average(backup_points, axis=0)
            dises = list(map(lambda p: np.linalg.norm(np.array(centers_attention - np.array(p))), backup_points))
            meas_dis = np.mean(dises) / r

            rda_single = 0.5 * (1 - meas_dis) + np.exp(- self.rda_eta * (len_backup_points + 2))

            rda.append(rda_single)

        return np.mean(rda)

    def rmi(self, rsu, rda, ras):
        """
        calculate the mean indicator
        :param rsu: rsu
        :param rda: rda
        :param ras: ras
        :return: rmi
        """
        return self.rmi_wsu * rsu + self.rmi_was * (1 - ras) + self.rmi_wda * rda

    def evaluate(self, prob_path, region_list):
        """
        evaluate the slm task
        :param probmap_path: probability map path
        :param region_list: region points
        :return: slm metrics
        """
        # read prob
        prob = self.read_gray_to_prob(prob_path)

        # generate mask
        mask = self.generate_mask_by_points(prob, region_list)
        # import os
        # cv2.imwrite(os.path.join(prob_path.rsplit("/", 1)[0], "maskbypt_0.jpg"), mask*255)
        # rsu
        rsu = self.rsu(prob, mask)

        # ras
        ras = self.ras(region_list, prob, visual=self.visual_ras, src_img=self.src_addmap_path)

        # rda
        rda = self.rda(region_list, prob)

        # mi
        rmi = self.rmi(rsu, rda, ras)

        # sort metrics
        metrics = self._format_output_dict(rsu, rda, ras, rmi)
        # self.logging_acc(metrics, prob_path)

        return metrics

    def append_metric(self, metric):
        """
        append metric to calc ave indicator
        :param metric: sort metrics
        """
        for k in metric.keys():
            self.all_metrics[k].append(metric[k])

    def get_the_mean_metric(self):
        """
        get the mean metric
        """
        mean_metric = {}
        for k in self.all_metrics:
            mean_metric[k] = np.mean(self.all_metrics[k])

        self.logging_acc(mean_metric, ave=True)
        return mean_metric


def semantic_localization_evaluation(model, selo_dataset, preprocess, args):
    assert selo_dataset == 'AIR-SLT'

    def collect_fn_selo(batch):
        assert len(batch) == 1
        source_img, subimages, text, points, subimg_name_list = batch[0]
        return source_img, subimages, text, points, subimg_name_list

    dataset = get_selo_dataset(
        root=args.test_dataset_dir, transform=preprocess, identifier=None
    )

    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=1,
        shuffle=False,
        num_workers=0,
        collate_fn=collect_fn_selo
    )
    tokenizer = open_clip.tokenize
    logger = dataset.logger
    slm_metric = SLM()

    with torch.no_grad():
        for idx, sample in tqdm.tqdm(enumerate(dataloader)):
            source_img, subimages, text, points, subimg_name_list = sample
            subimages = subimages.to(args.device)
            text = tokenizer(text).to(args.device)
            text_features = model.encode_text(text)
            text_features /= text_features.norm(dim=-1, keepdim=True)

            sim_results = []
            for subimage in subimages:
                subimage = subimage.unsqueeze(0)
                sub_img_feat = model.encode_image(subimage)
                sub_img_feat /= sub_img_feat.norm(dim=-1, keepdim=True)
                similarity = (sub_img_feat * text_features).sum().detach().cpu().numpy()
                sim_results.append(similarity)

            # print("Start generate heatmap ...")
            img_row = np.shape(source_img)[0]
            img_col = np.shape(source_img)[1]

            # mkdir map
            heat_map = np.zeros([img_row, img_col], dtype=float)
            heat_num = np.zeros([img_row, img_col], dtype=float)
            for idx, file in enumerate(subimg_name_list):
                r_start, r_end, c_start, c_end = file.replace(".jpg", "").split("_")
                heat_map[int(r_start):int(r_end), int(c_start):int(c_end)] += sim_results[idx]
                heat_num[int(r_start):int(r_end), int(c_start):int(c_end)] += 1

            for i in range(np.shape(heat_map)[0]):
                for j in range(np.shape(heat_map)[1]):
                    heat_map[i, j] = heat_map[i, j] / heat_num[i, j]

            # logger.info("Generation finished, start operating blur, colormap, etc. ...")
            # filter
            adaptive = np.asarray(heat_map)
            adaptive = adaptive - np.min(adaptive)
            probmap = adaptive / np.max(adaptive)
            # must convert to type unit8
            probmap = np.uint8(255 * probmap)
            probmap = cv2.medianBlur(probmap, 251)
            heatmap = cv2.applyColorMap(probmap, cv2.COLORMAP_JET)
            img_add = cv2.addWeighted(source_img, 0.7, heatmap, 0.3, 0)

            probmap_path = os.path.join(dataset.cache_path, "probmap_{}.jpg".format(idx))
            heatmap_path = os.path.join(dataset.cache_path, "heatmap_{}.jpg".format(idx))
            addmap_path = os.path.join(dataset.cache_path, "addmap_{}.jpg".format(idx))

            # logger.info("Saving heatmap in {} ...".format(heatmap_path))
            # logger.info("Saving probmap in {} ...".format(probmap_path))
            # logger.info("Saving addmap in {} ...".format(addmap_path))

            cv2.imwrite(probmap_path, probmap)
            cv2.imwrite(heatmap_path, heatmap)
            cv2.imwrite(addmap_path, img_add)
            # logger.info("Saved ok.")

            metrics = slm_metric.evaluate(probmap_path, region_list=points)
            slm_metric.append_metric(metrics)

        mean_metric = slm_metric.get_the_mean_metric()

    results = {}
    logging.info(f'{selo_dataset} selo metrics: {mean_metric}')

    for key, item in mean_metric.items():
        results[key] = float(item)

    return results


class AIR_SLT(Dataset):
    # Ref: https://github.com/xiaoyuan1996/SemanticLocalizationMetrics/blob/master/predict/generate_selo.py
    def __init__(self, root, subimage_transform, identifier):
        super().__init__()
        self.json_path = os.path.join(root, "annotations", "anno.json")
        # self.cache_path = os.path.join(root, "selo_cache_{}_{}".format(identifier, str(datetime.now()).replace(" ", "-").replace(":", "-").replace(".", "-")))
        self.cache_path = os.path.join(root, "selo_cache")
        os.makedirs(self.cache_path, exist_ok=True)
        with open(self.json_path, 'r', encoding='utf8') as fp:
            self.json_data = json.load(fp)
        self.img_root = os.path.join(root, "imgs")
        self.subimage_transform = subimage_transform
        self.logger = get_logger(os.path.join(self.cache_path, 'log.txt'))
        self.step = "256_512_768"

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

    def __getitem__(self, index):
        item = self.json_data[index]
        img_name = item['jpg_name']
        text = item['caption']
        points = item['points']
        steps = [int(step) for step in self.step.split("_")]
        img_path = os.path.join(self.img_root, img_name)

        # logging
        # self.logger.info("Processing {}/{}: {}".format(index, len(self.json_data), img_name))
        # self.logger.info("Corresponding text: {}".format(text))

        # processing
        self.split_image(img_path, steps)
        with torch.no_grad():
            subimages_dir = os.path.join(self.cache_path, os.path.basename(img_path).split(".")[0]) + '_subimages'
            subimages = os.listdir(subimages_dir)

            img = cv2.imread(img_path)
            subimg_list = []
            subimg_name_list = []
            for subimage_name in subimages:
                subimage_path = os.path.join(subimages_dir, subimage_name)
                subimg = Image.open(subimage_path)
                subimg = self.subimage_transform(subimg).unsqueeze(0)
                subimg_list.append(subimg)
                subimg_name_list.append(subimage_name)
            subimgs = torch.vstack(subimg_list)
        return img, subimgs, [text], points, subimg_name_list

    def split_image(self, img_path, steps):
        subimage_files_dir = os.path.join(self.cache_path, os.path.basename(img_path).split(".")[0])

        # 裁切图像文件夹
        subimages_dir = subimage_files_dir + '_subimages'
        if os.path.exists(subimages_dir):
            delete_dire(subimages_dir)
        else:
            os.makedirs(subimages_dir)

        # Read Image
        source_img = cv2.imread(img_path)
        img_weight = np.shape(source_img)[0]
        img_height = np.shape(source_img)[1]
        # self.logger.info("img size:{}x{}".format(img_weight, img_height))

        for step in steps:
            # self.logger.info("Start split images with step {}".format(step))
            for gap in [step, 0.5 * step]:
                gap = int(gap)

                # Cut img
                for h in range(0 + (step - gap), img_height, step):
                    h_start, h_end = h, h + step
                    # bound?
                    if h_end >= img_height:
                        h_start, h_end = img_height - step, img_height

                    for w in range(0 + (step - gap), img_weight, step):
                        w_start, w_end = w, w + step
                        # bound?
                        if w_end >= img_weight:
                            w_start, w_end = img_weight - step, img_weight

                        cut_img_name = str(w_start) + "_" + str(w_end) + "_" + str(h_start) + "_" + str(h_end) + ".jpg"
                        cut_img = source_img[w_start:w_end, h_start:h_end]
                        cut_img = cv2.resize(cut_img, (256, 256), interpolation=cv2.INTER_CUBIC)

                        cv2.imwrite(os.path.join(subimages_dir, cut_img_name), cut_img)

        # self.logger.info("Image {} has been split successfully.".format(img_path))


def delete_dire(dire):
    dir_list = []
    for root, dirs, files in os.walk(dire):
        for afile in files:
            os.remove(os.path.join(root, afile))
        for adir in dirs:
            dir_list.append(os.path.join(root, adir))
    for bdir in dir_list:
        os.rmdir(bdir)


# logger
def get_logger(save_path=None):
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)  # 设置打印级别
    formatter = logging.Formatter('%(asctime)s %(message)s')

    # 设置屏幕打印的格式
    sh = logging.StreamHandler()
    sh.setFormatter(formatter)
    logger.addHandler(sh)

    # 设置log保存
    if save_path != None:
        fh = logging.FileHandler(save_path, encoding='utf8')
        fh.setFormatter(formatter)
        logger.addHandler(fh)

    return logger


def get_selo_dataset(root, transform, identifier):

    AIR_SLT_root = os.path.join(root, "AIR-SLT")
    dataset = AIR_SLT(
        root=AIR_SLT_root,
        subimage_transform=transform,
        identifier=identifier
    )

    return dataset