inference.py

#!/usr/bin/env python
# Copyright (c) Facebook, Inc. and its affiliates.
"""
Detectron2 training script with a plain training loop.

This script reads a given config file and runs the training or evaluation.
It is an entry point that is able to train standard models in detectron2.

In order to let one script support training of many models,
this script contains logic that are specific to these built-in models and therefore
may not be suitable for your own project.
For example, your research project perhaps only needs a single "evaluator".

Therefore, we recommend you to use detectron2 as a library and take
this file as an example of how to use the library.
You may want to write your own script with your datasets and other customizations.

Compared to "train_net.py", this script supports fewer default features.
It also includes fewer abstraction, therefore is easier to add custom logic.
"""

import logging
import pickle
import os
import yaml
import time
import datetime
import itertools
from collections import OrderedDict
import torch
import torch.nn as nn
from tqdm import tqdm
from yacs.config import CfgNode as CN
from torch.nn.parallel import DistributedDataParallel
from contextlib import ExitStack, contextmanager
from pathlib import Path
from matplotlib import pyplot as plt

import detectron2.utils.comm as comm
import detectron2.data.transforms as T
from detectron2.checkpoint import DetectionCheckpointer, PeriodicCheckpointer
from detectron2.config import get_cfg
from detectron2.data import (
    MetadataCatalog,
    DatasetMapper,
    build_detection_test_loader,
    build_detection_train_loader,
)
from detectron2.engine import default_argument_parser, default_setup, default_writers, launch
from detectron2.evaluation import (
    CityscapesInstanceEvaluator,
    CityscapesSemSegEvaluator,
    COCOEvaluator,
    COCOPanopticEvaluator,
    DatasetEvaluators,
    LVISEvaluator,
    PascalVOCDetectionEvaluator,
    SemSegEvaluator,
    inference_on_dataset,
    print_csv_format,
)
from detectron2.modeling import build_model
from detectron2.solver import build_lr_scheduler, build_optimizer
from detectron2.utils.events import EventStorage

from dataset import setup_dataset
from model_zoo.RetinaNetPoint import RetinaNetPoint
from evaluation.COCOPointEvaluator import COCOPointEvaluator
from detectron2.utils.comm import get_world_size
from detectron2.utils.logger import log_every_n_seconds
torch.multiprocessing.set_sharing_strategy('file_system')

logger = logging.getLogger("detectron2")

import pdb


@contextmanager
def inference_context(model):
    """
    A context where the model is temporarily changed to eval mode,
    and restored to previous mode afterwards.

    Args:
        model: a torch Module
    """
    training_mode = model.training
    model.eval()
    yield
    model.train(training_mode)

def plot_preds(preds, save_dir):
    for pred in tqdm(preds):
        image = pred['image']
        xs = pred['points'][:, 0]
        ys = pred['points'][:, 1]
        save_path = os.path.join(save_dir, "images", pred['file_name'].split('/')[-1])
        
        # Create the plot
        fig, ax = plt.subplots()
        ax.imshow(image.permute(1, 2, 0))
        ax.plot(xs, ys, 'rx')
        fig.savefig(save_path, dpi=150)
        plt.close(fig)
        del fig, ax
    
def record_preds(results, preds, data, optimal_confidence, dataset_name):
    result = {k: v for k, v in data.items()}
    
    # Remove predictions with confidence level below optimal confidence
    if len(preds) > 0:
        # Get the index of the first element in scores which has score below optimal confidence
        idx_lim = 0
        while idx_lim < preds.scores.shape[0] and preds.scores[idx_lim] > optimal_confidence:
            idx_lim += 1
        valid_preds = preds[:idx_lim]
    else:
        valid_preds = preds
    result['points'] = valid_preds.pred_boxes.get_centers().to('cpu')
    result['scores'] = valid_preds.scores.to('cpu')
    
    # Append to the results dictionary
    results[dataset_name].append(result)
    
def do_inference(cfg, inference_cfg, model, resume=False, distributed=False):
    logger = logging.getLogger("detectron2")
    model.eval()
    
    # Initialize results dict
    results = {dataset_name: [] for dataset_name in cfg.DATASETS.TEST}
    
    # Detectron2 style preparation
    num_devices = get_world_size()
    start_time = time.perf_counter()
    total_data_time = 0
    total_compute_time = 0
    total_record_time = 0
    
    # compared to "train_net.py", we do not support accurate timing and
    # precise BN here, because they are not trivial to implement in a small training loop
    for dataset_name in cfg.DATASETS.TEST:
        data_loader = build_detection_test_loader(cfg, dataset_name)
        total = len(data_loader)  # inference data loader must have a fixed length
        num_warmup = min(5, total - 1)
        logger.info("Start inference on {} batches of dataset {}".format(len(data_loader), dataset_name))
        with ExitStack() as stack:
            if isinstance(model, nn.Module):
                stack.enter_context(inference_context(model))
            stack.enter_context(torch.no_grad())
            with torch.no_grad():
                # Run inference
                start_data_time = time.perf_counter()
                for idx, data in enumerate(data_loader):
                    total_data_time += time.perf_counter() - start_data_time
                    if idx == num_warmup:
                        start_time = time.perf_counter()
                        total_data_time = 0
                        total_compute_time = 0
                        total_record_time = 0

                    # Run inference
                    start_compute_time = time.perf_counter()
                    preds = model(data)
                    if torch.cuda.is_available():
                        torch.cuda.synchronize()
                    total_compute_time += time.perf_counter() - start_compute_time

                    # Record the predictions
                    start_record_time = time.perf_counter()
                    record_preds(results, preds[0]['instances'], data[0], inference_cfg.OPTIMAL_CONFIDENCE, dataset_name)
                    total_record_time += time.perf_counter() - start_record_time

                    iters_after_start = idx + 1 - num_warmup * int(idx >= num_warmup)
                    data_seconds_per_iter = total_data_time / iters_after_start
                    compute_seconds_per_iter = total_compute_time / iters_after_start
                    record_seconds_per_iter = total_record_time / iters_after_start
                    total_seconds_per_iter = (time.perf_counter() - start_time) / iters_after_start

                    if idx >= num_warmup * 2 or compute_seconds_per_iter > 5:
                        eta = datetime.timedelta(seconds=int(total_seconds_per_iter * (total - idx - 1)))
                        log_every_n_seconds(
                            logging.INFO,
                            (
                                f"Inference done {idx + 1}/{total}. "
                                f"Dataloading: {data_seconds_per_iter:.4f} s/iter. "
                                f"Inference: {compute_seconds_per_iter:.4f} s/iter. "
                                f"Record: {record_seconds_per_iter:.4f} s/iter. "
                                f"Total: {total_seconds_per_iter:.4f} s/iter. "
                                f"ETA={eta}"
                            ),
                            n=5,
                            name="detectron2"
                        )
                    start_data_time = time.perf_counter()
        
        # Measure the time only for this worker (before the synchronization barrier)
        total_time = time.perf_counter() - start_time
        total_time_str = str(datetime.timedelta(seconds=total_time))
        # NOTE this format is parsed by grep
        logger.info(
            "Total inference time: {} ({:.6f} s / iter per device, on {} devices)".format(
                total_time_str, total_time / (total - num_warmup), num_devices
            )
        )
        total_compute_time_str = str(datetime.timedelta(seconds=int(total_compute_time)))
        logger.info(
            "Total inference pure compute time: {} ({:.6f} s / iter per device, on {} devices)".format(
                total_compute_time_str, total_compute_time / (total - num_warmup), num_devices
            )
        )
        
        # Synchronize distributed processes
        if distributed:
            comm.synchronize()
        
        # Create save directories
        results_dir = os.path.join(inference_cfg.RESULTS_DIR)
        Path(results_dir).mkdir(parents=True, exist_ok=True)
        for dataset_name in cfg.DATASETS.TEST:
            Path(os.path.join(results_dir, dataset_name)).mkdir(parents=True, exist_ok=True)
        
        # Loop through all predictions and save them as PKL files
        for dataset_name in cfg.DATASETS.TEST:
            logger.info(f"Saving predictions for {dataset_name}")
            results_dataset = results[dataset_name]
            save_dir = os.path.join(results_dir, dataset_name)
            Path(os.path.join(save_dir, "annotations")).mkdir(parents=True, exist_ok=True)
            Path(os.path.join(save_dir, "images")).mkdir(parents=True, exist_ok=True)
            for pred in results_dataset:
                encoding = pred['file_name'].split('/')[-1].split('.')[0]
                save_path = os.path.join(save_dir, "annotations", encoding + ".pkl")
                save_data = {
                    'points': pred['points'],
                    'scores': pred['scores']
                }
                with open(save_path, 'wb') as f:
                    pickle.dump(save_data, f, protocol=pickle.HIGHEST_PROTOCOL)
        
            if torch.cuda.is_available():
                torch.cuda.synchronize()
        
            # Plot predictions
            if inference_cfg.PLOT_PREDS:
                logger.info(f"Plotting predictions for {dataset_name}")
                plot_preds(results_dataset, save_dir)

def setup_inference(args):
    assert "INFERENCE" in args.opts, "No inference config file supplied!"
    inference_cfg_idx = args.opts.index("INFERENCE")
    inference_cfg_path = args.opts[inference_cfg_idx + 1]
    with open(inference_cfg_path) as f:
        inference_cfg = CN(yaml.load(f, Loader=yaml.UnsafeLoader))
    
    # Remove custom parameters from args
    del args.opts[inference_cfg_idx]
    del args.opts[inference_cfg_idx]
    
    return inference_cfg
        
def setup(args):
    """
    Create configs and perform basic setups.
    """
    # Extract custom configs
    inference_cfg = setup_inference(args)
    
    cfg = get_cfg()
    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    cfg.freeze()
    default_setup(
        cfg, args
    )  # if you don't like any of the default setup, write your own setup code
    return cfg, inference_cfg


def main(args):
    cfg, inference_cfg = setup(args)
    
    # Setup the dataset
    datasets_dir = "/var/storage/myeghiaz/Detection" # all datasets are stored in this directory
    setup_dataset(cfg, datasets_dir)
    
    model = build_model(cfg)
    DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
        cfg.MODEL.WEIGHTS, resume=True
    )

    distributed = comm.get_world_size() > 1
    if distributed:
        model = DistributedDataParallel(
            model, device_ids=[comm.get_local_rank()], broadcast_buffers=False
        )
    
    do_inference(cfg, inference_cfg, model, resume=True, distributed=distributed)


if __name__ == "__main__":
    args = default_argument_parser().parse_args()
    print("Command Line Args:", args)
    launch(
        main,
        args.num_gpus,
        num_machines=args.num_machines,
        machine_rank=args.machine_rank,
        dist_url=args.dist_url,
        args=(args,),
    )