test_regdb.py

# -*- coding: utf-8 -*-
from __future__ import print_function, absolute_import

import argparse
import os.path as osp
import random
import time

import numpy as np
import torch
import torch.utils.data as data
from PIL import Image
from torch import nn
from torch.autograd import Variable
from torch.backends import cudnn
from torch.utils.data import DataLoader

from clustercontrast import datasets
from clustercontrast import models
from clustercontrast.utils.data import IterLoader
from clustercontrast.utils.data import transforms as T
from clustercontrast.utils.data.preprocessor import Preprocessor, Preprocessor_color
from clustercontrast.utils.data.sampler import RandomMultipleGallerySampler, RandomMultipleGallerySamplerNoCam
from clustercontrast.utils.serialization import load_checkpoint

start_epoch = best_mAP = 0


def get_data(name, data_dir, trial=0):
    root = osp.join(data_dir, name)
    dataset = datasets.create(name, root, trial=trial)
    return dataset


def get_train_loader_ir(args, dataset, height, width, batch_size, workers, num_instances, iters, trainset=None,
                        no_cam=False, train_transformer=None):
    train_set = sorted(dataset.train) if trainset is None else sorted(trainset)
    rmgs_flag = num_instances > 0
    if rmgs_flag:
        if no_cam:
            sampler = RandomMultipleGallerySamplerNoCam(train_set, num_instances)
        else:
            sampler = RandomMultipleGallerySampler(train_set, num_instances)
    else:
        sampler = None
    train_loader = IterLoader(
        DataLoader(Preprocessor(train_set, root=dataset.images_dir, transform=train_transformer), batch_size=batch_size,
                   num_workers=workers, sampler=sampler, shuffle=not rmgs_flag, pin_memory=True, drop_last=True),
        length=iters)

    return train_loader


def get_train_loader_color(args, dataset, height, width, batch_size, workers, num_instances, iters, trainset=None,
                           no_cam=False, train_transformer=None, train_transformer1=None):
    train_set = sorted(dataset.train) if trainset is None else sorted(trainset)
    rmgs_flag = num_instances > 0
    if rmgs_flag:
        if no_cam:
            sampler = RandomMultipleGallerySamplerNoCam(train_set, num_instances)
        else:
            sampler = RandomMultipleGallerySampler(train_set, num_instances)
    else:
        sampler = None
    if train_transformer1 is None:
        train_loader = IterLoader(
            DataLoader(Preprocessor(train_set, root=dataset.images_dir, transform=train_transformer),
                       batch_size=batch_size, num_workers=workers, sampler=sampler, shuffle=not rmgs_flag,
                       pin_memory=True, drop_last=True), length=iters)
    else:
        train_loader = IterLoader(DataLoader(
            Preprocessor_color(train_set, root=dataset.images_dir, transform=train_transformer,
                               transform1=train_transformer1), batch_size=batch_size, num_workers=workers,
            sampler=sampler, shuffle=not rmgs_flag, pin_memory=True, drop_last=True), length=iters)

    return train_loader


def get_test_loader(dataset, height, width, batch_size, workers, testset=None, test_transformer=None):
    normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    if test_transformer is None:
        test_transformer = T.Compose([T.Resize((height, width), interpolation=3), T.ToTensor(), normalizer])

    if testset is None:
        testset = list(set(dataset.query) | set(dataset.gallery))

    test_loader = DataLoader(Preprocessor(testset, root=dataset.images_dir, transform=test_transformer),
                             batch_size=batch_size, num_workers=workers, shuffle=False, pin_memory=True)

    return test_loader


def create_model(args):
    model = models.create(args.arch, num_features=args.features, norm=True, dropout=args.dropout, num_classes=0,
                          pooling_type=args.pooling_type)
    # use CUDA
    model.cuda()
    model = nn.DataParallel(model)  # ,output_device=1)
    return model


def main():
    args = parser.parse_args()

    if args.seed is not None:
        random.seed(args.seed)
        np.random.seed(args.seed)
        torch.manual_seed(args.seed)
        cudnn.deterministic = True
    main_worker(args)


class TestData(data.Dataset):
    def __init__(self, test_img_file, test_label, transform=None, img_size=(144, 288)):
        test_image = []
        for i in range(len(test_img_file)):
            img = Image.open(test_img_file[i])
            img = img.resize((img_size[0], img_size[1]), Image.ANTIALIAS)
            pix_array = np.array(img)
            test_image.append(pix_array)
        test_image = np.array(test_image)
        self.test_image = test_image
        self.test_label = test_label
        self.transform = transform

    def __getitem__(self, index):
        img1, target1 = self.test_image[index], self.test_label[index]
        img1 = self.transform(img1)
        return img1, target1

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


def fliplr(img):
    '''flip horizontal'''
    inv_idx = torch.arange(img.size(3) - 1, -1, -1).long()  # N x C x H x W
    img_flip = img.index_select(3, inv_idx)
    return img_flip


def extract_gall_feat(model, gall_loader, ngall):
    pool_dim = 2048
    net = model
    net.eval()
    print('Extracting Gallery Feature...')
    start = time.time()
    ptr = 0
    gall_feat_pool = np.zeros((ngall, pool_dim))
    gall_feat_fc = np.zeros((ngall, pool_dim))
    with torch.no_grad():
        for batch_idx, (input, label) in enumerate(gall_loader):
            batch_num = input.size(0)
            flip_input = fliplr(input)
            input = Variable(input.cuda())
            feat_fc = net(input, input, 2)
            flip_input = Variable(flip_input.cuda())
            feat_fc_1 = net(flip_input, flip_input, 2)
            feature_fc = (feat_fc.detach() + feat_fc_1.detach()) / 2
            fnorm_fc = torch.norm(feature_fc, p=2, dim=1, keepdim=True)
            feature_fc = feature_fc.div(fnorm_fc.expand_as(feature_fc))
            gall_feat_fc[ptr:ptr + batch_num, :] = feature_fc.cpu().numpy()
            ptr = ptr + batch_num
    print('Extracting Time:\t {:.3f}'.format(time.time() - start))
    return gall_feat_fc


def extract_query_feat(model, query_loader, nquery):
    pool_dim = 2048
    net = model
    net.eval()
    print('Extracting Query Feature...')
    start = time.time()
    ptr = 0
    query_feat_pool = np.zeros((nquery, pool_dim))
    query_feat_fc = np.zeros((nquery, pool_dim))
    with torch.no_grad():
        for batch_idx, (input, label) in enumerate(query_loader):
            batch_num = input.size(0)
            flip_input = fliplr(input)
            input = Variable(input.cuda())
            feat_fc = net(input, input, 1)
            flip_input = Variable(flip_input.cuda())
            feat_fc_1 = net(flip_input, flip_input, 1)
            feature_fc = (feat_fc.detach() + feat_fc_1.detach()) / 2
            fnorm_fc = torch.norm(feature_fc, p=2, dim=1, keepdim=True)
            feature_fc = feature_fc.div(fnorm_fc.expand_as(feature_fc))
            query_feat_fc[ptr:ptr + batch_num, :] = feature_fc.cpu().numpy()

            ptr = ptr + batch_num
    print('Extracting Time:\t {:.3f}'.format(time.time() - start))
    return query_feat_fc


def pairwise_distance(features_q, features_g):
    x = torch.from_numpy(features_q)
    y = torch.from_numpy(features_g)
    m, n = x.size(0), y.size(0)
    x = x.view(m, -1)
    y = y.view(n, -1)
    dist_m = torch.pow(x, 2).sum(dim=1, keepdim=True).expand(m, n) + torch.pow(y, 2).sum(dim=1, keepdim=True).expand(n,
                                                                                                                     m).t()
    dist_m.addmm_(1, -2, x, y.t())
    return dist_m.numpy()


def process_test_regdb(img_dir, trial=1, modal='visible'):
    if modal == 'visible':
        input_data_path = img_dir + 'idx/test_visible_{}'.format(trial) + '.txt'
    elif modal == 'thermal':
        input_data_path = img_dir + 'idx/test_thermal_{}'.format(trial) + '.txt'

    with open(input_data_path) as f:
        data_file_list = open(input_data_path, 'rt').read().splitlines()
        # Get full list of image and labels
        file_image = [img_dir + '/' + s.split(' ')[0] for s in data_file_list]
        file_label = [int(s.split(' ')[1]) for s in data_file_list]

    return file_image, np.array(file_label)


def eval_regdb(distmat, q_pids, g_pids, max_rank=20):
    num_q, num_g = distmat.shape
    if num_g < max_rank:
        max_rank = num_g
        print("Note: number of gallery samples is quite small, got {}".format(num_g))
    indices = np.argsort(distmat, axis=1)
    matches = (g_pids[indices] == q_pids[:, np.newaxis]).astype(np.int32)

    # compute cmc curve for each query
    all_cmc = []
    all_AP = []
    all_INP = []
    num_valid_q = 0.  # number of valid query

    # only two cameras
    q_camids = np.ones(num_q).astype(np.int32)
    g_camids = 2 * np.ones(num_g).astype(np.int32)

    for q_idx in range(num_q):
        # get query pid and camid
        q_pid = q_pids[q_idx]
        q_camid = q_camids[q_idx]

        # remove gallery samples that have the same pid and camid with query
        order = indices[q_idx]
        remove = (g_pids[order] == q_pid) & (g_camids[order] == q_camid)
        keep = np.invert(remove)

        # compute cmc curve
        raw_cmc = matches[q_idx][keep]  # binary vector, positions with value 1 are correct matches
        if not np.any(raw_cmc):
            # this condition is true when query identity does not appear in gallery
            continue

        cmc = raw_cmc.cumsum()

        # compute mINP
        # refernece Deep Learning for Person Re-identification: A Survey and Outlook
        pos_idx = np.where(raw_cmc == 1)
        pos_max_idx = np.max(pos_idx)
        inp = cmc[pos_max_idx] / (pos_max_idx + 1.0)
        all_INP.append(inp)

        cmc[cmc > 1] = 1

        all_cmc.append(cmc[:max_rank])
        num_valid_q += 1.

        # compute average precision
        # reference: https://en.wikipedia.org/wiki/Evaluation_measures_(information_retrieval)#Average_precision
        num_rel = raw_cmc.sum()
        tmp_cmc = raw_cmc.cumsum()
        tmp_cmc = [x / (i + 1.) for i, x in enumerate(tmp_cmc)]
        tmp_cmc = np.asarray(tmp_cmc) * raw_cmc
        AP = tmp_cmc.sum() / num_rel
        all_AP.append(AP)

    assert num_valid_q > 0, "Error: all query identities do not appear in gallery"

    all_cmc = np.asarray(all_cmc).astype(np.float32)
    all_cmc = all_cmc.sum(0) / num_valid_q
    mAP = np.mean(all_AP)
    mINP = np.mean(all_INP)
    return all_cmc, mAP, mINP


def main_worker(args):
    log_name = 'regdb_s2'  # model path
    model = create_model(args)
    for trial in range(1, 11):  # (1,11):

        normalize = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        transform_test = T.Compose([T.ToPILImage(), T.Resize((args.height, args.width)), T.ToTensor(), normalize, ])
        logs_dir_root = osp.join(args.logs_dir + '/' + log_name)
        # args.logs_dir = osp.join(logs_dir_root,str(trial))
        print('==> Test with the best model:')

        checkpoint = load_checkpoint(osp.join(logs_dir_root + '/' + str(trial), "model_best.pth.tar"))
        model.load_state_dict(checkpoint['state_dict'])

        mode = 'visible to thermal'
        print(mode)
        data_path = './data/regdb/'
        query_img, query_label = process_test_regdb(data_path, trial=trial, modal='visible')
        gall_img, gall_label = process_test_regdb(data_path, trial=trial, modal='thermal')

        gallset = TestData(gall_img, gall_label, transform=transform_test, img_size=(args.width, args.height))
        gall_loader = data.DataLoader(gallset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers)
        nquery = len(query_label)
        ngall = len(gall_label)
        queryset = TestData(query_img, query_label, transform=transform_test, img_size=(args.width, args.height))
        query_loader = data.DataLoader(queryset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers)
        query_feat_fc = extract_query_feat(model, query_loader, nquery)
        # for trial in range(1):
        ngall = len(gall_label)
        gall_feat_fc = extract_gall_feat(model, gall_loader, ngall)
        # fc feature
        distmat = np.matmul(query_feat_fc, np.transpose(gall_feat_fc))
        cmc, mAP, mINP = eval_regdb(-distmat, query_label, gall_label)
        if trial == 1:
            all_cmc = cmc
            all_mAP = mAP
            all_mINP = mINP

        else:
            all_cmc = all_cmc + cmc
            all_mAP = all_mAP + mAP
            all_mINP = all_mINP + mINP

        print('Test Trial: {}'.format(trial))
        print(
            'FC:   Rank-1: {:.2%} | Rank-5: {:.2%} | Rank-10: {:.2%}| Rank-20: {:.2%}| mAP: {:.2%}| mINP: {:.2%}'.format(
                cmc[0], cmc[4], cmc[9], cmc[19], mAP, mINP))
    cmc = all_cmc / 10
    mAP = all_mAP / 10
    mINP = all_mINP / 10
    print('All Average:')
    print(
        'FC:     Rank-1: {:.2%} | Rank-5: {:.2%} | Rank-10: {:.2%}| Rank-20: {:.2%}| mAP: {:.2%}| mINP: {:.2%}'.format(
            cmc[0], cmc[4], cmc[9], cmc[19], mAP, mINP))
    #################################
    for trial in range(1, 11):  # (1,11):
        normalize = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        transform_test = T.Compose([T.ToPILImage(), T.Resize((args.height, args.width)), T.ToTensor(), normalize, ])
        logs_dir_root = osp.join(args.logs_dir + '/' + log_name)
        # args.logs_dir = osp.join(logs_dir_root,str(trial))
        print('==> Test with the best model:')
        model = create_model(args)
        checkpoint = load_checkpoint(osp.join(logs_dir_root + '/' + str(trial), 'model_best.pth.tar'))

        model.load_state_dict(checkpoint['state_dict'])

        mode = 'thermal to visible'
        print(mode)
        data_path = './data/regdb/'
        query_img, query_label = process_test_regdb(data_path, trial=trial, modal='thermal')
        gall_img, gall_label = process_test_regdb(data_path, trial=trial, modal='visible')

        gallset = TestData(gall_img, gall_label, transform=transform_test, img_size=(args.width, args.height))
        gall_loader = data.DataLoader(gallset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers)
        nquery = len(query_label)
        ngall = len(gall_label)
        queryset = TestData(query_img, query_label, transform=transform_test, img_size=(args.width, args.height))
        query_loader = data.DataLoader(queryset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers)
        query_feat_fc = extract_gall_feat(model, query_loader, nquery)
        # for trial in range(1):
        ngall = len(gall_label)
        gall_feat_fc = extract_query_feat(model, gall_loader, ngall)
        # fc feature
        distmat = np.matmul(query_feat_fc, np.transpose(gall_feat_fc))
        cmc, mAP, mINP = eval_regdb(-distmat, query_label, gall_label)
        if trial == 1:
            all_cmc = cmc
            all_mAP = mAP
            all_mINP = mINP

        else:
            all_cmc = all_cmc + cmc
            all_mAP = all_mAP + mAP
            all_mINP = all_mINP + mINP

        print('Test Trial: {}'.format(trial))
        print(
            'FC:   Rank-1: {:.2%} | Rank-5: {:.2%} | Rank-10: {:.2%}| Rank-20: {:.2%}| mAP: {:.2%}| mINP: {:.2%}'.format(
                cmc[0], cmc[4], cmc[9], cmc[19], mAP, mINP))
    cmc = all_cmc / 10
    mAP = all_mAP / 10
    mINP = all_mINP / 10
    print('All Average:')
    print(
        'FC:     Rank-1: {:.2%} | Rank-5: {:.2%} | Rank-10: {:.2%}| Rank-20: {:.2%}| mAP: {:.2%}| mINP: {:.2%}'.format(
            cmc[0], cmc[4], cmc[9], cmc[19], mAP, mINP))


#################################


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Unsupervised Visible-Infrared Person Re-Identification via Dual-Modality-Shared Modality-Shared Learning and Label Refinement")

    # data
    parser.add_argument('-d', '--dataset', type=str, default='regdb', choices=datasets.names())
    parser.add_argument('-b', '--batch_size', type=int, default=8)
    parser.add_argument('-j', '--workers', type=int, default=8)
    parser.add_argument('--height', type=int, default=288, help="input height")
    parser.add_argument('--width', type=int, default=144, help="input width")
    parser.add_argument('--num_instances', type=int, default=8,
                        help="each mini-batch consist of (batch_size // num_instances) identities, and each identity has num_instances instances, default: 0 (NOT USE)")
    parser.add_argument('--mode', type=str, default='v2t', help='all or indoor (sysu test), t2v or v2t (regdb test)')
    parser.add_argument('--shot', default=1, type=int, help='1 for single shot;10 for multi shot')
    parser.add_argument('--trial', default=1, type=int, help='trial (only for RegDB dataset)')

    # cluster
    parser.add_argument('--eps_ir', type=float, default=0.3,
                        help="max neighbor distance of thermal modality for DBSCAN")
    parser.add_argument('--eps_rgb', type=float, default=0.3,
                        help="max neighbor distance of visible modality for DBSCAN")
    parser.add_argument('--eps_sh', type=float, default=0.3, help="max neighbor distance of shared modality for DBSCAN")
    parser.add_argument('--eps_gap', type=float, default=0.02,
                        help="multi-scale criterion for measuring cluster reliability")
    parser.add_argument('--k1', type=int, default=30, help="hyperparameter for jaccard distance")
    parser.add_argument('--k2', type=int, default=6, help="hyperparameter for jaccard distance")

    # model
    parser.add_argument('-a', '--arch', type=str, default='resnet', choices=models.names())
    parser.add_argument('--features', type=int, default=0)
    parser.add_argument('--dropout', type=float, default=0)
    parser.add_argument('--pooling-type', type=str, default='gem')
    parser.add_argument('--momentum', type=float, default=0.1, help="update momentum for the hybrid memory")

    # optimizer
    parser.add_argument('--lr', type=float, default=0.00035, help="learning rate")
    parser.add_argument('--weight-decay', type=float, default=5e-4)
    parser.add_argument('--epochs', type=int, default=50)
    parser.add_argument('--iters', type=int, default=400)
    parser.add_argument('--step-size', type=int, default=20)

    # training configs
    parser.add_argument('--seed', type=int, default=1)
    parser.add_argument('--print-freq', type=int, default=10)
    parser.add_argument('--eval-step', type=int, default=1)

    # path
    working_dir = osp.dirname(osp.abspath(__file__))
    parser.add_argument('--data-dir', type=str, metavar='PATH', default=osp.join(working_dir, 'data'))
    parser.add_argument('--logs-dir', type=str, metavar='PATH', default=osp.join(working_dir, 'logs'))
    parser.add_argument('--use-hard', action="store_true")
    parser.add_argument('--no-cam', action="store_true")

    main()