O2020/predict_model.py

import sys
import yaml
import torch
import torch.nn as nn
import json
import os
import numpy as np
import pandas as pd
from tqdm.auto import tqdm
from utils.misc import DETRModel, Logger, sample_metrics, get_data, AverageMeter
from utils.dataset import CustomDataset
import utils.misc as util
import torchvision.transforms as T

'''
Main function of this script. It is explained in more detail in the tutorial.
'''
def run(parameters, directories, ids_dict):

    #Retrieving needed directories
    csv_path = directories['csv_path']
    train_test_path = directories['train_test_path']
    test_dir = os.path.join(train_test_path, "test")
    output_path = directories['output_path']
    test_model = parameters['test_model']
    trained_model_file  = os.path.join(output_path, f'{test_model}.pth')
    detr_path = directories['detr_path']
    sys.path.extend([detr_path])

    #Retrieving needed parameters
    num_classes = parameters['num_classes']
    model_name = parameters['model_name']
    dev = parameters['dev']
    weight_dict = parameters['weight_dict']
    null_class_coef = parameters['null_class_coef']
    pre_trained = parameters['pre_trained']
    from_scratch = parameters['from_scratch']

    device = torch.device(dev)

    if pre_trained or from_scratch:

        num_queries = parameters['num_queries']

        #Importing classes that are required from DETR's repository
        from models.detr import DETR,SetCriterion
        from models.backbone import Backbone,Joiner
        from models.position_encoding import PositionEmbeddingSine
        from models.transformer import Transformer

        #Creating empty model in the case of being pre-trained or being trained from scratch
        #because some parameters might be different from the one that is on DETR's repository
        hidden_dim = 256
        backbone = Backbone(model_name, train_backbone=True, return_interm_layers=False, dilation=True)
        pos_enc = PositionEmbeddingSine(hidden_dim // 2, normalize=True)
        backbone_with_pos_enc = Joiner(backbone, pos_enc)
        backbone_with_pos_enc.num_channels = backbone.num_channels
        transformer = Transformer(d_model=hidden_dim, return_intermediate_dec=True)
        model = DETR(backbone_with_pos_enc, transformer, num_classes=num_classes, num_queries=num_queries, aux_loss=False)
    else:
        #Initializing a model with the characteristics of that from DETR's repository
        model = DETRModel(model_name=model_name, num_classes=num_classes)

    #Loading the fine-tuned or trained models
    try:
        model.load_state_dict(torch.load(trained_model_file, map_location=torch.device('cpu')))
        model.to(device)
        model.eval()
    except IOError as e:
        print("I/O error{}: {}".format(e.errno, e.strerror))

    #Creation of the loss function
    cost_bbox = weight_dict['loss_bbox']
    cost_giou = weight_dict['loss_giou']
    cost_class = weight_dict['loss_ce']

    from models.matcher import HungarianMatcher
    from models.detr import SetCriterion

    matcher = HungarianMatcher(cost_class=cost_class, cost_bbox=cost_bbox, cost_giou=cost_giou)
    losses = ['labels', 'boxes', 'cardinality']

    criterion = SetCriterion(num_classes,
                             matcher, weight_dict,
                             eos_coef=null_class_coef,
                             losses=losses)
    criterion = criterion.to(device)
    criterion.eval()

    #Retrieving information about the images in the test set
    train_ids, val_ids, test_ids = ids_dict.values()

    test_csv = os.path.join(csv_path, 'test.csv')
    test_marking = pd.read_csv(test_csv)
    test_image_data = test_marking.groupby('image_id')

    test_info = [get_data(img_id, test_image_data) for img_id in test_ids]
    null_test_images = [x['image_id'] for x in test_info if len(x['boxes'])==0]

    print(f'Total number of images in test set: {len(test_info)}')
    print('Images with no boxes in test set: ', len(null_test_images))

    null_test_boxes = np.zeros(num_classes)
    for img_id in test_info:
        for i in range(len(img_id['boxes'])):
            null_test_boxes[img_id['labels'][i]] += 1

    #Creating test dataset
    test_ds = CustomDataset(test_info, test_dir, image_set = 'test')
    size = len(test_ds)

    logger = util.Logger('test_metrics', format='json')
    log = None

    bar = tqdm(test_ds, total=size, leave=False)

    #Iterating through the images in test set to compute the metrics
    for i, (img, target) in enumerate(bar):

        img = img.unsqueeze(0)
        img = img.to(device)
        target = {k: v.to(device) for k, v in target.items()}
        output = model(img)
        loss_dict = criterion(output, [target])

        if log is None:
            log = {k:AverageMeter() for k in loss_dict}
            log['total_loss'] = AverageMeter()
            for i in range(len(null_test_boxes)):
                log[f'avg_prec_{i}'] = AverageMeter(id=i)
                log[f'avg_rec_{i}'] = AverageMeter(id=i)
                log[f'avg_f1_{i}'] = AverageMeter(id=i)

        for k,v in loss_dict.items():
            log[k].update(v.item())

        total_loss = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)
        log['total_loss'].update(total_loss.item())

        metrics = sample_metrics(output, [target])
        for i in range(len(null_test_boxes)):
            class_metric = metrics[:, i]
            log[f'avg_prec_{i}'].update(class_metric[0])
            log[f'avg_rec_{i}'].update(class_metric[1])
            log[f'avg_f1_{i}'].update(class_metric[2])

        bar.set_postfix({k:v.avg(null_test_boxes) for k,v in log.items() if v.id is None})

    log = {k:v.avg(null_test_boxes) for k,v in log.items()}
    logger.save(log)

    for k,v in log.items():
        print(k, v)

if __name__ == "__main__":

    #Attempting to read dir.yaml file that contains directories information
    try:
        with open(os.path.join("configs","dir.yaml"), "r") as fp:
            directories = yaml.load(fp, Loader=yaml.FullLoader)
    except:
        sys.exit('File dir.yaml could not be read.')

    #Attempting to read parameters.yaml file that contains parameters information
    try:
        with open(os.path.join("configs","parameters.yaml"), "r") as fp:
            parameters = yaml.load(fp, Loader=yaml.FullLoader)
    except:
        sys.exit('File parameters.yaml could not be read.')

    #Attempting to read ids_dict.json file that contains sets' images ids
    try:
        with open("ids_dict.json", "r") as fp:
            ids_dict = json.load(fp)
    except:
        sys.exit('File ids_dict.json could not be read.')

    run(parameters, directories, ids_dict)