nuscene_seqnet.py

import torch
import torchvision.transforms as transforms
import torch.utils.data as data

from os.path import join, exists
from scipy.io import loadmat
import numpy as np
from random import randint, random
from collections import namedtuple
from PIL import Image

from sklearn.neighbors import NearestNeighbors
import h5py
import faiss
import os


def load_pc_file(opt, filename):
    pcs = []
    # edge_indices = set(range(-(opt.seqLen // 2), opt.seqLen // 2 + 1, 1))
    if opt.seqLen == 1:
        edge_indices = [0]
    elif opt.seqLen == 2:
        edge_indices = [-1, 0]
    elif opt.seqLen == 3:
        edge_indices = [-2, -1, 0]
    elif opt.seqLen == 4:
        edge_indices = [-3, -2, -1, 0]
    elif opt.seqLen == 5:
        edge_indices = [-4, -3, -2, -1, 0]
    file_index = int(filename[-9:-4])
    for offset in edge_indices:
        filename_spread = os.path.join(filename[:-9], '{:0>5d}.bin'.format(int(file_index + offset)))
        pc = np.fromfile(filename_spread, dtype=np.float64).reshape(-1, 1)
        pc = pc.astype(np.float32)
        pcs.append(pc)
    pcs = np.array(pcs)
    pcs = torch.from_numpy(pcs)
    return pcs


def get_whole_training_set(opt, onlyDB=False):
    return WholeDatasetFromStruct(opt, join(opt.structDir, 'nuscenes_train.mat'), opt.imgDir, onlyDB=onlyDB)


def get_whole_val_set(opt):
    return WholeDatasetFromStruct(opt, join(opt.structDir, 'nuscenes_val.mat'), opt.imgDir)

def get_whole_test_set(opt):
    return WholeDatasetFromStruct(opt, join(opt.structDir, 'nuscenes_test.mat'), opt.imgDir)

def get_training_query_set(opt, margin=0.1):
    return QueryDatasetFromStruct(opt, join(opt.structDir, 'nuscenes_train.mat'), opt.imgDir, margin=margin)


dbStruct = namedtuple('dbStruct', ['whichSet', 'dataset', 'dbImage', 'utmDb', 'qImage', 'utmQ', 'numDb', 'numQ', 'posDistThr', 'posDistSqThr', 'nonTrivPosDistSqThr'])


def parse_dbStruct(path):
    mat = loadmat(path)
    matStruct = mat['dbStruct'].item()
    dataset = 'pitts30k'

    whichSet = matStruct[0].item()

    # .mat file is generated by python, I replace the use of cell (in Matlab) with char (in Python)
    # dbImage = [f[0].item() for f in matStruct[1]]
    dbImage = matStruct[1]
    utmDb = matStruct[2].T

    # .mat file is generated by python, I replace the use of cell (in Matlab) with char (in Python)
    # qImage = [f[0].item() for f in matStruct[3]]
    qImage = matStruct[3]
    utmQ = matStruct[4].T

    numDb = matStruct[5].item()
    numQ = matStruct[6].item()

    posDistThr = matStruct[7].item()
    posDistSqThr = matStruct[8].item()
    nonTrivPosDistSqThr = matStruct[9].item()

    return dbStruct(whichSet, dataset, dbImage, utmDb, qImage, utmQ, numDb, numQ, posDistThr, posDistSqThr, nonTrivPosDistSqThr)


class WholeDatasetFromStruct(data.Dataset):
    def __init__(self, opt, structFile, img_dir, onlyDB=False):
        super().__init__()
        self.opt = opt
        self.dbStruct = parse_dbStruct(structFile)

        self.images = [join(img_dir, dbIm.split('.')[0] + '.bin') for dbIm in self.dbStruct.dbImage]
        if not onlyDB:
            self.images += [join(img_dir, qIm.split('.')[0] + '.bin') for qIm in self.dbStruct.qImage]

        self.whichSet = self.dbStruct.whichSet
        self.dataset = self.dbStruct.dataset

        self.positives = None
        self.distances = None

    def __getitem__(self, index):
        pcl = load_pc_file(self.opt, self.images[index])

        return pcl, index

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

    def getPositives(self):
        # positives for evaluation are those within trivial threshold range
        # fit NN to find them, search by radius
        if self.positives is None:
            knn = NearestNeighbors(n_jobs=-1)
            knn.fit(self.dbStruct.utmDb)

            # !! attention: radius dramatically change the PR performance!!
            # self.distances, self.positives = knn.radius_neighbors(self.dbStruct.utmQ, radius=self.dbStruct.posDistThr)

            # !! attention: note that the results are not sorted by default
            self.distances, self.positives = knn.radius_neighbors(self.dbStruct.utmQ, radius=self.dbStruct.nonTrivPosDistSqThr**0.5)

        return self.positives


def collate_fn(batch):
    """Creates mini-batch tensors from the list of tuples (query, positive, negatives).

    Args:
        data: list of tuple (query, positive, negatives). 
            - query: torch tensor of shape (3, h, w).
            - positive: torch tensor of shape (3, h, w).
            - negative: torch tensor of shape (n, 3, h, w).
    Returns:
        query: torch tensor of shape (batch_size, 3, h, w).
        positive: torch tensor of shape (batch_size, 3, h, w).
        negatives: torch tensor of shape (batch_size, n, 3, h, w).
    """

    batch = list(filter(lambda x: x is not None, batch))
    if len(batch) == 0: return None, None, None, None, None

    query, positive, negatives, indices = zip(*batch)

    query = data.dataloader.default_collate(query)
    positive = data.dataloader.default_collate(positive)
    negCounts = data.dataloader.default_collate([x.shape[0] for x in negatives])
    negatives = torch.cat(negatives, 0)
    import itertools
    indices = list(itertools.chain(*indices))

    return query, positive, negatives, negCounts, indices


class QueryDatasetFromStruct(data.Dataset):
    def __init__(self, opt, structFile, img_dir, nNegSample=1000, nNeg=10, margin=0.1):
        super().__init__()
        self.opt = opt
        self.img_dir = img_dir
        self.margin = margin

        self.dbStruct = parse_dbStruct(structFile)
        self.whichSet = self.dbStruct.whichSet
        self.dataset = self.dbStruct.dataset
        self.nNegSample = nNegSample  # number of negatives to randomly sample
        self.nNeg = nNeg  # number of negatives used for training

        # potential positives are those within nontrivial threshold range
        # fit NN to find them, search by radius
        knn = NearestNeighbors(n_jobs=-1)
        knn.fit(self.dbStruct.utmDb)

        # TODO use sqeuclidean as metric?
        self.nontrivial_positives = list(knn.radius_neighbors(self.dbStruct.utmQ, radius=self.dbStruct.nonTrivPosDistSqThr**0.5, return_distance=False))
        # radius returns unsorted, sort once now so we dont have to later
        for i, posi in enumerate(self.nontrivial_positives):
            self.nontrivial_positives[i] = np.sort(posi)
        # its possible some queries don't have any non trivial potential positives
        # lets filter those out
        self.queries = np.where(np.array([len(x) for x in self.nontrivial_positives]) > 0)[0]

        # potential negatives are those outside of posDistThr range
        potential_positives = knn.radius_neighbors(self.dbStruct.utmQ, radius=self.dbStruct.posDistThr, return_distance=False)

        self.potential_negatives = []
        for pos in potential_positives:
            self.potential_negatives.append(np.setdiff1d(np.arange(self.dbStruct.numDb), pos, assume_unique=True))

        self.cache = None  # filepath of HDF5 containing feature vectors for images

        self.negCache = [np.empty((0, )) for _ in range(self.dbStruct.numQ)]

    def __getitem__(self, index):

        index = self.queries[index]  # re-map index to match dataset
        with h5py.File(self.cache, mode='r') as h5:
            h5feat = h5.get("features")
            qOffset = self.dbStruct.numDb

            # Method 1: use sklearn ========================================
            qFeat = h5feat[index + qOffset]
            posFeat = h5feat[self.nontrivial_positives[index].tolist()]
            # knn = NearestNeighbors(n_jobs=1)  # TODO replace with faiss?
            # knn.fit(posFeat)
            # dPos, posNN = knn.kneighbors(qFeat.reshape(1, -1), 1)
            # dPos = dPos.item()
            # posIndex = self.nontrivial_positives[index][posNN[0]].item()

            # Method2: use Torch
            qFeat_t = torch.tensor(qFeat)
            posFeat_t = torch.tensor(posFeat)
            dist = torch.norm(qFeat_t - posFeat_t, dim=1, p=None)
            result = dist.topk(1, largest=False)
            dPos_t, posNN_t = result.values, result.indices
            posIndex_t = self.nontrivial_positives[index][posNN_t].item()

            dPos = dPos_t.numpy()
            posIndex = posIndex_t
            # End of Method 1/Method2 =======================================

            negSample = np.random.choice(self.potential_negatives[index], self.nNegSample)
            negSample = np.unique(np.concatenate([self.negCache[index], negSample]))

            # Method 1: use sklearn ========================================
            negFeat = h5feat[negSample.tolist()]

            # knn.fit(negFeat)
            # dNeg, negNN = knn.kneighbors(qFeat.reshape(1, -1), self.nNeg * 10)  # to quote netvlad paper code: 10x is hacky but fine
            # dNeg = dNeg.reshape(-1)
            # negNN = negNN.reshape(-1)

            # Method2: use Torch
            negFeat_t = torch.tensor(negFeat)
            dist = torch.norm(qFeat_t - negFeat_t, dim=1, p=None)
            result = dist.topk(self.nNeg * 10, largest=False)
            dNeg_t, negNN_t = result.values, result.indices
            dNeg = dNeg_t.numpy()
            negNN = negNN_t.numpy()
            # End of Method 1/Method2 =======================================

            # try to find negatives that are within margin, if there aren't any return none
            violatingNeg = dNeg < dPos + self.margin**0.5

            if np.sum(violatingNeg) < 1:
                # if none are violating then skip this query
                return None

            negNN = negNN[violatingNeg][:self.nNeg]
            negIndices = negSample[negNN].astype(np.int32)
            self.negCache[index] = negIndices

        query = load_pc_file(self.opt, join(self.img_dir, self.dbStruct.qImage[index].split('.')[0] + '.bin'))
        positive = load_pc_file(self.opt, join(self.img_dir, self.dbStruct.dbImage[posIndex].split('.')[0] + '.bin'))

        negatives = []
        for negIndex in negIndices:
            negative = load_pc_file(self.opt, join(self.img_dir, self.dbStruct.dbImage[negIndex].split('.')[0] + '.bin'))

            negatives.append(negative)

        negatives = torch.stack(negatives, 0)
        return query, positive, negatives, [index, posIndex] + negIndices.tolist()

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