eval_airloc.py

import pickle
import yaml
import argparse
from scipy.spatial.distance import cdist

import sys
sys.path.append('.')

import torch
from torch.utils import data
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data.sampler import SubsetRandomSampler
from tqdm import tqdm

from model.build_model import build_netvlad, build_airloc_edge, build_gcn
from datasets.mp3d_airloc.mp3d_triplet_v3_edge import mp3d
from datasets.utils.batch_collator import eval_custom_collate
from statistics import mean
from utils.generate_reference import generate,points_to_obj_desc

torch.autograd.set_detect_anomaly(False)
torch.autograd.profiler.profile(False)
torch.autograd.profiler.emit_nvtx(False)

def eval(configs):
    base_dir = configs['base_dir']
    batch_size = configs['batch_size']
    scenes = configs["scenes"]
    ap_thres = configs['ap_thres']
    m = configs["ap_wght"]
    method = configs['method']

    test_dataset = mp3d(base_dir=base_dir,test_scenes = scenes)
    test_loader = data.DataLoader(dataset=test_dataset, batch_size=batch_size, collate_fn=eval_custom_collate,shuffle = True)

    if method in ["airloc","airloc_without_edge",'netvlad_mean','netvlad']:
        model = build_netvlad(configs)
        model.eval() 

    if method in ["airloc"]:
        edge_model = build_airloc_edge(configs)
        edge_model.eval()
    
    if method in ["gcn","gcn_mean"]:
        model = build_gcn(configs)
        model.eval() 

    
    f = open(configs['database']['config_path'], 'r', encoding='utf-8')
    ref_configs = yaml.safe_load(f.read())
    ref_configs['db_path'] = configs['database']['db_path']
    ref_configs['base_dir'] = configs['database']['db_raw_path']
    ref_configs['K'] = configs['database']['K']
    ref_configs['netvlad_model_path'] = configs['netvlad_model_path'] 
    ref_configs['graph_model_path'] = configs['graph_model_path']
    ref_configs['method'] = method
    
    ref = generate(ref_configs)

    rooms = []
    ref_data = []
    ref_points = []
    for key in ref.keys():
        scene =  key.split("_")[1]
        if scene in scenes:
            rooms.append(key)
            ref_data.append(ref[key][0])
            ref_points.append(ref[key][1])
    
    with torch.no_grad():
        test_accuracy = []

        for step, anchor_pts in enumerate(tqdm(test_loader)):
            anchor_pt, query = points_to_obj_desc(anchor_pts,model,method)
            if query == None:
                continue
            
            if method in ["airloc","airloc_without_edge",'gcn']:
                room_sim = []
                for i in range(len(ref_data)):
                    ob_sim = []
                    for j in range(len(query)):
                        prod = (ref_data[i] @ query[j].T.to(ref_data[i].device))
                        ref = torch.norm(ref_data[i], dim = 1).to(prod.device)        # Normalising
                        qry = torch.norm(query[j], dim = 1).to(prod.device)  
                        mat = ((prod/qry).T/ref).T
                        prod = torch.max(mat, dim = 0).values
                        ob_sim.append(torch.sum(prod))
                    room_sim.append(torch.stack(ob_sim))
                    
                room_sim = torch.stack(room_sim)
            
            elif method in ["netvlad_mean",'gcn_mean','netvlad']:
                room_sim = 1-torch.cdist(torch.stack(ref_data),torch.stack(query).to(torch.stack(ref_data)),p = 2)
                
            mInds = torch.argsort(room_sim,axis=0)
            mValues, indices = torch.sort(room_sim,axis=0)

            positive = 0

            if method == "airloc":
                for i, obj in enumerate(anchor_pts):
                    query_room = obj[0]["room_image_name"][0]
            
                    if (mValues[-1,i]-mValues[-2,i])> ap_thres : 
                        ref_room = rooms[mInds[-1][i]]
                    
                    else:
                        anc_e = edge_model(anchor_pt)
                        ref_e = edge_model(ref_points)
                        prod = (ref_e@ anc_e.T)
                        room_sim_ = m*torch.nn.functional.normalize(room_sim,dim = 0)+torch.nn.functional.normalize(prod,dim = 0).to(room_sim.device)
                        mInds_ = torch.argsort(room_sim_,axis=0)
                        ref_room = rooms[mInds_[-1][i]]

                    if (ref_room==query_room):
                        positive+=1

            elif method in  ["airloc_without_edge","netvlad_mean",'gcn','gcn_mean','netvlad']   :  
                for i, obj in enumerate(anchor_pts):
                    query_room = obj[0]["room_image_name"][0]
                    ref_room = rooms[mInds[-1][i]]
                    
                    if (ref_room==query_room):
                        positive+=1


            acc= positive/len(query)
            test_accuracy.append(acc)
        
        print("Test_accuracy : ", mean(test_accuracy) )



def main():
    parser = argparse.ArgumentParser(description="Evaluating")
    parser.add_argument(
        "-c", "--config_file",
        dest = "config_file",
        type = str, 

        default = ""
    )
    parser.add_argument(
        "-g", "--gpu",
        dest = "gpu",
        type = int, 
        default = 1
    )
    args = parser.parse_args()

    config_file = args.config_file
    f = open(config_file, 'r', encoding='utf-8')
    configs = f.read()
    configs = yaml.safe_load(configs)
    configs['use_gpu'] = args.gpu

    # for m in ["airloc","netvlad_mean","airloc_without_edge","gcn","gcn_mean","netvlad"]:
    #     configs["method"] = m
    #     print("method = ",m)
    eval(configs)
    
if __name__ == "__main__":
    main()