Meehai/dronescapes

1#!/usr/bin/env python32"""3Evaluation script for semantic segmentation for dronescapes. Outputs F1Score and mIoU for the classes and each frame.4Usage: ./evaluate_semantic_segmentation.py y_dir gt_dir --classes C1 .. Cn [--class_weights W1 .. Wn] -o results.csv5"""6import sys7import os8from loggez import loggez_logger as logger9from pathlib import Path10from argparse import ArgumentParser, Namespace11from tempfile import TemporaryDirectory12from multiprocessing import Pool13from functools import partial14from torchmetrics.functional.classification import multiclass_stat_scores15from tqdm import tqdm16import torch as tr17import numpy as np18import pandas as pd19 20sys.path.append(Path(__file__).parents[1].__str__())21from dronescapes_reader import MultiTaskDataset22from dronescapes_reader.dronescapes_representations import SemanticRepresentation23 24def compute_metrics(tp: np.ndarray, fp: np.ndarray, tn: np.ndarray, fn: np.ndarray) -> pd.DataFrame:25    precision = tp / (tp + fp)26    recall = tp / (tp + fn)27    f1 = 2 * precision * recall / (precision + recall)28    iou = tp / (tp + fp + fn)29    return pd.DataFrame([precision, recall, f1, iou], index=["precision", "recall", "f1", "iou"]).T30 31def compute_metrics_by_class(df: pd.DataFrame, class_name: str) -> pd.DataFrame:32    df = df.query("class_name == @class_name").drop(columns="class_name")33    df.loc["all"] = df.sum()34    df[["precision", "recall", "f1", "iou"]] = compute_metrics(df["tp"], df["fp"], df["tn"], df["fn"])35    df.insert(0, "class_name", class_name)36    df = df.fillna(0).round(3)37    return df38 39def _do_one(i: int, reader: MultiTaskDataset, num_classes: int) -> tuple[tr.Tensor, str]:40    data, name = reader[i][0:2]41    y = data["pred"].argmax(-1) if data["pred"].dtype != tr.int64 else data["pred"]42    gt = data["gt"].argmax(-1) if data["gt"].dtype != tr.int64 else data["gt"]43    return multiclass_stat_scores(y, gt, num_classes=num_classes, average=None)[:, 0:4], name44 45def compute_raw_stats_per_frame(reader: MultiTaskDataset, classes: list[str], n_workers: int = 1) -> pd.DataFrame:46    res = tr.zeros((len(reader), len(classes), 4)).long() # (N, NC, 4)47 48    map_fn = map if n_workers == 1 else Pool(n_workers).imap49    do_one_fn = partial(_do_one, reader=reader, num_classes=len(classes))50    map_res = list(tqdm(map_fn(do_one_fn, range(len(reader))), total=len(reader)))51    res, index = tr.stack([x[0] for x in map_res]).reshape(len(reader) * len(classes), 4), [x[1] for x in map_res]52 53    df = pd.DataFrame(res, index=np.repeat(index, len(classes)), columns=["tp", "fp", "tn", "fn"])54    df.insert(0, "class_name", np.array(classes)[:, None].repeat(len(index), 1).T.flatten())55    return df56 57def compute_final_per_scene(res: pd.DataFrame, scene: str, classes: list[str],58                            class_weights: list[float]) -> tuple[float, float]:59    df = res.iloc[[x.startswith(scene) for x in res.index]]60    # aggregate for this class all the individual predictions61    df_scene = df[["class_name", "tp", "fp", "tn", "fn"]].groupby("class_name") \62        .apply(lambda x: x.sum(), include_groups=False).loc[classes]63    df_metrics = compute_metrics(df_scene["tp"], df_scene["fp"], df_scene["tn"], df_scene["fn"])64    iou_weighted = (df_metrics["iou"] * class_weights).sum()65    f1_weighted = (df_metrics["f1"] * class_weights).sum()66    return scene, iou_weighted, f1_weighted67 68def _check_and_symlink_dirs(y_dir: Path, gt_dir: Path) -> Path:69    """checks whether the two provided paths are actual full of npz directories and links them together in a tmp dir"""70    assert (l := {x.name for x in y_dir.iterdir()}) == (r := {x.name for x in gt_dir.iterdir()}), f"{l} \n vs \n {r}"71    assert all(x.endswith(".npz") for x in [*l, *r]), f"Not dirs of only .npz files: {l} \n {r}"72    (temp_dir := Path(TemporaryDirectory().name)).mkdir(exist_ok=False)73    os.symlink(y_dir, temp_dir / "pred")74    os.symlink(gt_dir, temp_dir / "gt")75    return temp_dir76 77def get_args() -> Namespace:78    parser = ArgumentParser()79    parser.add_argument("y_dir", type=lambda p: Path(p).absolute())80    parser.add_argument("gt_dir", type=lambda p: Path(p).absolute())81    parser.add_argument("--output_path", "-o", type=Path, required=True)82    parser.add_argument("--classes", required=True, nargs="+")83    parser.add_argument("--class_weights", nargs="+", type=float)84    parser.add_argument("--scenes", nargs="+", default=["all"], help="each scene will get separate metrics if provided")85    parser.add_argument("--overwrite", action="store_true")86    parser.add_argument("--n_workers", type=int, default=1)87    args = parser.parse_args()88    if args.class_weights is None:89        logger.info("No class weights provided, defaulting to equal weights.")90        args.class_weights = [1 / len(args.classes)] * len(args.classes)91    assert (a := len(args.class_weights)) == (b := len(args.classes)), (a, b)92    assert np.fabs(sum(args.class_weights) - 1) < 1e-3, (args.class_weights, sum(args.class_weights))93    assert args.output_path.suffix == ".csv", f"Prediction file must end in .csv, got: '{args.output_path.suffix}'"94    if len(args.scenes) > 0:95        logger.info(f"Scenes: {args.scenes}")96    if args.output_path.exists() and args.overwrite:97        os.remove(args.output_path)98    assert args.n_workers >= 1 and isinstance(args.n_workers, int), args.n_workers99    return args100 101def main(args: Namespace):102    # setup to put both directories in the same parent directory for the reader to work.103    temp_dir = _check_and_symlink_dirs(args.y_dir, args.gt_dir)104    pred_repr = SemanticRepresentation("pred", classes=args.classes, color_map=[[0, 0, 0]] * len(args.classes))105    gt_repr = SemanticRepresentation("gt", classes=args.classes, color_map=[[0, 0, 0]] * len(args.classes))106    reader = MultiTaskDataset(temp_dir, task_names=["pred", "gt"], task_types={"pred": pred_repr, "gt": gt_repr},107                              handle_missing_data="drop", normalization=None)108    assert (a := len(reader.files_per_repr["gt"])) == (b := len(reader.files_per_repr["pred"])), f"{a} vs {b}"109 110    # Compute TP, FP, TN, FN for each frame111    raw_stats = compute_raw_stats_per_frame(reader, args.classes, args.n_workers)112    logger.info(f"Stored raw metrics file to: '{args.output_path}'")113    Path(args.output_path).parent.mkdir(exist_ok=True, parents=True)114    raw_stats.to_csv(args.output_path)115 116    # Compute Precision, Recall, F1, IoU for each class and put them together in the same df.117    metrics_per_class = pd.concat([compute_metrics_by_class(raw_stats, class_name) for class_name in args.classes])118 119    # Aggregate the class-level metrics to the final metrics based on the class weights (compute globally by stats)120    final_agg = []121    for scene in args.scenes: # if we have >1 scene in the test set, aggregate the results for each of them separately122        final_agg.append(compute_final_per_scene(metrics_per_class, scene, args.classes, args.class_weights))123    final_agg = pd.DataFrame(final_agg, columns=["scene", "iou", "f1"]).set_index("scene")124    if len(args.scenes) > 1:125        final_agg.loc["mean"] = final_agg.mean()126    final_agg = (final_agg * 100).round(3)127    print(final_agg)128 129if __name__ == "__main__":130    main(get_args())131