utils.py

import os
import numpy as np
from torch.utils.data.sampler import Sampler
import sys
import os.path as osp
import torch
import matplotlib.pyplot as plt
import random

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

def GenIdx( train_color_label, train_thermal_label):
    color_pos = []
    unique_label_color = np.unique(train_color_label)
    for i in range(len(unique_label_color)):
        tmp_pos = [k for k,v in enumerate(train_color_label) if v==unique_label_color[i]]
        color_pos.append(tmp_pos)
        
    thermal_pos = []
    unique_label_thermal = np.unique(train_thermal_label)
    for i in range(len(unique_label_thermal)):
        tmp_pos = [k for k,v in enumerate(train_thermal_label) if v==unique_label_thermal[i]]
        thermal_pos.append(tmp_pos)
    return color_pos, thermal_pos
    
def GenCamIdx(gall_img, gall_label, mode):
    if mode =='indoor':
        camIdx = [1,2]
    else:
        camIdx = [1,2,4,5]
    gall_cam = []
    for i in range(len(gall_img)):
        gall_cam.append(int(gall_img[i][-10]))
    
    sample_pos = []
    unique_label = np.unique(gall_label)
    for i in range(len(unique_label)):
        for j in range(len(camIdx)):
            id_pos = [k for k,v in enumerate(gall_label) if v==unique_label[i] and gall_cam[k]==camIdx[j]]
            if id_pos:
                sample_pos.append(id_pos)
    return sample_pos
    
def ExtractCam(gall_img):
    gall_cam = []
    for i in range(len(gall_img)):
        cam_id = int(gall_img[i][-10])
        # if cam_id ==3:
            # cam_id = 2
        gall_cam.append(cam_id)
    
    return np.array(gall_cam)


class IdentitySampler(Sampler):
    """Sample person identities evenly in each batch.
        Args:
            train_color_label, train_thermal_label: labels of two modalities
            color_pos, thermal_pos: positions of each identity
            batchSize: batch size
    """

    def __init__(self, train_color_label, train_thermal_label, color_pos, thermal_pos, num_pos, batchSize, epoch):
        uni_label = np.unique(train_color_label)
        self.n_classes = len(uni_label)

        N = np.maximum(len(train_color_label), len(train_thermal_label))
        for j in range(int(N / (batchSize * num_pos)) + 1):
            batch_idx = np.random.choice(uni_label, batchSize, replace=False)
            for i in range(batchSize):
                while len(color_pos[batch_idx[i]]) < 4 or len(thermal_pos[batch_idx[i]]) < 4:
                    batch_idx[i] = np.random.choice(uni_label, 1, replace=False)
                    print("re-sampling")
                sample_color = np.random.choice(color_pos[batch_idx[i]], num_pos)
                sample_thermal = np.random.choice(thermal_pos[batch_idx[i]], num_pos)

                if j == 0 and i == 0:
                    index1 = sample_color
                    index2 = sample_thermal
                else:
                    index1 = np.hstack((index1, sample_color))
                    index2 = np.hstack((index2, sample_thermal))
        self.index1 = index1
        self.index2 = index2
        self.N = N

    def __iter__(self):
        return iter(np.arange(len(self.index1)))

    def __len__(self):
        return self.N


class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count
 
def mkdir_if_missing(directory):
    if not osp.exists(directory):
        try:
            os.makedirs(directory)
        except OSError as e:
            if e.errno != errno.EEXIST:
                raise

class Logger(object):
    """
    Write console output to external text file.
    Code imported from https://github.com/Cysu/open-reid/blob/master/reid/utils/logging.py.
    """  
    def __init__(self, fpath=None):
        self.console = sys.stdout
        self.file = None
        if fpath is not None:
            mkdir_if_missing(osp.dirname(fpath))
            self.file = open(fpath, 'w')

    def __del__(self):
        self.close()

    def __enter__(self):
        pass

    def __exit__(self, *args):
        self.close()

    def write(self, msg):
        self.console.write(msg)
        if self.file is not None:
            self.file.write(msg)

    def flush(self):
        self.console.flush()
        if self.file is not None:
            self.file.flush()
            os.fsync(self.file.fileno())

    def close(self):
        self.console.close()
        if self.file is not None:
            self.file.close()
            
def set_seed(seed, cuda=True):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if cuda:
        torch.backends.cudnn.deterministic = True
        torch.cuda.manual_seed(seed)

def set_requires_grad(nets, requires_grad=False):
            """Set requies_grad=Fasle for all the networks to avoid unnecessary computations
            Parameters:
                nets (network list)   -- a list of networks
                requires_grad (bool)  -- whether the networks require gradients or not
            """
            if not isinstance(nets, list):
                nets = [nets]
            for net in nets:
                if net is not None:
                    for param in net.parameters():
                        param.requires_grad = requires_grad
def plot_pair_distribution(type, X, clean_index, noisy_index, save_path=''):
    plt.clf()
    ax = plt.gca()

    font1 = {'family': 'Times New Roman',
             'weight': 'normal',
             'size': 13,
             }

    # Plot data histogram
    if type == 0:
        ax.hist(X[clean_index], bins=100, density=True, histtype='stepfilled', color='darkorchid', alpha=0.3,
                label='True Pos. Pairs')
        ax.hist(X[noisy_index], bins=100, density=True, histtype='stepfilled', color='blue', alpha=0.3,
                label='False Pos. Pairs')
    if type == 1:
        ax.hist(X[clean_index], bins=100, density=True, histtype='stepfilled', color='teal', alpha=0.3,
                label='True Neg. Pairs')
        ax.hist(X[noisy_index], bins=100, density=True, histtype='stepfilled', color='peru', alpha=0.3,
                label='False Neg. Pairs')

    ax.set_xlabel('Normalized Distances', fontdict=font1)
    ax.set_ylabel('Frequency', fontdict=font1)
    x_ticks = np.array([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
    plt.xticks(x_ticks)
    plt.tick_params(labelsize=11)
    ax.legend(loc='upper left', prop=font1)

    if save_path:
        plt.savefig(save_path, dpi=300)
    else:
        plt.show()
                        
class AllSampler(Sampler):
    def __init__(self, dataset, train_color_label, train_thermal_label, shuffle=True):
        N1 = len(train_color_label)
        N2 = len(train_thermal_label)
        # N = np.maximum(len(train_color_label), len(train_thermal_label))
        if dataset == 'regdb':
            index1 = np.concatenate((np.arange(N1), np.arange(20)))
            index2 = np.concatenate((np.arange(N2), np.arange(20)))
        else:
            index1 = np.concatenate((np.arange(N1), np.arange(14)))
            index2 = np.concatenate((np.arange(N2), np.arange(N1 - N2 + 14)))

        if shuffle:
            np.random.shuffle(index1)
            np.random.shuffle(index2)

        self.index1 = index1
        self.index2 = index2
        self.N = len(index1)

    def __iter__(self):
        return iter(np.arange(len(self.index1)))

    def __len__(self):
        return self.N