evaluate.py

import argparse
import os
import pandas as pd
import numpy as np
import pprint
import logging

from pathlib import Path
from tqdm.autonotebook import tqdm 

from utils.inference import load_model_from_config, setup_inference
from utils.eval_utils import MIDOGEvaluation
from utils.factory import ConfigCreator, ModelFactory


# Set default configurations
BATCH_SIZE = 8
DEVICE = 'cuda'
NMS_THRESH = 0.3
NUM_WORKERS = 8
OVERLAP = 0.3
VERBOSE = False
SPLIT = 'test'
BOX_FORMAT = 'cxcy'

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--batch_size",     type=int, default=BATCH_SIZE, help="Batch size.")
    parser.add_argument("--box_format",     type=str, default=BOX_FORMAT, help='Box format (default: xyxy).')
    parser.add_argument("--config_file",    type=str, help="Existing config file.", required=True)
    parser.add_argument("--dataset",        type=str, help="Dataset filepath.", required=True)
    parser.add_argument("--device",     	type=str, default=DEVICE, help="Device.")
    parser.add_argument("--img_dir",        type=str, help="Image directory.", required=True)
    parser.add_argument("--nms_thresh",     type=float, default=NMS_THRESH, help="Final NMS threshold.")
    parser.add_argument("--num_workers",    type=int, default=NUM_WORKERS, help="Number of processes.")
    parser.add_argument("--overlap",        type=float, default=OVERLAP, help="Overlap between patches.")
    parser.add_argument("--overwrite",      action="store_true", help="If true, existing results are overwritten.")
    parser.add_argument("--split",          type=str, default=SPLIT, help="Data split to evaluate.")
    parser.add_argument("--wsi",            action="store_true", help="Processes WSI")
    return parser.parse_args()


def evaluate(
        batch_size: int,
        box_format: str, 
        config_file: str, 
        dataset: str | pd.DataFrame, 
        nms_thresh: float,
        device: str,  
        img_dir: str, 
        num_workers: int,
        overlap: float, 
        overwrite: bool,
        split: str,
        wsi: bool,
        logger: logging.Logger = None
):
    
    # Check existing config_file 
    if not Path(config_file).exists():
        raise FileNotFoundError(f"Could not find this config_file: {config_file}. Provide existing config_file.")
    
    # Check dataset 
    if isinstance(dataset, str):
        if not Path(dataset).exists():
            raise FileNotFoundError(f"Could not find dataset: {dataset}")
        dataset = pd.read_csv(dataset)
    
    # Convert to xmin, ymin, xmax, ymax
    if box_format == 'cxcy':
        radius = 25
        dataset = dataset.assign(xmin=dataset['x'] - radius)
        dataset = dataset.assign(ymin=dataset['y'] - radius)
        dataset = dataset.assign(xmax=dataset['x'] + radius)
        dataset = dataset.assign(ymax=dataset['y'] + radius)

    # Check image directory
    img_dir = Path(img_dir)
    if not img_dir.is_dir():
        raise ValueError(f"This is not a directory: {str(img_dir)}")
    
    # Load the model 
    model, config = load_model_from_config(config_file, det_thresh=0.05)  # We set low det_thresh here to measure AP
 
    # Setup the inference 
    processor, patch_config = setup_inference(
        model=model,
        is_wsi=wsi,
        batch_size=batch_size,
        num_workers=num_workers,
        nms_thresh=nms_thresh,
        device=device,
        patch_size=config.patch_size,
        overlap=overlap,
        overwrite=overwrite,
        logger=logger
        )

    print('Loaded model configurations:')
    pprint.pprint(config)
    print()

    # Filter dataset
    test_dataset = dataset.query("split == @split")
    filenames = test_dataset.filename.unique()

    # Initialize predictions
    preds = {}

    # Run inference over all images 
    for file in tqdm(filenames, desc="Running inference"):

        # Create image path
        image_path = img_dir.joinpath(file)
        
        # Process images individually
        results = processor.process_single(image_path, patch_config)

        # Collect predictions
        preds[file] = results

    # Filter only mitotic figures
    filtered_dataset = test_dataset.query('label == 1')

    # Init evaluation
    evaluation = MIDOGEvaluation(
        gt_file=filtered_dataset,
        preds=preds,
        output_file=Path(config_file).with_suffix('.json'),
        det_thresh=config.det_thresh,
        split=split
    )

    # evaluate
    evaluation.evaluate()
    print('Evaluation done.')
    

    print(f'Evaluation results for {args.split} split')
    # show aggregate metrics
    pprint.pprint(evaluation._metrics['aggregates'])


    
def main(args):

    # Convert args
    eval_kwargs = vars(args)

    # Run optimize
    evaluate(**eval_kwargs)


if __name__ == "__main__":
    args = get_args()
    main(args)
    print('End of script.')